"""
Some common graphs and lattices
"""
# pylint: disable=invalid-name
from functools import partial
from typing import Any, Optional, Sequence, Tuple
import networkx as nx
Graph = Any
[docs]def Line1D(
n: int,
node_weight: Optional[Sequence[float]] = None,
edge_weight: Optional[Sequence[float]] = None,
pbc: bool = True,
) -> Graph:
"""
1D chain with ``n`` sites
:param n: [description]
:type n: int
:param pbc: [description], defaults to True
:type pbc: bool, optional
:return: [description]
:rtype: Graph
"""
g = nx.Graph()
if edge_weight is None:
edge_weight = 1.0 # type: ignore
if not isinstance(edge_weight, list):
edge_weight = [edge_weight for _ in range(n)] # type: ignore
if node_weight is None:
node_weight = 0.0 # type: ignore
if not isinstance(node_weight, list):
node_weight = [node_weight for _ in range(n)] # type: ignore
for i in range(n):
g.add_node(i, weight=node_weight[i])
for i in range(n - 1):
g.add_edge(i, i + 1, weight=edge_weight[i])
if pbc is True:
g.add_edge(n - 1, 0, weight=edge_weight[i])
return g
[docs]def Even1D(n: int, s: int = 0) -> Graph:
g = nx.Graph()
for i in range(n):
g.add_node(i, weight=1.0)
for i in range(s, n, 2):
g.add_edge(i, (i + 1) % n, weight=1.0)
return g
Odd1D = partial(Even1D, s=1)
[docs]class Grid2DCoord:
"""
Two-dimensional grid lattice
"""
[docs] def __init__(self, n: int, m: int):
"""
:param n: number of rows
:type n: int
:param m: number of cols
:type m: int
"""
# row first
self.m = m
self.n = n
self.mn = m * n
[docs] def one2two(self, i: int) -> Tuple[int, int]:
x = i // self.n
y = i % self.n
return x, y
[docs] def two2one(self, x: int, y: int) -> int:
return x * self.n + y
[docs] def all_rows(self, pbc: bool = False) -> Sequence[Tuple[int, int]]:
"""
return all row edge with 1d index encoding
:param pbc: whether to include pbc edges (periodic boundary condition),
defaults to False
:type pbc: bool, optional
:return: list of row edge
:rtype: Sequence[Tuple[int, int]]
"""
r = []
for i in range(self.mn):
if (i + 1) % self.n != 0:
r.append((i, i + 1))
elif pbc:
r.append((i, i - self.n + 1))
return r
[docs] def all_cols(self, pbc: bool = False) -> Sequence[Tuple[int, int]]:
"""
return all col edge with 1d index encoding
:param pbc: whether to include pbc edges (periodic boundary condition),
defaults to False
:type pbc: bool, optional
:return: list of col edge
:rtype: Sequence[Tuple[int, int]]
"""
r = []
for i in range(self.mn):
if i + self.n < self.mn:
r.append((i, i + self.n))
elif pbc:
r.append((i, i - (self.m - 1) * self.n))
return r
[docs] def lattice_graph(self, pbc: bool = True) -> Graph:
"""
Get the 2D grid lattice in ``nx.Graph`` format
:param pbc: whether to include pbc edges (periodic boundary condition),
defaults to True
:type pbc: bool, optional
:return: _description_
:rtype: Graph
"""
g = nx.Graph()
for i in range(self.mn):
g.add_node(i, weight=0)
for i in range(self.mn):
x, y = self.one2two(i)
if pbc is False and x - 1 < 0:
pass
else:
g.add_edge(i, self.two2one((x - 1) % self.m, y), weight=1)
if pbc is False and y - 1 < 0:
pass
else:
g.add_edge(i, self.two2one(x, (y - 1) % self.n), weight=1)
return g