API reference

Estimator

graphdml.GraphDML

Bases: BaseEstimator

Direct and peer effects on a network via graph double machine learning.

Estimates theta (average direct effect) and alpha (average peer effect) in the partially linear network model::

T = m(X, A) + e_T
Y = theta * T + alpha * (E @ T) + g(X, A) + e_Y

where E is a known exposure map and m, g are unknown functions of a node's own and its neighbors' covariates, learned with any scikit-learn model.

Parameters:

Name	Type	Description	Default
model_y	scikit-learn regressor	Learns E[Y | X, A]. Defaults to :func:graphdml.learners.default_regressor (regularised histogram gradient boosting).	None
model_t	scikit-learn regressor or classifier	Learns E[T | X, A]. For binary T a classifier's predict_proba is used. Defaults to :func:graphdml.learners.default_classifier for binary T, else the regressor.	None
model_g	scikit-learn regressor	Learns g(X, A) for score="iv-type". Defaults to model_y.	None
featurizer	object with ``transform(data)``, "own" or "neighborhood"	Builds the tabular features the nuisance models see. Defaults to :class:~graphdml.NeighborhoodFeatures (own + 1- and 2-hop neighbor aggregates).	None
exposure	(str, ExposureMap, list or None)	How neighbors' treatments reach a node: "sum", "mean", "weighted", ... A list fits one peer coefficient per exposure. None fits the direct effect only.	"sum"
score	('iv-type', 'partialling-out')	"iv-type" stays consistent if either the treatment model or the outcome-confounding model g is right, and had lower bias in 7 of 8 benchmark cells. "partialling-out" is the original score; it is attenuated toward zero when the treatment model is wrong (see docs/methodology.md).	"iv-type"
focal_set	('min_degree', 'random')	How to choose the focal set of nodes with independent scores. An array gives the focal nodes explicitly (validated). None uses every node and treats them as independent; intervals are then not valid under interference.	"min_degree"
nuisance_training	('buffered', 'focal')	Which nodes train the nuisance models for each fold. "focal" trains on the other folds' focal nodes only. "buffered" trains on every node whose data is independent of the held-out fold's scores, which is typically several times more data.	"buffered"
n_folds	int		5
n_rep	int	Repetitions of cross-fitting (with new focal sets and folds). Estimates are aggregated by the median, with the variance adjusted for split-to-split spread.	1
aggregation	('dml2', 'dml1')	"dml2" solves the pooled moment condition; "dml1" averages per-fold estimates.	"dml2"
clip_propensity	float	Clip classifier propensities to [c, 1 - c].	0.01
mode	(None, 'original')	"original" overrides the settings above with the original GDML procedure (see docs/design.md), for comparison and replication.	None
estimation_nodes	array of node positions or boolean mask	Restrict the estimate to a sub-population: only these nodes can be focal. All nodes still supply neighbors' treatments and nuisance training data.	None
random_state	(int, RandomState or None)		None

Attributes:

Name	Type	Description
coef_, se_, cov_	estimates, standard errors and covariance, ordered as	coef_names_ ("direct" then "peer:<exposure>").
ate_	the average total effect, treating every node vs none (direct + peer).
ade_, ape_	the direct and (first) peer effect.
total_weights_	weights turning ``coef_`` into the total effect: 1 for the direct	effect and, per exposure, its average value when every node is treated.
focal_nodes_, fold_assignment_	focal node positions and their folds (first rep).
n_focal_	effective sample size.
residuals_	per-focal-node predictions and residuals (first rep).
diagnostics_	dict of diagnostics; warnings_ : list of warning messages.

ate_ property

Average total effect: treating every node vs no node (direct + peer).

ade_ property

Average direct effect.

ape_ property

Average peer effect for the first exposure map (None if no exposure).

fit(data)

Estimate direct and peer effects from data.

summary_frame(alpha=0.05, total=True)

Estimates, standard errors, z statistics, p-values and confidence intervals.

Rows: direct (ADE), one peer:<exposure> row per exposure (APE), and, when there is a peer effect and total=True, total: the average total effect of treating every node vs no node (the ATE).

cluster_summary_frame(groups, alpha=0.05, total=True)

Like :meth:summary_frame, with standard errors clustered by groups.

groups gives a cluster label for every node (e.g. village). Use it when nodes may share common shocks beyond their dependency sets. Requires n_rep=1.

summary(alpha=0.05)

Human-readable results, diagnostics, warnings and identifying assumptions.

Data

graphdml.GraphData dataclass

Node covariates, treatment and outcome, plus the network that connects the nodes.

Parameters:

Name	Type	Description	Default
X	array of shape (n_nodes, n_features)	Pre-treatment covariates. Must be numeric and finite.	required
T	array of shape (n_nodes,)	Treatment, binary (0/1) or continuous.	required
Y	array of shape (n_nodes,)	Outcome.	required
adjacency	sparse matrix of shape (n_nodes, n_nodes)	adjacency[i, j] != 0 means node j can influence node i (j is an in-neighbor of i). Undirected graphs must be symmetric. Values are only used as weights by :class:~graphdml.WeightedExposure; everything else uses the pattern.	required
feature_names	sequence of str		None
node_ids	array of shape (n_nodes,)	External identifiers such as Neo4j element ids. Defaults to 0..n_nodes-1.	None
directed	bool		False

Notes

Most users build this with :meth:from_edges, :meth:from_pandas or :meth:from_networkx. Self-loops are dropped (a node does not "influence itself" through the network; its own treatment enters the model directly).

n_edges property

Number of edges (unordered pairs for undirected graphs).

pattern cached property

Binary adjacency: pattern[i, j] == 1 iff j is an in-neighbor of i.

degree cached property

In-degree of every node (number of nodes that can influence it).

from_edges(X, T, Y, edges, *, weights=None, directed=False, feature_names=None, node_ids=None) classmethod

Build from an (n_edges, 2) array of (source, target) node positions.

source influences target. For undirected graphs the order is irrelevant. Duplicate edges keep their first occurrence; self-loops are dropped.

from_pandas(nodes, edges, *, features, treatment, outcome, node_id=None, source='source', target='target', weight=None, directed=False) classmethod

Build from a node table and an edge table.

node_id names the id column of nodes (defaults to the index). edges refers to nodes by those ids; source influences target.

from_networkx(graph, *, features, treatment, outcome, weight=None) classmethod

Build from a networkx graph whose nodes carry the given attributes.

For a DiGraph, an edge u -> v means u influences v.

with_values(*, T=None, Y=None, X=None, feature_names=None)

Return a copy with the treatment, outcome and/or covariates replaced.

If X changes its number of columns and no feature_names are given, the names are reset to X0, X1, ....

to_pandas()

Return (nodes, edges) DataFrames. Edges have source influencing target.

Exposure maps

graphdml.SumExposure

Bases: ExposureMap

Z_i = sum of in-neighbors' treatments (Z = A T).

graphdml.MeanExposure

Bases: ExposureMap

Z_i = average of in-neighbors' treatments (0 for nodes without neighbors).

graphdml.WeightedExposure

Bases: ExposureMap

Z_i = sum_j w_ij T_j using adjacency weights, optionally row-normalised.

graphdml.TwoHopExposure

Bases: ExposureMap

Treatments of nodes at distance exactly two (summed or averaged).

Used by :func:graphdml.two_hop_test to check the 1-hop interference assumption.

graphdml.MatrixExposure

Bases: ExposureMap

A user-supplied exposure operator E (sparse or dense (n, n)).

Features

graphdml.NeighborhoodFeatures

Bases: BaseEstimator

Own covariates plus aggregated neighbor covariates over one or more hops.

Hop 1 applies each aggregation in aggs to in-neighbors' covariates. Each further hop averages the previous hop's features over in-neighbors (mean propagation, as in SGC/SIGN-style graph features), so aggs=("mean",), hops=2 yields [X, Ā X, Ā² X] with Ā the row-normalised adjacency.

Parameters:

Name	Type	Description	Default
aggs	tuple of {"mean", "sum", "max", "min"}		("mean", "max", "min")
hops	int		2
include_degree	bool	Append each node's in-degree as a feature.	True

graphdml.OwnFeatures

Bases: BaseEstimator

Only the node's own covariates (ignores the network; i.i.d. DML baseline).

graphdml.PrecomputedFeatures

Bases: BaseEstimator

A fixed (n_nodes, n_features) matrix, e.g. graph embeddings computed elsewhere.

The features must be functions of covariates and graph structure only, never of treatments or outcomes.

Diagnostics and checks

graphdml.placebo_test(estimator, data, n_permutations=5, alpha=0.05, random_state=0)

Refit with randomly permuted treatments; every effect should be ~0.

This checks the estimation pipeline (e.g. leakage, a broken featurizer), not the causal assumptions. Uses a Bonferroni correction across permutations and effects.

graphdml.negative_control_test(estimator, data, outcome, alpha=0.05)

Refit with an outcome the treatment cannot have affected; every effect should be ~0.

Good negative controls share the confounders of the real outcome but are causally unaffected by treatment, e.g. the same outcome measured before treatment.

graphdml.two_hop_test(estimator, data, alpha=0.05)

Test whether two-hop neighbors' treatments matter once one-hop exposure is modelled.

The refit uses a focal set that is independent at the two-hop scale (pairwise distance

= 5), so it has less power than the main analysis.

graphdml.selftest

graphdml.selftest(): check on your own machine that the estimator finds real signals and does not invent false ones.

It repeatedly simulates two worlds with the same strong network confounding, one with planted effects and one with none (the null twin), fits graphdml, and compares the estimates with the known truth. Methods that ignore the network are shown alongside, to show what the confounding would do to them.

SelfTestResult dataclass

Outcome of :func:selftest. passed is True when every check below holds.

selftest(n_reps=30, n=6000, *, direct_effect=1.0, peer_effect=0.6, estimator=None, alpha=0.05, n_jobs=-1, random_state=0)

Run the signal check and report whether graphdml passed.

Parameters:

Name	Type	Description	Default
n_reps	int	Repetitions per world; each fits one model on a fresh n-node graph.	30
n	int	Nodes per simulated graph.	6000
direct_effect	float	The effects planted in the signal world (the null world has none).	1.0
peer_effect	float	The effects planted in the signal world (the null world has none).	1.0
estimator	GraphDML	An unfitted estimator to test (default: GraphDML(exposure="mean") with the library's default learners). Pass your own to test your settings.	None
alpha	float	Significance level used for detection and false alarms.	0.05
n_jobs	int	Parallel workers (joblib).	-1
random_state	int	Seeds the repetitions.	0

Returns:

Type	Description
SelfTestResult	passed, a results table and the individual checks.

Notes

The criteria are set from statistics, not tuned to results. The bands use 3 Monte Carlo standard errors, so they loosen when n_reps is small.

• coverage: the share of 95% intervals containing the truth is at least 0.95 - 3 * SE in the planted world, for every effect.
• power: the planted direct and peer effects are detected in at least 80% of runs, with the correct sign.
• false alarms: in the null world, the share of runs that claim an effect is at most alpha + 3 * SE, for every effect.

Datasets

graphdml.datasets.GraphDataset dataclass

A ready-to-fit :class:~graphdml.GraphData plus its ground truth and story.

Attributes:

Name	Type	Description
name	str
data	GraphData
truth	dict	{"ade": ..., "ape": ...}: the true average direct and peer effects.
exposure	str	The exposure map the data were generated with.
treatment_name, outcome_name	str	What T and Y mean in the story.
DESCR	str	The story card: what the variables are, what confounds, what naive methods miss.
node_effects	dict of arrays	Per-node effects ("ade", "ape") when effects vary across nodes.
extras	dict	Anything else: oracle nuisances, negative-control outcomes, positions, latents.

focal_truth(nodes)

Average true effects over a subset of nodes (e.g. model.focal_nodes_).

Equals :attr:truth when effects are the same for every node.

graphdml.datasets.make_flu_town(n_residents=5000, *, avg_neighbors=6.0, confounding=1.0, direct_effect=-2.0, peer_effect=-0.5, random_state=0)

Flu shots in a town: direct protection, herd protection and network confounding.

Treatment flu_shot (binary), outcome sick_days this season, exposure "sum" (number of vaccinated neighbors). Health-conscious residents cluster together, and neighbors' health-consciousness drives both vaccination and sickness. Naive regression overstates how much the shot protects you. See DESCR for the full story.

graphdml.datasets.make_referral_app(n_users=8000, *, links_per_user=3, confounding=1.0, direct_effect=12.0, peer_effect=40.0, random_state=0)

Promo coupons in a social app: targeting on friends' engagement, mean exposure.

Treatment coupon (binary), outcome monthly_spend in dollars, exposure "mean" (share of friends who got a coupon): the peer effect of 40 means +$4 for every additional 10% of friends with a coupon. The friendship graph has influencer hubs (preferential attachment), and the targeting model favours users with engaged friends.

graphdml.datasets.make_classroom_tutoring(n_classes=200, class_size=25, *, friends_per_student=6.0, confounding=1.0, direct_effect=1.5, peer_effect=0.4, heterogeneous=False, random_state=0)

Tutoring hours in classrooms: continuous treatment, communities, focal-set estimand.

Treatment tutoring_hours (continuous), outcome test_score, exposure "mean" (friends' average tutoring hours). Friendships form mostly within classes, and students with higher prior GPA have more friends. With heterogeneous=True tutoring helps lower-GPA students more, so the focal set (which favours low-degree students) has a larger average effect than the whole school: compare truth with focal_truth(model.focal_nodes_).

graphdml.datasets.make_homophily_trap(n_people=5000, *, avg_friends=6.0, homophily=3.0, direct_effect=1.0, random_state=0)

An online course and salary growth, where the true peer effect is zero.

Treatment took_course (binary), outcome salary_growth (%), exposure "mean". An unobserved trait ("ambition") drives enrolment, salary growth and who befriends whom. GraphDML reports a spurious peer effect; extras["negative_control"] (last year's salary growth, which the course cannot affect) exposes the problem via :func:graphdml.negative_control_test.

graphdml.datasets.make_signal_check(n=6000, *, direct_effect=1.0, peer_effect=0.6, confounding=1.0, avg_degree=4.0, noise_sd=2.0, random_state=0)

Planted direct and peer effects, hidden behind strong network confounding.

Set direct_effect=0, peer_effect=0 for the null twin: the same data-generating process with no true effect. Methods that ignore the network still "find" effects in the null twin, because the confounding is strong; a correct method should not.

T is binary. Both treatment and outcome depend nonlinearly on a node's own covariates and on its neighbors' (mean and maximum), with strength confounding. The exposure is "mean": the share of treated neighbors.

graphdml.datasets.make_lastfm_promo(*, confounding=1.0, direct_effect=2.0, peer_effect=3.0, noise_sd=3.0, random_state=0, data_home=None)

A concert promotion on the real LastFM Asia social network (semi-synthetic).

The network (7,624 users, 27,806 mutual follows) and covariates (activity, taste) are real; the promotion promo and outcome listening_hours are simulated, so the true effects are known. Exposure "mean": the share of friends who received the promo. The targeting rule and the outcome both depend nonlinearly on friends' taste.

graphdml.datasets.load_insurance_experiment(data_dir)

Load the insurance experiment as a :class:GraphDataset (no ground truth).

• Nodes: the 4,902 surveyed households. An edge j -> i means i named j as a friend (j can influence i).
• T = assigned to an intensive session (randomized); Y = bought insurance.
• Covariates: household characteristics (median-imputed, with missing indicators), the randomized design variables (second round, information treatments), the number of friends named, and administrative-village dummies.
• extras["exposure"]: the study's peer exposure, the share of named friends who attended a first-round intensive session (the denominator counts every named friend, surveyed or not, as in the study). It is linear in T, so it is given as a :class:~graphdml.MatrixExposure.
• extras also has second_round, info_none, natural_village (for clustered SEs), the study's own exposure variable for checking, and the published estimates.

Neo4j

graphdml.io.neo4j.Neo4jGraphSource

Bases: _Client

Read a graph from Neo4j into :class:~graphdml.GraphData.

Parameters:

Name	Type	Description	Default
driver	Driver	An existing driver; otherwise one is opened from uri/auth (or the NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD environment variables).	None
database	str		None

load(node_label, rel_type, *, features, treatment, outcome, id_property=None, influence='undirected')

Load all node_label nodes and the rel_type relationships among them.

influence says how a relationship (a)-[:REL]->(b) transmits effects: "undirected" (both ways, e.g. FRIENDS_WITH), "source_to_target" (a influences b) or "target_to_source" (b influences a, e.g. a FOLLOWS b). Node ids are id_property values if given (recommended; they should be unique), otherwise Neo4j element ids.

load_query(node_query, edge_query, *, features, treatment, outcome, influence='undirected', parameters=None)

Load from your own Cypher (e.g. one village, or a computed property).

node_query must return a column id plus the named property columns; edge_query must return source and target ids. Edges whose endpoints are not returned by the node query are dropped.

graphdml.io.neo4j.Neo4jResultWriter

Bases: _Client

Write a fitted :class:~graphdml.GraphDML back to Neo4j.

Each call creates one (:GDMLRun {name, ...estimates}) node and an (:GDMLRun)-[:ESTIMATED_ON {fold, t_hat, y_hat, res_t, ..., res_y}]->(n) relationship to every focal node, so runs can be compared, visualised (focal nodes are exactly the ones with an ESTIMATED_ON edge) and deleted without touching your data.

write(model, *, run_name, node_label, id_property=None, description='', dry_run=False, batch_size=1000)

Write the run. id_property must match what the data were loaded with.

Returns a summary dict. With dry_run=True nothing is written and the summary includes the first rows that would be.

list_runs(name=None)

All runs (or one by name), newest first.

delete_run(run_name)

Delete a run node and its relationships. Returns the number of runs deleted.

graphdml.io.neo4j.GDSEmbeddings

Bases: BaseEstimator

Featurizer: a base featurizer's output plus Neo4j GDS FastRP node embeddings.

FastRP embeddings summarise each node's position in the graph (structure only, never treatment or outcome), so they can stand in for unobserved structural confounders such as centrality or community. Rows are aligned to data.node_ids, which must be the id_property values used when loading.

Parameters:

Name	Type	Description	Default
source	Neo4jGraphSource		required
node_label	the same graph that was loaded		required
rel_type	the same graph that was loaded		required
id_property	the same graph that was loaded		required
dimension	embedding size		16
base	featurizer whose output is kept alongside the embeddings	(default :class:~graphdml.NeighborhoodFeatures; None for embeddings only)	'default'
random_seed	FastRP seed (embeddings are deterministic given the seed)		42