dibs.kernel module

class dibs.kernel.AdditiveFrobeniusSEKernel(*, h=20.0, scale=1.0)[source]

Squared exponential kernel defined as

\(k(Z, Z') = \text{scale} \cdot \exp(- \frac{1}{h} ||Z - Z'||^2_F )\)

Parameters

h (float) – bandwidth parameter
scale (float) – scale parameter

eval(*, x, y)[source]

Evaluates kernel function

Parameters

x (ndarray) – any shape [...]
y (ndarray) – any shape [...], but same as x

Returns

kernel value of shape [1,]

class dibs.kernel.JointAdditiveFrobeniusSEKernel(*, h_latent=5.0, h_theta=500.0, scale_latent=1.0, scale_theta=1.0)[source]

Squared exponential kernel defined as

\(k((Z, \Theta), (Z', \Theta')) = \text{scale}_z \cdot \exp(- \frac{1}{h_z} ||Z - Z'||^2_F ) + \text{scale}_{\theta} \cdot \exp(- \frac{1}{h_{\theta}} ||\Theta - \Theta'||^2_F )\)

Parameters

h_latent (float) – bandwidth parameter for \(Z\) term
h_theta (float) – bandwidth parameter for \(\Theta\) term
scale_latent (float) – scale parameter for \(Z\) term
scale_theta (float) – scale parameter for \(\Theta\) term

eval(*, x_latent, x_theta, y_latent, y_theta)[source]

Evaluates kernel function k(x, y)

Parameters

x_latent (ndarray) – any shape [...]
x_theta (Any) – any PyTree of jnp.array tensors
y_latent (ndarray) – any shape [...], but same as x_latent
y_theta (Any) – any PyTree of jnp.array tensors, but same as x_theta

Returns

kernel value of shape [1,]

dibs.metrics module

class dibs.metrics.ParticleDistribution(logp: Any, g: Any, theta: Optional[Any] = None)[source]

NamedTuple for structuring sampled particles \((G, \Theta)\) (or \(G\)) and their assigned log probabilities

Parameters

logp (ndarray) – vector of log probabilities or weights of shape [M, ]
g (ndarray) – batch of graph adjacency matrix of shape [M, d, d]
theta (ndarray) – batch of parameter PyTrees with leading dimension M

g: Any: Alias for field number 1

logp: Any: Alias for field number 0

theta: Any: Alias for field number 2

dibs.metrics.expected_edges(*, dist)[source]

Computes expected number of edges, defined as

\(\text{expected edges}(p) := \sum_G p(G | D) |\text{edges}(G)|\)

Parameters: dist (dibs.metrics.ParticleDistribution) – particle distribution
Returns: expected number of edges [1, ]

dibs.metrics.expected_shd(*, dist, g)[source]

Computes expected structural hamming distance metric, defined as

\(\text{expected SHD}(p, G^*) := \sum_G p(G | D) \text{SHD}(G, G^*)\)

Parameters

dist (dibs.metrics.ParticleDistribution) – particle distribution
g (ndarray) – ground truth adjacency matrix of shape [d, d]

Returns

expected SHD [1, ]

dibs.metrics.neg_ave_log_likelihood(*, dist, eltwise_log_likelihood, x)[source]

Computes neg. ave log likelihood for a joint posterior over \((G, \Theta)\), defined as

\(\text{neg. LL}(p, G^*) := - \sum_G \int_{\Theta} p(G, \Theta | D) p(D^{\text{test}} | G, \Theta)\)

Parameters

dist (dibs.metrics.ParticleDistribution) – particle distribution
eltwise_log_likelihood (callable) – function evaluting the log likelihood \(p(D | G, \Theta)\) for a batch of graph samples given a data set of held-out observations; must satisfy the signature [:, d, d], PyTree(leading dim :), [N, d] -> [:,]
x (ndarray) – held-out observations of shape [N, d]

Returns

neg. ave log likelihood metric of shape [1,]

dibs.metrics.neg_ave_log_marginal_likelihood(*, dist, eltwise_log_marginal_likelihood, x)[source]

Computes neg. ave log marginal likelihood for a marginal posterior over \(G\), defined as

\(\text{neg. MLL}(p, G^*) := - \sum_G p(G | D) p(D^{\text{test}} | G)\)

Parameters

dist (dibs.metrics.ParticleDistribution) – particle distribution
eltwise_log_marginal_likelihood (callable) – function evaluting the marginal log likelihood \(p(D | G)\) for a batch of graph samples given a data set of held-out observations; must satisfy the signature [:, d, d], [N, d] -> [:,]
x (ndarray) – held-out observations of shape [N, d]

Returns

neg. ave log marginal likelihood metric of shape [1,]

dibs.metrics.pairwise_structural_hamming_distance(*, x, y)[source]

Computes pairwise Structural Hamming distance, i.e. the number of edge insertions, deletions or flips in order to transform one graph to another This means, edge reversals do not double count, and that getting an undirected edge wrong only counts 1

Parameters

x (ndarray) – batch of adjacency matrices [N, d, d]
y (ndarray) – batch of adjacency matrices [M, d, d]

Returns

matrix of shape [N, M] where elt i,j is SHD(x[i], y[j])

dibs.metrics.threshold_metrics(*, dist, g)[source]

Computes various threshold metrics (e.g. ROC, precision-recall, …)

Parameters

dist (dibs.metrics.ParticleDistribution) – sampled particle distribution
g (ndarray) – ground truth adjacency matrix of shape [d, d]

Returns

dict of metrics

dibs.target module

class dibs.target.Data(passed_key: Any, n_vars: int, n_observations: int, n_ho_observations: int, g: Any, theta: Any, x: Any, x_ho: Any, x_interv: Any)[source]

NamedTuple for structuring simulated synthetic data and their ground truth generative model

Parameters

passed_key (ndarray) – jax.random key passed into the function generating this object
n_vars (int) – number of variables in model
n_observations (int) – number of observations in x and used to perform inference
n_ho_observations (int) – number of held-out observations in x_ho and elements of x_interv used for evaluation
g (ndarray) – ground truth DAG
theta (Any) – ground truth parameters
x (ndarray) – i.i.d observations from the model of shape [n_observations, n_vars]
x_ho (ndarray) – i.i.d observations from the model of shape [n_ho_observations, n_vars]
x_interv (list) – list of (interv dict, i.i.d observations)

g: Any: Alias for field number 4

n_ho_observations: int: Alias for field number 3

n_observations: int: Alias for field number 2

n_vars: int: Alias for field number 1

passed_key: Any: Alias for field number 0

theta: Any: Alias for field number 5

x: Any: Alias for field number 6

x_ho: Any: Alias for field number 7

x_interv: Any: Alias for field number 8

dibs.target.make_graph_model(*, n_vars, graph_prior_str, edges_per_node=2)[source]

Instantiates graph model

Parameters

n_vars (int) – number of variables in graph
graph_prior_str (str) – specifier for random graph model; choices: er, sf
edges_per_node (int) – number of edges per node (in expectation when applicable)

Returns

Object representing graph model. For example ErdosReniDAGDistribution or ScaleFreeDAGDistribution

dibs.target.make_linear_gaussian_equivalent_model(*, key, n_vars=20, graph_prior_str='sf', bge_mean_obs=None, bge_alpha_mu=None, bge_alpha_lambd=None, obs_noise=0.1, mean_edge=0.0, sig_edge=1.0, min_edge=0.5, n_observations=100, n_ho_observations=100)[source]

Samples a synthetic linear Gaussian BN instance with Bayesian Gaussian equivalent (BGe) marginal likelihood as inference model to weight each DAG in an MEC equally

By marginalizing out the parameters, the BGe model does not allow inferring the parameters \(\Theta\).

Parameters

key (ndarray) – rng key
n_vars (int) – number of variables i
n_observations (int) – number of iid observations of variables
n_ho_observations (int) – number of iid held-out observations of variables
graph_prior_str (str) – graph prior (er or sf)
bge_mean_obs (float) – BGe score prior mean parameter of Normal
bge_alpha_mu (float) – BGe score prior precision parameter of Normal
bge_alpha_lambd (float) – BGe score prior effective sample size (degrees of freedom parameter of Wishart)
obs_noise (float) – observation noise
mean_edge (float) – edge weight mean
sig_edge (float) – edge weight stddev
min_edge (float) – min edge weight enforced by constant shift of sampled parameter

Returns

BGe inference model and observations from a linear Gaussian generative process

Return type

tuple(BGe, Data)

dibs.target.make_linear_gaussian_model(*, key, n_vars=20, graph_prior_str='sf', obs_noise=0.1, mean_edge=0.0, sig_edge=1.0, min_edge=0.5, n_observations=100, n_ho_observations=100)[source]

Samples a synthetic linear Gaussian BN instance

Parameters

key (ndarray) – rng key
n_vars (int) – number of variables
n_observations (int) – number of iid observations of variables
n_ho_observations (int) – number of iid held-out observations of variables
graph_prior_str (str) – graph prior (er or sf)
obs_noise (float) – observation noise
mean_edge (float) – edge weight mean
sig_edge (float) – edge weight stddev
min_edge (float) – min edge weight enforced by constant shift of sampled parameter

Returns

linear Gaussian inference model and observations from a linear Gaussian generative process

Return type

tuple(LinearGaussian, Data)

dibs.target.make_nonlinear_gaussian_model(*, key, n_vars=20, graph_prior_str='sf', obs_noise=0.1, sig_param=1.0, hidden_layers=(5,), n_observations=100, n_ho_observations=100)[source]

Samples a synthetic nonlinear Gaussian BN instance where the local conditional distributions are parameterized by fully-connected neural networks.

Parameters

key (ndarray) – rng key
n_vars (int) – number of variables
n_observations (int) – number of iid observations of variables
n_ho_observations (int) – number of iid held-out observations of variables
graph_prior_str (str) – graph prior (er or sf)
obs_noise (float) – observation noise
sig_param (float) – stddev of the BN parameters, i.e. here the neural net weights and biases
hidden_layers (tuple) – list of ints specifying the hidden layer (sizes) of the neural nets parameterizatin the local condtitionals

Returns

nonlinear Gaussian inference model and observations from a nonlinear Gaussian generative process

Return type

tuple(DenseNonlinearGaussian, Target)

dibs.target.make_synthetic_bayes_net(*, key, n_vars, graph_model, generative_model, n_observations=100, n_ho_observations=100, n_intervention_sets=10, perc_intervened=0.1)[source]

Returns an instance of Target for evaluation of a method on a ground truth synthetic causal Bayesian network

Parameters

key (ndarray) – rng key
n_vars (int) – number of variables
graph_model (Any) – graph model object. For example: ErdosReniDAGDistribution
generative_model (Any) – BN model object for generating the observations. For example: LinearGaussian
n_observations (int) – number of observations generated for posterior inference
n_ho_observations (int) – number of held-out observations generated for evaluation
n_intervention_sets (int) – number of different interventions considered overall for generating interventional data
perc_intervened (float) – percentage of nodes intervened upon (clipped to 0) in an intervention.

Returns

synthetic ground truth generative DAG and parameters as well observations sampled from the model

Return type

Data

dibs.graph_utils module

dibs.graph_utils.acyclic_constr_nograd(mat, n_vars)[source]

Differentiable acyclicity constraint from Yu et al. (2019) http://proceedings.mlr.press/v97/yu19a/yu19a.pdf

Parameters

mat (ndarray) – graph adjacency matrix of shape [n_vars, n_vars]
n_vars (int) – number of variables, to allow for jax.jit-compilation

Returns

constraint value [1, ]

dibs.graph_utils.adjmat_to_str(mat, max_len=40)[source]

Converts binary adjacency matrix to human-readable string

Parameters

mat (ndarray) – graph adjacency matrix
max_len (int) – maximum length of string

Returns

human readable description of edges in adjacency matrix

Return type

str

dibs.graph_utils.elwise_acyclic_constr_nograd(mat, n_vars)

Vectorized version of acyclic_constr_nograd. Takes similar arguments as acyclic_constr_nograd but with additional array axes over which acyclic_constr_nograd is mapped.

Original documentation:

Differentiable acyclicity constraint from Yu et al. (2019) http://proceedings.mlr.press/v97/yu19a/yu19a.pdf

Args:
mat (ndarray): graph adjacency matrix of shape [n_vars, n_vars] n_vars (int): number of variables, to allow for jax.jit-compilation

Returns:
constraint value [1, ]

dibs.graph_utils.graph_to_mat(g)[source]

Returns adjacency matrix of ig.Graph object

Parameters: g (igraph.Graph) – graph
Returns: adjacency matrix
Return type: ndarray

dibs.graph_utils.mat_is_dag(mat)[source]

Returns True iff adjacency matrix represents a DAG

Parameters: mat (ndarray) – graph adjacency matrix
Returns: True iff mat represents a DAG
Return type: bool

dibs.graph_utils.mat_to_graph(mat)[source]

Returns ig.Graph object for adjacency matrix

Parameters: mat (ndarray) – adjacency matrix
Returns: graph
Return type: igraph.Graph