abp/examples/lattice_demo.py

from abp.viz import VisibleGraphState
import numpy as np
import time
import itertools

square_unit_cell = (((0, 0), (0, 1)), ((0, 0), (1, 0)), ((1, 0), (1, 1)), ((0, 1), (1, 1)))
funny_unit_cell = (((0, 0), (0, 1)), ((0, 0), (1, 0)), ((1, 0), (1, 1)), ((0, 1), (1, 1)), ((0, 0), (.5, .5)))

def add_offset(vector, offset):
    """ Offset a vector in n-dimensional space """
    return tuple(v+o for v, o in zip(vector, offset))

def offset_unit_cell(unit_cell, offset):
    """ Offset a unit cell """
    return {(add_offset(a, offset), add_offset(b, offset)) for a, b in unit_cell}

def lattice(unit_cell, size):
    """ Generate a lattice from a unit cell """
    edges = set()
    for offset in itertools.product(*map(range, size)):
        edges |= offset_unit_cell(unit_cell, offset)

    nodes = set(itertools.chain(*edges))
    return nodes, edges

#s = VisibleGraphState()
nodes, edges = lattice(funny_unit_cell, (10, 10))

psi = VisibleGraphState()
for node in nodes:
    pos = {"x": node[0], "y": node[1], "z": 0}
    psi.add_node(str(node), meta={"position":pos})

for edge in edges:
    psi.add_edge(str(edge[0]), str(edge[1]))