Quantum mechanics on graphs

physgraph.py

"""Quantum mechanics on graphs.

Todo:

* Implement dirichlet boundary conditions.
* Use node attributes to label boundary points and impose dirichlet boundary
  conditions.
* Develop plotting/viz routines using matplotlib/mayavi2.
* Use node attributes to store the physics coordinates of the point.
* Write a bisect function to put additional nodes between existing ones.
* Use node attributes to attache the value of the the potential at the node.
* Write code for generating the initial condition vetor for time dependent
  simulations.
* Write general code that implements the Crank-Nicholson method.
* Explore other sparse formats for implementing boundary conditions.
* Figure out more general, R-matrix type boundary conditions.
"""

import networkx as nx
import numpy as np
import scipy.sparse
from sympy import Matrix, sympify, S, symbols
from sympy.concrete.summations import summation
from sympy.physics import secondquant as sq


def sparse_laplacian(g, format='csr'):
    """Construct the Laplacian matrix for the graph as s scipy.sparse matrix."""
    o = g.order()
    A = nx.to_scipy_sparse_matrix(g)
    diag = A.sum(axis=1)
    diag.shape = (1,o)
    D = scipy.sparse.spdiags(diag, (0,), o, o, format=format)
    return D - A


def kinetic(g, sparse=False):
    """Construct the kinetic energy matrix for the graph.

    This is -(1/2)*Laplacian. To impose Dirichlet boundary conditions at the
    ith node, we need to zero out the ith row and ith column of the adjacency
    matrix. I am not sure what the best way of doing that is. We may have to
    use a different sparse storage format. We can use the neighbor method
    of Graph and write our own laplacian function that looks at a custom
    boundary condition attribute.
    """
    if sparse:
        return 0.5*sparse_laplacian(g)
    else:
        return 0.5*nx.laplacian(g)


def potential(g, attr, sparse=True):
    o = g.order()
    diag = np.array([g.node[i][attr] for i in range(g.order())])
    if sparse:
        return scipy.sparse.spdiags(diag, (0,), o, o, format='csr')
    else:
        return np.diag(diag)


def hamiltonian(g, attr, sparse=True):
    return kinetic(g, sparse) + potential(g, attr, sparse)


def set_node_attr(g, name, values):
    for i in range(g.order()):
        g.node[i][name] = values[i]


def symbolic_laplacian(g):
    return Matrix(nx.laplacian(g).astype(int))


def symbolic_kinetic(g):
    return symbolic_laplacian(g)/S(2)


g = nx.grid_graph(dim=[2])

def hamiltonian(g, nparticles):
    T1 = symbolic_kinetic(g)
    nlevels = T1.rows
    V = symbols('V')
    T2 = summation(sq.Bd(n)*T1[n,m]*sq.B(m),(n,0,nlevels-1),(m,0,nlevels-1))
    V2 = summation(Bd(n)**2*B(n)**2,(n,0,nlevels-1))*V/2
    return T2 + V2