How to use the netket.graph.CustomGraph function in NetKet

def test_is_connected():
    for i in range(5, 10):
        for j in range(i + 1, i * i):
            x = nx.dense_gnm_random_graph(i, j)
            y = nk.graph.CustomGraph(x.edges)
            if len(x) == len(
                set((i for (i, j) in x.edges)) | set((j for (i, j) in x.edges))
            ):
                assert y.is_connected == nx.is_connected(x)
            else:
                assert not nx.is_connected(x)

[7, 8],
    [8, 9],
    [9, 10],
    [10, 11],
    [11, 12],
    [12, 13],
    [13, 14],
    [14, 15],
    [15, 16],
    [16, 17],
    [17, 18],
    [18, 19],
    [19, 0],
]

g = nk.graph.CustomGraph(edges=edges)
hi = nk.hilbert.CustomHilbert(local_states=[-1, 1], graph=g)
ha = nk.operator.GraphOperator(
    hi, siteops=[sigmax], bondops=[mszsz], bondops_colors=[0]
)
operators["Graph Hamiltonian"] = ha

# Custom Hamiltonian
sx = [[0, 1], [1, 0]]
sy = [[0, 1.0j], [-1.0j, 0]]
sz = [[1, 0], [0, -1]]
g = nk.graph.CustomGraph(edges=[[i, i + 1] for i in range(20)])
hi = nk.hilbert.CustomHilbert(local_states=[1, -1], graph=g)

sx_hat = nk.operator.LocalOperator(hi, [sx] * 3, [[0], [1], [5]])
sy_hat = nk.operator.LocalOperator(hi, [sy] * 4, [[2], [3], [4], [9]])
szsz_hat = nk.operator.LocalOperator(hi, sz, [0]) * nk.operator.LocalOperator(

import networkx as nx
import numpy as np
import pytest
from pytest import approx
import os

herm_operators = {}
generic_operators = {}

# Custom Hamiltonian
sx = [[0, 1], [1, 0]]
sy = [[0, 1.0j], [-1.0j, 0]]
sz = [[1, 0], [0, -1]]
sm = [[0, 0], [1, 0]]
sp = [[0, 1], [0, 0]]
g = nk.graph.CustomGraph(edges=[[i, i + 1] for i in range(8)])
hi = nk.hilbert.CustomHilbert(local_states=[1, -1], graph=g)

sx_hat = nk.operator.LocalOperator(hi, [sx] * 3, [[0], [1], [4]])
sy_hat = nk.operator.LocalOperator(hi, [sy] * 4, [[1], [2], [3], [4]])
szsz_hat = nk.operator.LocalOperator(hi, sz, [0]) * nk.operator.LocalOperator(
    hi, sz, [1]
)
szsz_hat += nk.operator.LocalOperator(hi, sz, [4]) * nk.operator.LocalOperator(
    hi, sz, [5]
)
szsz_hat += nk.operator.LocalOperator(hi, sz, [6]) * nk.operator.LocalOperator(
    hi, sz, [8]
)
szsz_hat += nk.operator.LocalOperator(hi, sz, [7]) * nk.operator.LocalOperator(
    hi, sz, [0]
)

import netket as nk
import networkx as nx
import igraph as ig
import math

nxg = nx.star_graph(10)
graphs = [
    nk.graph.Hypercube(length=10, n_dim=1, pbc=True),
    nk.graph.Hypercube(length=4, n_dim=2, pbc=True),
    nk.graph.Hypercube(length=5, n_dim=1, pbc=False),
    nk.graph.CustomGraph(nxg.edges()),
    nk.graph.Lattice(
        basis_vectors=[[1.0, 0.0], [1.0 / 2.0, math.sqrt(3) / 2.0]],
        extent=[10, 10],
        pbc=[0, 0],
        atoms_coord=[[0, 0]],
    ),
    nk.graph.Lattice(
        basis_vectors=[[1.5, math.sqrt(3) / 2.0], [0, math.sqrt(3)]],
        extent=[3, 5],
        atoms_coord=[[0, 0], [1, 0]],
    ),
    nk.graph.Lattice(
        basis_vectors=[[2.0, 0.0], [1.0, math.sqrt(3)]],
        extent=[4, 4],
        atoms_coord=[[0, 0], [1.0 / 2.0, math.sqrt(3) / 2.0], [1.0, 0.0]],
    ),

from __future__ import print_function
import netket as nk
import networkx as nx

g = nk.graph.Hypercube(length=10, n_dim=1)
print(g.distances)
print(g.is_bipartite)
print(g.is_connected)

Gx = nx.star_graph(10)
g = nk.graph.CustomGraph(Gx.edges)
print(g.distances)
print(g.is_bipartite)
print(g.is_connected)

g = nk.graph.CustomGraph([[i, (i + 1) % 10] for i in range(10)])
print(g.is_connected)
print(g.distances)
print(g.n_sites)

edge_colors.append([w + L2 * h, w + 1, 2])
            if h > 0:
                edge_colors.append([w + L2 * h, w + 1 + L2 * (h - 1), 2])
            elif pbc:
                edge_colors.append([w + L2 * h, w + 1 + L2 * (L1 - 1), 2])
        w = L2 - 1
        if pbc:
            edge_colors.append([L2 - 1 + L2 * h, L2 * h, 1])
            edge_colors.append([w + L2 * h, L2 * ((h + 1) % L1), 2])
            edge_colors.append([w + L2 * h, L2 * ((L1 + h - 1) % L1), 2])
        if h < L1 - 1:
            edge_colors.append([w + L2 * h, w + L2 * (h + 1), 1])
        elif pbc:
            edge_colors.append([w + L2 * h, w, 1])

    g = netket.graph.CustomGraph(edge_colors)
    if L1 * L2 % 2 == 0:
        hi = netket.hilbert.Spin(s=0.5, total_sz=0.0, graph=g)
    else:
        hi = netket.hilbert.Spin(s=0.5,  graph=g)
    sigmaz = [[1, 0], [0, -1]]
    sigmax = [[0, 1], [1, 0]]
    sigmay = [[0, -1j], [1j, 0]]

    interaction = numpy.kron(sigmaz, sigmaz) + numpy.kron(sigmax, sigmax) + numpy.kron(sigmay, sigmay)

    bond_operator = [
        (J[0] * interaction).tolist(),
        (J[1] * interaction).tolist(),
    ]

    bond_color = [1, 2]

def total_spin_netket_operator(hilbert_state_shape):
    edge_colors = []
    for i in range(numpy.prod(hilbert_state_shape)):
        edge_colors.append([i, i, 1])

    g = netket.graph.CustomGraph(edge_colors)
    hi = netket.hilbert.Spin(s=0.5, graph=g)

    sigmaz = [[1, 0], [0, -1]]
    sigmax = [[0, 1], [1, 0]]
    sigmay = [[0, -1j], [1j, 0]]

    interaction = numpy.kron(sigmaz, sigmaz) + numpy.kron(sigmax, sigmax) + numpy.kron(sigmay, sigmay)

    bond_operator = [
        (interaction).tolist(),
    ]

    bond_color = [1]

    return netket.operator.GraphOperator(hi, bondops=bond_operator, bondops_colors=bond_color)

# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function
import netket as nk
import networkx as nx

g = nk.graph.Hypercube(length=10, n_dim=1)
print(g.distances)
print(g.is_bipartite)
print(g.is_connected)

Gx = nx.star_graph(10)
g = nk.graph.CustomGraph(Gx.edges)
print(g.distances)
print(g.is_bipartite)
print(g.is_connected)

g = nk.graph.CustomGraph([[i, (i + 1) % 10] for i in range(10)])
print(g.is_connected)
print(g.distances)
print(g.n_sites)