# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: claude
#     language: python
#     name: claude
# ---

import jax.numpy as jnp
from jax import grad, jit, vmap, jacfwd, jacrev
from jax import random
import jax.scipy.sparse as sp
import matplotlib.pyplot as plt
import matplotlib
from mpl_toolkits.axes_grid1 import make_axes_locatable
import networkx as nx
import pandas as pd
from tqdm.auto import tqdm, trange
from scipy.sparse import csc_matrix
from IPython.core.display import display, HTML
display(HTML("<style>.container { width:100% !important; }</style>"))

import claude as cl
from claude import observables as obs
from claude import constraints as const

# +

N = 1000
ba = nx.barabasi_albert_graph(N,30)
gn = nx.scale_free_graph(1000)
ba_adj = np.asarray(nx.adj_matrix(ba).todense())
gn_adj = np.asarray(nx.adj_matrix(gn).todense())
d = np.array(ba_adj.sum(axis=0)).flatten()
dinv = 1/d

din = np.array(gn_adj.sum(axis=0)).flatten()
dout = np.array(gn_adj.sum(axis=1)).flatten()

# -

def plot_adj(adj):
    plt.imshow(adj,cmap=matplotlib.cm.get_cmap('gray_r'),vmin=0,vmax=1)
    plt.colorbar()


# + [markdown] heading_collapsed=true
# # Connectivity

# + hidden=true
m = cl.GraphEnsemble(1000)
m.fit([const.Connectivity(343)], opt_kwargs={'nit':30,'fatol':1e-2,'disp':True})

# + hidden=true
plt.figure(figsize=[15,5])
plt.subplot(1,3,1)
plot_adj(ba_adj)
plt.subplot(1,3,2)
plot_adj(m.adj_matrix)
plt.subplot(1,3,3)
plt.hist([m.sample().sum() for i in trange(1000)])

# + [markdown] hidden=true
# ## Subgraph connectivity

# + hidden=true
m = cl.GraphEnsemble(1000)
m.fit([const.Connectivity(343, nodeset1=np.arange(100))], opt_kwargs={'nit':30,'fatol':1e-2,'disp':True})

# + hidden=true
plt.figure(figsize=[15,5])
plt.subplot(1,3,1)
plot_adj(ba_adj)
plt.subplot(1,3,2)
plot_adj(m.adj_matrix)
plt.subplot(1,3,3)
plt.hist([np.triu(m.sample()[:100,:100]).sum() for i in trange(1000)])

# + [markdown] hidden=true
# ## Connectivity between 2 groups of nodes

# + hidden=true
m = cl.GraphEnsemble(1000)
m.fit([const.Connectivity(343, nodeset1=np.arange(100), nodeset2=np.arange(150,200))], opt_kwargs={'nit':30,'fatol':1e-2,'disp':True})

# + hidden=true
plt.figure(figsize=[15,5])
plt.subplot(1,3,1)
plot_adj(ba_adj)
plt.subplot(1,3,2)
plot_adj(m.adj_matrix)
plt.subplot(1,3,3)
plt.hist([m.sample()[:100,150:200].sum() for i in trange(1000)])

# + [markdown] heading_collapsed=true
# # Fixed degree sequence

# + [markdown] hidden=true
# ### Full degree sequence

# + hidden=true
m = cl.GraphEnsemble(1000)
m.fit([const.DegreeSequence(d)], opt_kwargs={'nit':100,'fatol':1e-2,'disp':True})
#m.fit([const.DegreeSequence(get_din_dout(ba)[0]),const.Connectivity(250, nodeset1=np.arange(30), nodeset2=np.arange(40,50))], opt_kwargs={'nit':100,'fatol':1e-2})
#m.fit([const.DegreeSequence(d[:30], nodeset=np.arange(30))], opt_kwargs={'nit':30,'fatol':1e-2})

# + hidden=true
plt.figure(figsize=[15,5])
plt.subplot(1,3,1)
plot_adj(ba_adj)
plt.subplot(1,3,2)
plot_adj(m.adj_matrix)
plt.subplot(1,3,3)
plt.plot(d, m.adj_matrix.sum(axis=1),'.')
plt.plot(plt.xlim(),plt.xlim())

# + [markdown] hidden=true
# ### Partial degree sequence

# + hidden=true
m = cl.GraphEnsemble(1000)
#m.fit([const.DegreeSequence(get_din_dout(ba)[0]),const.Connectivity(250, nodeset1=np.arange(30), nodeset2=np.arange(40,50))], opt_kwargs={'nit':100,'fatol':1e-2})
m.fit([const.DegreeSequence(d[:50], nodeset=np.arange(50))], opt_kwargs={'nit':30,'fatol':1e-2,'disp':True})

# + hidden=true
plt.figure(figsize=[25,10])
plt.subplot(1,3,1)
plot_adj(ba_adj)
plt.subplot(1,3,2)
plot_adj(m.adj_matrix)
plt.subplot(1,3,3)
plt.plot(d, '.', markersize=15)
plt.plot(m.adj_matrix.sum(axis=1),'.',markersize=5)

# + [markdown] hidden=true
# ### Degree sequence in subgraph

# + hidden=true
m = cl.GraphEnsemble(1000)
#m.fit([const.DegreeSequence(get_din_dout(ba)[0]),const.Connectivity(250, nodeset1=np.arange(30), nodeset2=np.arange(40,50))], opt_kwargs={'nit':100,'fatol':1e-2})
m.fit([const.DegreeSequence(np.arange(50), nodeset=np.arange(50), subgraph_nodeset=np.arange(100))], opt_kwargs={'nit':100,'fatol':1e-2,'disp':True})
#m.fit([const.DegreeSequence(d[:50], nodeset=np.arange(50))], opt_kwargs={'nit':30,'fatol':1e-2})

# + hidden=true
plt.figure(figsize=[20,5])
plt.subplot(1,3,1)
plot_adj(m.adj_matrix)
plt.subplot(1,3,2)
plt.plot(m.adj_matrix[:50,:100].sum(axis=1),'.')
plt.subplot(1,3,3)
plt.plot(m.adj_matrix.sum(axis=1),'.')
plt.axvline(100,color='red')

# + [markdown] heading_collapsed=true
# # Joining constraints

# + hidden=true
m = cl.GraphEnsemble(1000)
m.fit([const.DegreeSequence(get_din_dout(ba)[0]),const.Connectivity(250, nodeset1=np.arange(30), nodeset2=np.arange(40,50))], opt_kwargs={'nit':100,'fatol':1e-2,'disp':True})

# + hidden=true
plt.figure(figsize=[20,5])
plt.subplot(1,3,1)
plot_adj(ba_adj)
plt.subplot(1,3,2)
plot_adj(m.adj_matrix)
plt.subplot(1,3,3)
plt.hist([m.sample()[:30,40:50].sum() for i in trange(1000)])
# -

# # Observables

# ### RWR

# +

x = np.zeros(N)
x[np.random.randint(0,N, int(N/100*20))] = 1
#x = np.random.rand(N)
x= x/x.sum()
x = jnp.asarray(x)

lambd = 0.3

def propagate(adj,x,lambd=lambd):
    p = adj / jnp.reshape(adj.sum(axis=0), [adj.shape[0],1])
    I = jnp.eye(p.shape[0])
    return lambd*x.dot(jnp.linalg.inv(I - (1-lambd)*p))

def propagate2(adj,x,i,lambd=lambd):
    p = adj / jnp.reshape(adj.sum(axis=0), [adj.shape[0],1])
    I = jnp.eye(p.shape[0])
    return lambd*x.dot(jnp.linalg.inv(I - (1-lambd)*p))[i]


# -

m = cl.GraphEnsemble(1000)
#m.fit([const.DegreeSequence(d)], opt_kwargs={'nit':100,'fatol':1e-2})
m.fit([const.DegreeSequence(d),const.Connectivity(250, nodeset1=np.arange(30), nodeset2=np.arange(40,50))], opt_kwargs={'nit':50,'fatol':1e-2,'disp':True})

k = obs.RandomWalkWithRestart.eval_transfer_matrix(ba_adj, lambd)

mu = m.predict_mean(obs.RandomWalkWithRestart.func, f_args=[x, lambd, None, k])

sigma = m.predict_std(obs.RandomWalkWithRestart.grad, f_args=[x, lambd, None, k], obs_dim_idx=1)

# +
N_samples = 100
y = propagate(ba_adj, x)
#samples = []
rdmy_ms = np.zeros([len(ba),N_samples])
for i in trange(N_samples):
    flag = False
    while not flag:
        smpl = m.sample()
        prop = propagate(smpl,x,lambd)
        if np.isnan(propagate(smpl,x)).sum() == 0:
            #samples.append(smpl)
            rdmy_ms[:,i] = propagate(smpl, x)
            flag = True
    
z = (y - rdmy_ms.mean(axis=1))/rdmy_ms.std(axis=1)
# -

plt.figure(figsize=[13,5])
plt.subplot(121)
plt.plot(rdmy_ms.mean(axis=1),mu,'.')
plt.xlabel('Numerical expected value',fontsize=15)
plt.ylabel('Predicted expected value',fontsize=15)
plt.plot(plt.xlim(),plt.xlim())
plt.subplot(122)
plt.plot(rdmy_ms.std(axis=1),sigma,'.')
plt.xlabel('Numerical std. dev.',fontsize=15)
plt.ylabel('Predicted std. dev.',fontsize=15)
plt.plot(plt.xlim(),plt.xlim())
plt.tight_layout()

# ### ANND

mu = m.predict_mean(obs.AverageNeighborDegree.func)

sigma = [m.predict_std(grad(obs.AverageNeighborDegree.func, argnums=0), f_args=[i]) for i in trange(N)]

# +
annd = obs.AverageNeighborDegree.func
N_samples = 100
y = np.asarray(annd(csc_matrix(adj))).flatten()
#samples = []
rdm_annd = np.zeros([N,N_samples])
for i in trange(N_samples):
    smpl = csc_matrix(m.sample())
    prop = np.asarray(annd(smpl)).flatten()
    rdm_annd[:,i] = prop
    
z = (y - rdm_annd.mean(axis=1))/rdm_annd.std(axis=1)
# -

plt.figure(figsize=[13,5])
plt.subplot(121)
plt.plot(rdm_annd.mean(axis=1),mu,'.')
plt.xlabel('Numerical expected value',fontsize=15)
plt.ylabel('Predicted expected value',fontsize=15)
plt.plot(plt.xlim(),plt.xlim())
plt.subplot(122)
plt.plot(rdm_annd.std(axis=1),sigma,'.')
plt.xlabel('Numerical std. dev.',fontsize=15)
plt.ylabel('Predicted std. dev.',fontsize=15)
plt.plot(plt.xlim(),plt.xlim())
plt.tight_layout()



StarcoderdataJupyterScriptsDedupFiltered