text: 
  import numpy as np
  import scipy.optimize as opt
  from tqdm.auto import trange
  from scipy import sparse as sp
  
  def has_method(o, name):
      return callable(getattr(o, name, None))
  
  def save_ensemble(ge, filepath, save_space=False):
  	if save_space:
  		from copy import copy
  		ge = copy(ge)
  		ge.adj_matrix=None
  		ge.sigma=None
  	ge.save(filepath)
  
  def load_ensemble(filepath):
  	import pickle
  	ge = pickle.load(open(filepath, 'rb'))
  	if ge.adj_matrix is None or ge.sigma is None:
  		ge.adj_matrix = ge.eval_adj_matrix(ge.theta)
  		ge.sigma = ge.eval_sigma(ge.theta)
  	return ge
  
  class GraphEnsemble:
  	def __init__(self, N_nodes, directed=False, self_loops=False, nodenames=None):
  		self.N = N_nodes
  		self.N1 = N_nodes
  		self.N2 = N_nodes
  		self.self_loops = self_loops
  		self.directed = directed
  		self.fixed_edges = None
  		self.nodenames = nodenames if nodenames is not None else list(range(N_nodes))
  
  	def fit(self, constraints, method='anderson', opt_kwargs=None, theta0=None):
  		self.N_theta = sum([len(c.c_vals) for c in constraints])
  		self.constraints = constraints
  		self.theta = self.eval_theta(method, opt_kwargs, theta0)
  		self.adj_matrix = self.eval_adj_matrix(self.theta)
  		self.sigma = self.eval_sigma(self.theta)
  
  	def predict_mean(self, obs, f_args=None):
  		if has_method(obs, 'func'): # obs is of class Observable
  			func = obs.func
  			f_args = obs.f_args
  		else:
  			func = obs
  		if f_args is None:
  			f_args = []
  		return func(self.adj_matrix, *f_args)
  
  	def predict_std(self, obs, g_args=None, obs_dim_idx=None, slice_param=None, batch_size=None):
  		if has_method(obs, 'grad'): # obs is of class Observable
  			if hasattr(obs, 'obs_dim_idx'):
  				obs_dim_idx = obs.obs_dim_idx
  			g_args = obs.g_args
  			func_grad = obs.grad
  		else:
  			func_grad = obs
  		if g_args is None:
  			g_args = []
  		if batch_size is None:
  			if obs.sparse:
  				grad_term = func_grad(self.adj_matrix, *g_args, sparse=True)
  				sum_dims = None if obs_dim_idx is None else 1-obs_dim_idx # scipy sparse does not support tuple axes
  				std_vec = np.squeeze(np.asarray(np.sqrt(grad_term.multiply(self.sigma).power(2).sum(axis=sum_dims))))
  				if obs_dim_idx is None:
  					std_vec = std_vec.item() # sparse returns a numpy scalar, so convert to float
  			else:
  				grad_term = func_grad(self.adj_matrix, *g_args)
  				sum_dims = tuple([dim for dim in range(len(grad_term.shape)) if dim != obs_dim_idx])
  				std_vec = np.sqrt(((self.sigma * grad_term) ** 2).sum(axis=sum_dims))
  		else:
  			if obs.sparse:
  				raise NotImplementedError('Sparse mode with batch_size calculation not implemented yet')
  			slice_param = obs.slice_param
  			std_vec = np.zeros(obs.obs_dim)
  			for b in trange(int(obs.obs_dim / batch_size)):
  				bslice = slice(b*batch_size,(b+1)*batch_size)
  				grad_term = func_grad(self.adj_matrix, *g_args, {slice_param:bslice})
  				std_vec[bslice] = np.sqrt(((self.sigma[...,None] * grad_term) ** 2).sum(axis=(0,1)))
  		return std_vec
  
  	def predict_zscore(self, value, obs, obs_grad=None, f_args=None, g_args=None, batch_size=None):
  		if has_method(obs, 'func') and has_method(obs, 'grad'): # obs is of class Observable
  			mu = self.predict_mean(obs)
  			sigma = self.predict_std(obs, batch_size=batch_size)
  		else:
  			mu = self.predict_mean(obs, f_args=f_args)
  			sigma = self.predict_std(obs_grad, g_args=g_args, batch_size=batch_size)
  		return (value-mu)/sigma
  
  	def fix_edges_value(self, edgelist, values):
  		el = np.asarray(edgelist)
  		if hasattr(values, '__len__'):
  			values = np.asarray(values)
  		else:
  			values = values * np.ones(el.shape[0])
  		self.fixed_edges = np.concatenate([el, values[:,None]],axis=1)
  
  	def sample(self):
  		if self.directed:
  			s = (np.random.rand(self.N, self.N) <= self.adj_matrix).astype(int)
  			return s
  		else:
  			s = np.triu((np.random.rand(self.N, self.N) <= self.adj_matrix).astype(int), 1)
  			return s + s.T
  
  	def eval_coupling_matrix(self, theta):
  		coupling = np.ones([self.N, self.N])
  		for c_idx, c in enumerate(self.constraints):
  			i0, i1 = self.get_multipliers_idx(c_idx)
  			coupling *= c.coupling_matrix(theta[i0:i1], self)
  		return coupling
  
  	def eval_adj_matrix(self, theta): # Fermi-Dirac distribution
  		coupling = self.eval_coupling_matrix(theta)
  		adj_matrix = (coupling / (1 + coupling))
  		if not self.self_loops:
  			np.fill_diagonal(adj_matrix, 0.)
  		if self.fixed_edges is not None:
  			adj_matrix[self.fixed_edges[:,0].astype(int),self.fixed_edges[:,1].astype(int)] = self.fixed_edges[:,2]
  			if not self.directed:
  				adj_matrix[self.fixed_edges[:, 1].astype(int), self.fixed_edges[:, 0].astype(int)] = self.fixed_edges[:, 2]
  		return adj_matrix
  
  	def eval_sigma(self, theta):
  		coupling = self.eval_coupling_matrix(theta)
  		sigma = np.sqrt(coupling) / (1 + coupling)
  		if not self.self_loops:
  			np.fill_diagonal(sigma, 0)
  		return sigma
  
  	def eval_ml_eqs(self, theta):
  		adj = self.eval_adj_matrix(theta)
  		errors = np.array([])
  		for c_idx, c in enumerate(self.constraints):
  			errors = np.concatenate([errors, c.eval_ml_eqs(adj, self)])
  		return errors
  
  	def eval_theta(self, method='anderson', opt_kwargs=None, theta0=None):
  		if theta0 is None:
  			theta0 = np.ones(self.N_theta)
  		return opt.root(self.eval_ml_eqs,
  						x0=theta0,
  						method=method,
  						options=opt_kwargs
  						).x
  
  	def get_multipliers_idx(self, c_idx):
  		i0 = sum([len(c.c_vals) for c in self.constraints[:c_idx]])
  		i1 = i0 + len(self.constraints[c_idx].c_vals)
  		return (i0, i1)
  
  	def check_constraint_deviations(self, c_idx):
  		return self.constraints[c_idx].eval_ml_eqs(self.adj_matrix, self)
  
  	def get_graph_dims(self):
  		return np.array([self.N,self.N])
  
  	def save(self, filepath):
  		import pickle
  		pickle.dump(self, open(filepath,'wb'), pickle.HIGHEST_PROTOCOL)
  
  class BipartiteGraphEnsemble(GraphEnsemble):
  	def __init__(self, N1_nodes, N2_nodes):
  		self.N1 = N1_nodes
  		self.N2 = N2_nodes
  		self.fixed_edges = None
  
  	def eval_theta(self, method='anderson', opt_kwargs=None, theta0=None):
  		if theta0 is None:
  			theta0 = np.ones(self.N_theta)
  		return opt.root(self.eval_ml_eqs,
  						x0=theta0,
  						method=method,
  						options=opt_kwargs
  						).x
  
  	def eval_coupling_matrix(self, theta):
  		coupling = np.ones([self.N1, self.N2])
  		for c_idx, c in enumerate(self.constraints):
  			i0, i1 = self.get_multipliers_idx(c_idx)
  			coupling *= c.coupling_matrix(theta[i0:i1], self)
  		return coupling
  
  	def eval_adj_matrix(self, theta): # Fermi-Dirac distribution
  		coupling = self.eval_coupling_matrix(theta)
  		adj_matrix = (coupling / (1 + coupling))
  		if self.fixed_edges is not None:
  			adj_matrix[self.fixed_edges[:,0].astype(int),self.fixed_edges[:,1].astype(int)] = self.fixed_edges[:,2]
  		return adj_matrix
  
  	def eval_sigma(self, theta):
  		coupling = self.eval_coupling_matrix(theta)
  		sigma = np.sqrt(coupling) / (1 + coupling)
  		return sigma
  
  	def predict_std(self, obs, g_args=None, obs_dim_idx=None, batch_size=None, obs_dim=None):
  		if has_method(obs, 'grad'): # obs is of class Observable
  			if hasattr(obs, 'obs_dim_idx'):
  				obs_dim_idx = obs.obs_dim_idx
  			g_args = obs.g_args
  			func_grad = obs.grad
  		else:
  			func_grad = obs
  		if g_args is None:
  			g_args = []
  		if batch_size is None:
  			grad_term = func_grad(self.adj_matrix, *g_args)
  			sum_dims = tuple([dim for dim in range(len(grad_term.shape)) if dim != obs_dim_idx])
  			std_vec = np.sqrt(((self.sigma * grad_term) ** 2).sum(axis=sum_dims))
  		else:
  			std_vec = np.zeros(obs_dim)
  			for b in trange(int(obs_dim / batch_size)):
  				bslice = slice(b*batch_size,(b+1)*batch_size)
  				grad_term = func_grad(self.adj_matrix, *g_args, bslice=bslice)
  				std_vec[bslice] = np.sqrt(((self.sigma[...,None] * grad_term) ** 2).sum(axis=(0,1)))
  		if hasattr(obs, 'output_nodeset') and obs.output_nodeset is not None:
  			std_vec = std_vec[obs.output_nodeset]
  		return std_vec
  
  	def sample(self):
  		s = (np.random.rand(self.N1, self.N2) <= self.adj_matrix).astype(int)
  		return s
  
  	def get_graph_dims(self):
  		return np.array([self.N1,self.N2])
  
  class MultiGraphEnsemble(GraphEnsemble):
  
  	def eval_adj_matrix(self, theta): # Bose-Einstein condensate
  		coupling = self.eval_coupling_matrix(theta)
  		adj_matrix = (coupling / (1 - coupling))
  		if not self.self_loops:
  			np.fill_diagonal(adj_matrix, 0.)
  		return adj_matrix
  
  	def eval_sigma(self, theta):
  		coupling = self.eval_coupling_matrix(theta)
  		sigma = np.sqrt(coupling) / (1 - coupling)
  		if not self.self_loops:
  			np.fill_diagonal(sigma, 0)
  		return sigma
  
  	def sample(self):
  		raise NotImplementedError
  
  class GraphEnsembleSet:
  	def __init__(self, ge_list, nid_pair_list):
  		if not self.check_consistency(ge_list, nid_pair_list):
  			raise ValueError('Node ids are not associated to consistent GraphEnsemble dimensions')
  		self.ge_list = ge_list
  		self.nid_pair_list = nid_pair_list
  		self.nidpair2geid = {(nid_pair_list[i][0],nid_pair_list[i][1]):i for i in range(len(ge_list))}
  
  	def __getitem__(self, nid_pair):
  		nid1,nid2 = nid_pair
  		return self.ge_list[self.nidpair2geid[(nid1,nid2)]]
  
  	def __len__(self):
  		return len(self.ge_list)
  
  	def fit(self, constraints_list, method='anderson', opt_kwargs=None, theta0=None):
  		if isinstance(method, str):
  			method = [method for i in range(len(self.ge_list))]
  		if isinstance(opt_kwargs, dict):
  			opt_kwargs = [opt_kwargs for i in range(len(self.ge_list))]
  		if not isinstance(theta0, list):
  			theta0 = [theta0 for i in range(len(self.ge_list))]
  
  		for i in range(len(self.ge_list)):
  			self.ge_list[i].fit(constraints_list[i], method=method[i], opt_kwargs=opt_kwargs[i], theta0=theta0[i])
  
  	def predict_mean(self, obs):
  		adj_list = [self[obs.nid_pair_list[i]].adj_matrix for i in range(len(obs))]
  		return obs.func(adj_list)
  
  	def predict_std(self, obs):
  		adj_list = [self[obs.nid_pair_list[i]].adj_matrix for i in range(len(obs))]
  		std = 0
  		for i in range(len(obs)):
  			grad_term = obs.grad(adj_list, term_num=i)
  			std = std + ((self[obs.nid_pair_list[i]].sigma * grad_term) ** 2).sum()
  		return np.sqrt(std)
  
  
  	def predict_mean_list(self, obs_list, f_args_list):
  		return [self.ge_list[i].predict_mean(obs_list[i], f_args_list[i]) for i in range(len(self.ge_list))]
  
  	def check_consistency(self, ge_list, nid_pair_list):
  		nid2dim = {}
  		for i,(nid1,nid2) in enumerate(nid_pair_list):
  			ge = ge_list[i]
  			for j,nid in enumerate([nid1,nid2]):
  				if nid not in nid2dim:
  					nid2dim[nid] = ge.get_graph_dims()[j]
  				else:
  					if nid2dim[nid] != ge.get_graph_dims()[j]:
  						return False
  		return True
  
  