Commit 9ad36a5c authored by Joseph Weston's avatar Joseph Weston
Browse files

add magnetic gauge fixing for finite systems


Co-authored-by: Pablo Piskunow's avatarPablo Piskunow <pablo.perez.piskunow@gmail.com>
Co-authored-by: Dániel Varjas's avatarDaniel Varjas <dvarjas@gmail.com>
parent 77b2153f
#
# (2018) Modified by Kwant Authors
#
# Modifications
# =============
# Merged and modified from scipy/sparse/csgraph/_shortest_path.pyx
#
# All shortest path algorithms except for Dijkstra's removed.
# Implementation of Dijkstra's algorithm modified to allow for specific
# use-cases required by Flux. The changes are documented in the docstring.
#
#
# Copyright (c) 2001, 2002 Enthought, Inc.
# All rights reserved.
#
# Copyright (c) 2003-2017 SciPy Developers.
# All rights reserved.
#
# Copyright (c) 2011 Jake Vanderplas <vanderplas@astro.washington.edu>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# a. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
# b. Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution.
# c. Neither the name of Enthought nor the names of the SciPy Developers
# may be used to endorse or promote products derived from this software
# without specific prior written permission.
#
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS
# BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
# OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
import numpy as np
cimport numpy as np
from numpy.math cimport INFINITY as inf
from libc.stdlib cimport malloc, free
from libc.string cimport memset
ctypedef np.float64_t DTYPE_t
ctypedef np.int32_t ITYPE_t
ITYPE = np.int32
def dijkstra_directed(
object graph,
ITYPE_t[:] sources,
ITYPE_t[:] targets,
bint return_paths=True,
bint return_predecessors=False):
"""Modified directed Dijkstra algorithm.
Edges with infinite weight are treated as if they do not exist.
The shortest paths between edge pairs 'zip(sources, targets)'
are found.
If 'len(sources) == 1' then this routine can be used to
compute the one-to-all distances by passing an integer
greater than any node in 'graph' as the 'target'.
In this case 'return_predecessors' may be specified, and
the predecessor matrix will also be returned.
'return_predecessors' and 'return_paths' are mutually exclusive.
Returns
-------
if return_paths:
(paths, path_lengths)
elif return_predecessors:
(path_lengths, predecessors)
else:
path_lengths
"""
# Implementation of Dijkstra's algorithm modified to allow for
# early exit (when target is found) and return the path from source
# to target, rather than the whole predecessor matrix. In addition
# graph edges with infinite weight are treated as if they do not exist.
cdef ITYPE_t[:] csr_indices = graph.indices, csr_indptr = graph.indptr
cdef DTYPE_t[:] csr_weights = graph.data
# This implementation of Dijkstra's algorithm is very tightly coupled
# to our use-case in 'flux', so we allow ourselves to assert
assert sources.shape[0] == targets.shape[0]
assert graph.shape[0] == graph.shape[1]
assert not (return_predecessors and return_paths)
assert not (return_predecessors and (sources.shape[0] > 1))
cdef unsigned int num_links = sources.shape[0], num_nodes = graph.shape[0]
cdef unsigned int i, k, j_source, j_target, j_current
cdef ITYPE_t j
cdef DTYPE_t next_val
cdef FibonacciHeap heap
cdef FibonacciNode *v
cdef FibonacciNode *current_node
cdef FibonacciNode* nodes = <FibonacciNode*> malloc(num_nodes * sizeof(FibonacciNode))
cdef ITYPE_t[:] pred = np.empty((num_nodes,), dtype=ITYPE)
# outputs
cdef DTYPE_t[:] path_lengths = np.zeros((num_links,), float)
cdef list paths
if return_paths:
paths = []
for i in range(num_links):
j_source = sources[i]
j_target = targets[i]
for k in range(num_nodes):
initialize_node(&nodes[k], k)
pred[:] = -1 # only useful for debugging
heap.min_node = NULL
insert_node(&heap, &nodes[j_source])
while heap.min_node:
v = remove_min(&heap)
v.state = SCANNED
if v.index == j_target:
path_lengths[i] = v.val
if return_paths:
paths.append(_calculate_path(pred, j_source, j_target))
break # next iteration of outer 'for' loop
for j in range(csr_indptr[v.index], csr_indptr[v.index + 1]):
if csr_weights[j] == inf:
# Treat infinite weight links as missing
continue
j_current = csr_indices[j]
current_node = &nodes[j_current]
if current_node.state != SCANNED:
next_val = v.val + csr_weights[j]
if current_node.state == NOT_IN_HEAP:
current_node.state = IN_HEAP
current_node.val = next_val
insert_node(&heap, current_node)
pred[j_current] = v.index
elif current_node.val > next_val:
decrease_val(&heap, current_node,
next_val)
pred[j_current] = v.index
free(nodes)
if return_paths:
return paths, path_lengths
elif return_predecessors:
return path_lengths, pred
else:
return path_lengths
cdef list _calculate_path(ITYPE_t[:] pred, int j_source, int j_target):
visited = []
cdef int node = j_target
while node != j_source:
visited.append(node)
node = pred[node]
visited.append(j_source)
return visited
######################################################################
# FibonacciNode structure
# This structure and the operations on it are the nodes of the
# Fibonacci heap.
#
cdef enum FibonacciState:
SCANNED
NOT_IN_HEAP
IN_HEAP
cdef struct FibonacciNode:
unsigned int index
unsigned int rank
FibonacciState state
DTYPE_t val
FibonacciNode* parent
FibonacciNode* left_sibling
FibonacciNode* right_sibling
FibonacciNode* children
cdef void initialize_node(FibonacciNode* node,
unsigned int index,
DTYPE_t val=0):
# Assumptions: - node is a valid pointer
# - node is not currently part of a heap
node.index = index
node.val = val
node.rank = 0
node.state = NOT_IN_HEAP
node.parent = NULL
node.left_sibling = NULL
node.right_sibling = NULL
node.children = NULL
cdef FibonacciNode* rightmost_sibling(FibonacciNode* node):
# Assumptions: - node is a valid pointer
cdef FibonacciNode* temp = node
while(temp.right_sibling):
temp = temp.right_sibling
return temp
cdef FibonacciNode* leftmost_sibling(FibonacciNode* node):
# Assumptions: - node is a valid pointer
cdef FibonacciNode* temp = node
while(temp.left_sibling):
temp = temp.left_sibling
return temp
cdef void add_child(FibonacciNode* node, FibonacciNode* new_child):
# Assumptions: - node is a valid pointer
# - new_child is a valid pointer
# - new_child is not the sibling or child of another node
new_child.parent = node
if node.children:
add_sibling(node.children, new_child)
else:
node.children = new_child
new_child.right_sibling = NULL
new_child.left_sibling = NULL
node.rank = 1
cdef void add_sibling(FibonacciNode* node, FibonacciNode* new_sibling):
# Assumptions: - node is a valid pointer
# - new_sibling is a valid pointer
# - new_sibling is not the child or sibling of another node
cdef FibonacciNode* temp = rightmost_sibling(node)
temp.right_sibling = new_sibling
new_sibling.left_sibling = temp
new_sibling.right_sibling = NULL
new_sibling.parent = node.parent
if new_sibling.parent:
new_sibling.parent.rank += 1
cdef void remove(FibonacciNode* node):
# Assumptions: - node is a valid pointer
if node.parent:
node.parent.rank -= 1
if node.left_sibling:
node.parent.children = node.left_sibling
elif node.right_sibling:
node.parent.children = node.right_sibling
else:
node.parent.children = NULL
if node.left_sibling:
node.left_sibling.right_sibling = node.right_sibling
if node.right_sibling:
node.right_sibling.left_sibling = node.left_sibling
node.left_sibling = NULL
node.right_sibling = NULL
node.parent = NULL
######################################################################
# FibonacciHeap structure
# This structure and operations on it use the FibonacciNode
# routines to implement a Fibonacci heap
ctypedef FibonacciNode* pFibonacciNode
cdef struct FibonacciHeap:
FibonacciNode* min_node
pFibonacciNode[100] roots_by_rank # maximum number of nodes is ~2^100.
cdef void insert_node(FibonacciHeap* heap,
FibonacciNode* node):
# Assumptions: - heap is a valid pointer
# - node is a valid pointer
# - node is not the child or sibling of another node
if heap.min_node:
add_sibling(heap.min_node, node)
if node.val < heap.min_node.val:
heap.min_node = node
else:
heap.min_node = node
cdef void decrease_val(FibonacciHeap* heap,
FibonacciNode* node,
DTYPE_t newval):
# Assumptions: - heap is a valid pointer
# - newval <= node.val
# - node is a valid pointer
# - node is not the child or sibling of another node
# - node is in the heap
node.val = newval
if node.parent and (node.parent.val >= newval):
remove(node)
insert_node(heap, node)
elif heap.min_node.val > node.val:
heap.min_node = node
cdef void link(FibonacciHeap* heap, FibonacciNode* node):
# Assumptions: - heap is a valid pointer
# - node is a valid pointer
# - node is already within heap
cdef FibonacciNode *linknode
cdef FibonacciNode *parent
cdef FibonacciNode *child
if heap.roots_by_rank[node.rank] == NULL:
heap.roots_by_rank[node.rank] = node
else:
linknode = heap.roots_by_rank[node.rank]
heap.roots_by_rank[node.rank] = NULL
if node.val < linknode.val or node == heap.min_node:
remove(linknode)
add_child(node, linknode)
link(heap, node)
else:
remove(node)
add_child(linknode, node)
link(heap, linknode)
cdef FibonacciNode* remove_min(FibonacciHeap* heap):
# Assumptions: - heap is a valid pointer
# - heap.min_node is a valid pointer
cdef FibonacciNode *temp
cdef FibonacciNode *temp_right
cdef FibonacciNode *out
cdef unsigned int i
# make all min_node children into root nodes
if heap.min_node.children:
temp = leftmost_sibling(heap.min_node.children)
temp_right = NULL
while temp:
temp_right = temp.right_sibling
remove(temp)
add_sibling(heap.min_node, temp)
temp = temp_right
heap.min_node.children = NULL
# choose a root node other than min_node
temp = leftmost_sibling(heap.min_node)
if temp == heap.min_node:
if heap.min_node.right_sibling:
temp = heap.min_node.right_sibling
else:
out = heap.min_node
heap.min_node = NULL
return out
# remove min_node, and point heap to the new min
out = heap.min_node
remove(heap.min_node)
heap.min_node = temp
# re-link the heap
for i in range(100):
heap.roots_by_rank[i] = NULL
while temp:
if temp.val < heap.min_node.val:
heap.min_node = temp
temp_right = temp.right_sibling
link(heap, temp)
temp = temp_right
return out
# Copyright 2011-2013 Kwant authors.
#
# Copyright 2011-2018 Kwant authors.
# This file is part of Kwant. It is subject to the license terms in the file
# LICENSE.rst found in the top-level directory of this distribution and at
# http://kwant-project.org/license. A list of Kwant authors can be found in
......@@ -10,7 +9,7 @@
# Merge the public interface of all submodules.
__all__ = []
for module in ['leads', 'dispersion', 'noise', 'symmetry']:
for module in ['leads', 'dispersion', 'noise', 'symmetry', 'gauge']:
exec('from . import {0}'.format(module))
exec('from .{0} import *'.format(module))
exec('__all__.extend({0}.__all__)'.format(module))
# Copyright 2011-2018 Kwant authors.
#
# This file is part of Kwant. It is subject to the license terms in the file
# LICENSE.rst found in the top-level directory of this distribution and at
# http://kwant-project.org/license. A list of Kwant authors can be found in
# the file AUTHORS.rst at the top-level directory of this distribution and at
# http://kwant-project.org/authors.
"""Functions for fixing the magnetic gauge automatically in a Kwant system."""
import functools as ft
import numpy as np
import scipy
from scipy.integrate import dblquad
from .. import system, builder
from ..graph.dijkstra import dijkstra_directed
__all__ = ['magnetic_gauge']
### Integation
# Integrate vector field over triangle, for internal use by 'surface_integral'
# Triangle is (origin, origin + v1, origin + v2), 'n' is np.cross(v1, v2)
def _quad_triangle(f, origin, v1, v2, n, tol):
if np.dot(n, n) < tol**2: # does triangle have significant area?
return 0
def g(x, y):
return np.dot(n, f(origin + x * v1 + y * v2))
result, *_ = dblquad(g, 0, 1, lambda x: 0, lambda x: 1-x)
return result.real
def _const_triangle(f, origin, v1, v2, n, tol):
return np.dot(f, n) / 2
def _average_triangle(f, origin, v1, v2, n, tol):
return np.dot(n, f(origin + 1/3 * (v1 + v2))) / 2
def surface_integral(f, loop, tol=1e-8, average=False):
"""Calculate the surface integral of 'f' over a surface enclosed by 'loop'.
This function only works for *divergence free* vector fields, where the
surface integral depends only on the boundary.
Parameters
----------
f : callable or real n-vector
The vector field for which to calculate the surface integral.
If callable, takes a real n-vector as argument and returns a
real n-vector.
loop : sequence of vectors
Ordered sequence of real n-vectors (positions) that define the
vertices of the polygon that encloses the surface to integrate over.
tol : float, default: 1e-8
Error tolerance on the result.
average : bool, default: False
If True, approximate the integral over each triangle using a
single function evaluation at the centre of the triangle.
"""
if callable(f):
integrator = _average_triangle if average else _quad_triangle
else:
integrator = _const_triangle
origin, *points = loop
integral = 0
# Split loop into triangles with 1 vertex on 'origin', evaluate
# the integral over each triangle and sum the result
for p1, p2 in zip(points, points[1:]):
v1 = p1 - origin
v2 = p2 - origin
n = np.cross(v1, v2)
integral += integrator(f, origin, v1, v2, n, tol)
return integral
### Loop finding graph algorithm
def find_loops(graph, subgraph):
"""
Parameters
----------
graph : COO matrix
The complete undirected graph, where the values of the matrix are
the weights of the corresponding graph links.
subgraph : COO matrix
An subgraph of 'graph', with missing edges denoted by infinities.
Must have the same sparsity structure as 'graph'.
Returns
-------
A sequence of paths which are partially contained in the subgraph.
The loop is formed by adding a link between the first and last node.
The ordering is such that the paths are made of links that belong to
the subgraph or to the previous closed loops.
"""
# For each link we do 1 update of 'subgraph' and a call to
# 'csgraph.shortest_path'. It is cheaper to update the CSR
# matrix rather than convert to LIL and back every iteration.
subgraph = subgraph.tocsr()
graph = graph.tocsr()
assert same_sparsity_structure(subgraph, graph)
# Links in graph, but not in subgraph.
links_to_find = scipy.sparse.triu(graph - subgraph).tocoo()
links_to_find = np.vstack((links_to_find.row, links_to_find.col)).transpose()
links_to_find, min_length = order_links(subgraph, links_to_find)
# Find shortest path between each link in turn, updating the subgraph with
# the links as we go.
loops = []
while len(links_to_find) > 0:
frm, to = links_to_find[0]
(path,), (path_length,) = dijkstra_directed(subgraph,
sources=np.array([frm]),
targets=np.array([to]))
# Reorder links that are still to find based on the loop length in
# the updated graph. We only reorder when the path length for *this*
# link is a "little bit" longer that the perviously determined minimum.
# The "little bit" is needed so we don't needlessly re-order the links
# on amorphous lattices.
if path_length > min_length * 1.1:
links_to_find, min_length = order_links(subgraph, links_to_find)
else:
# Assumes that 'graph' and 'subgraph' have the same sparsity structure.
assign_csr(subgraph, graph, (frm, to))
assign_csr(subgraph, graph, (to, frm))
loops.append(path)
links_to_find = links_to_find[1:]
return loops
def order_links(subgraph, links_to_find):
if len(links_to_find) == 0:
return [], None
# Order 'links_to_find' by length of shortest path between the nodes of the link
path_lengths = dijkstra_directed(subgraph,
sources=links_to_find[:, 0],
targets=links_to_find[:, 1],
return_paths=False)
idxs = np.argsort(path_lengths)
return links_to_find[idxs], path_lengths[idxs[0]]
### Generic sparse matrix utilities
def assign_csr(a, b, element):
"""Assign a single element from a CSR matrix to another.
Parameters
----------
a : CSR matrix
b : CSR matrix or scalar
If a CSR matrix, must have the same sparsity structure
as 'a'. If a scalar, must be the same dtype as 'a'.
element: (int, int)
Row and column indices of the element to set.
"""
assert isinstance(a, scipy.sparse.csr_matrix)
row, col = element
for j in range(a.indptr[row], a.indptr[row + 1]):
if a.indices[j] == col:
break
else:
raise ValueError('{} not in sparse matrix'.format(element))
if isinstance(b, scipy.sparse.csr_matrix):
a.data[j] = b.data[j]
else:
a.data[j] = b
def same_sparsity_structure(a, b):
a = a.tocsr().sorted_indices()
b = b.tocsr().sorted_indices()
return (np.array_equal(a.indices, b.indices)
and np.array_equal(a.indptr, b.indptr))
def add_coo_matrices(*mats, shape):
"""Add a sequence of COO matrices by appending their constituent arrays."""
values = np.hstack([mat.data for mat in mats])
rows = np.hstack([mat.row for mat in mats])
cols = np.hstack([mat.col for mat in mats])
return scipy.sparse.coo_matrix((values, (rows, cols)), shape=shape)
def shift_diagonally(mat, shift, shape):
"""Shift the row/column indices of a COO matrix."""
return scipy.sparse.coo_matrix(
(mat.data, (mat.row + shift, mat.col + shift)),
shape=shape)