#!/usr/bin/python
#
# This file provides a variety of parallel input/output functions.
# ----------------------------------------------------------------------------
# pyCTQW - Distributed memory CTQW Fortran library and Python module
# Copyright (C) 2013-2014, Joshua Izaac
#
# pyCTQW is free software: you can redistribute it and/or modify it under the
# terms of version 3 of the GNU General Public License as published by
# the Free Software Foundation.
#
# pyCTQW is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License
# along with pyCTQW. If not, see <http://www.gnu.org/licenses/>.
# ----------------------------------------------------------------------------
"""
Input and output functions.
"""
import os as _os
import errno as _errno
import sys as _sys
import numpy as _np
from petsc4py import PETSc as _PETSc
import fileinput as _fl
def createPath(filename):
if _os.path.isabs(filename):
outDir = _os.path.dirname(filename)
else:
outDir = './'+_os.path.dirname(filename)
# create output directory if it doesn't exist
try:
_os.mkdir(outDir)
except OSError as exception:
if exception.errno != _errno.EEXIST:
raise
[docs]def vecToArray(obj):
""" Converts a PETSc vector to a numpy array, available on *all* MPI nodes.
Args:
obj (petsc4py.PETSc.Vec): input vector.
Returns:
numpy.array :
"""
# scatter vector 'obj' to all processes
comm = obj.getComm()
scatter, obj0 = _PETSc.Scatter.toAll(obj)
scatter.scatter(obj, obj0, False, _PETSc.Scatter.Mode.FORWARD)
return _np.asarray(obj0)
# deallocate
comm.barrier()
scatter.destroy()
obj0.destroy()
[docs]def vecToArray0(obj):
""" Converts a PETSc vector to a numpy array available on MPI node 0.
Args:
obj (petsc4py.PETSc.Vec): input vector.
Returns:
numpy.array :
"""
# scatter vector 'obj' to process 0
comm = obj.getComm()
rank = comm.getRank()
scatter, obj0 = _PETSc.Scatter.toZero(obj)
scatter.scatter(obj, obj0, False, _PETSc.Scatter.Mode.FORWARD)
if rank == 0: return _np.asarray(obj0)
# deallocate
comm.barrier()
scatter.destroy()
obj0.destroy()
[docs]def arrayToVec(vecArray):
""" Converts a (global) array to a PETSc vector over :attr:`petsc4py.PETSc.COMM_WORLD`.
Args:
vecArray (array or numpy.array): input vector.
Returns:
petsc4py.PETSc.Vec() :
"""
vec = _PETSc.Vec().create(comm=_PETSc.COMM_WORLD)
vec.setSizes(len(vecArray))
vec.setUp()
(Istart,Iend) = vec.getOwnershipRange()
return vec.createWithArray(vecArray[Istart:Iend],
comm=_PETSc.COMM_WORLD)
vec.destroy()
[docs]def arrayToMat(matArray):
""" Converts a (global) 2D array to a PETSc matrix over :attr:`petsc4py.PETSc.COMM_WORLD`.
Args:
matArray (array or numpy.array): input square array.
:rtype: petsc4py.PETSc.Mat()
.. important::
Requires `SciPy <http://www.scipy.org>`_.
"""
try:
import scipy.sparse as sparse
except:
print '\nERROR: loading matrices from txt files requires Scipy!'
return
matSparse = sparse.csr_matrix(matArray)
mat = _PETSc.Mat().createAIJ(size=matSparse.shape,comm=_PETSc.COMM_WORLD)
(Istart,Iend) = mat.getOwnershipRange()
ai = matSparse.indptr[Istart:Iend+1] - matSparse.indptr[Istart]
aj = matSparse.indices[matSparse.indptr[Istart]:matSparse.indptr[Iend]]
av = matSparse.data[matSparse.indptr[Istart]:matSparse.indptr[Iend]]
mat.setValuesCSR(ai,aj,av)
mat.assemble()
return mat
mat.destroy()
[docs]def matToSparse(mat):
""" Converts a PETSc matrix to a (global) sparse matrix.
Args:
mat (petsc4py.PETSc.Mat): input PETSc matrix.
:rtype: scipy.sparse.csr_matrix
.. important::
Requires `SciPy <http://www.scipy.org>`_.
"""
import scipy.sparse as sparse
data = mat.getValuesCSR()
(Istart,Iend) = mat.getOwnershipRange()
columns = mat.getSize()[0]
sparseSubMat = sparse.csr_matrix(data[::-1],shape=(Iend-Istart,columns))
comm = _PETSc.COMM_WORLD
sparseSubMat = comm.tompi4py().allgather(sparseSubMat)
return sparse.vstack(sparseSubMat)
[docs]def adjToH(adj,d=[0],amp=[0.]):
""" Creates a 1 particle PETSc-type Hamiltonian matrix from a PETSc adjacency matrix.
Args:
adj (petsc4py.PETSc.Mat): input PETSc-type adjacency matrix.
d (array of ints): an array containing *integers* indicating the nodes
where diagonal defects are to be placed (e.g. ``d=[0,1,4]``).
amp (array of floats): an array containing *floats* indicating the diagonal defect
amplitudes corresponding to each element in ``d`` (e.g. ``amp=[0.5,-1,4.2]``).
Returns:
: 1 particle Hamiltonian matrix
:rtype: petsc4py.PETSc.Mat()
Warning:
* The size of ``a`` and ``d`` must be identical
>>> amp = [0.5,-1.,4.2]
>>> len(d) == len(amp)
True
* Elements of ``d`` can range from :math:`[0,N-1]` where the adjacency matrix is :math:`N\\times N`.
"""
(Istart,Iend) = adj.getOwnershipRange()
diagSum = []
for i in range(Istart,Iend):
diagSum.append(_np.sum(adj.getRow(i)[-1]))
for j,val in enumerate(d):
if i==val: diagSum[i-Istart] += amp[j]
mat = _PETSc.Mat().create(comm=_PETSc.COMM_WORLD)
mat.setSizes(adj.getSize())
mat.setUp()
for i in range(Istart,Iend):
mat.setValue(i,i,diagSum[i-Istart])
mat.assemble()
mat.axpy(-1,adj)
return mat
mat.destroy()
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#---------------------- Vec I/O functions ---------------------------
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
[docs]def exportVec(vec,filename,filetype):
""" Export a PETSc vector to a file.
Args:
vec (petsc4py.PETSc.Vec): input vector.
filename (str): path to desired output file.
filetype (str): the filetype of the exported vector.
* ``'txt'`` - a column vector in text format.
* ``'bin'`` - a PETSc binary vector.
"""
createPath(filename)
if filetype == 'txt':
# scatter prob to process 0
comm = vec.getComm()
rank = comm.getRank()
scatter, vec0 = _PETSc.Scatter.toZero(vec)
scatter.scatter(vec, vec0, False, _PETSc.Scatter.Mode.FORWARD)
# use process 0 to write to text file
if rank == 0:
array0 = _np.asarray(vec0)
with open(filename,'w') as f:
for i in range(len(array0)):
f.write('{0: .12e}\n'.format(array0[i]))
# deallocate
comm.barrier()
scatter.destroy()
vec0.destroy()
elif filetype == 'bin':
binSave = _PETSc.Viewer().createBinary(filename, 'w')
binSave(vec)
binSave.destroy()
vec.comm.barrier()
[docs]def loadVec(filename,filetype):
""" Import a PETSc vector from a file.
Args:
filename (str): path to input file.
filetype (str): the filetype.
* ``'txt'`` - a column vector in text format.
* ``'bin'`` - a PETSc binary vector.
"""
if filetype == 'txt':
try:
vecArray = _np.loadtxt(filename,dtype=_PETSc.ScalarType)
return arrayToVec(vecArray)
except:
print "\nERROR: input state space file " + filename\
+ " does not exist or is in an incorrect format"
_sys.exit()
elif filetype == 'bin':
binLoad = _PETSc.Viewer().createBinary(filename, 'r')
try:
return _PETSc.Vec().load(binLoad)
except:
print "\nERROR: input state space file " + filename\
+ " does not exist or is in an incorrect format"
_sys.exit()
binLoad.destroy()
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#---------------------- Mat I/O functions ---------------------------
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
[docs]def exportMat(mat,filename,filetype,mattype=None):
""" Export a PETSc matrix to a file.
Args:
mat (petsc4py.PETSc.Mat): input matrix.
filename (str): path to desired output file.
filetype (str): the filetype of the exported vector.
* ``'txt'`` - a 2D matrix array in text format.
* ``'bin'`` - a PETSc binary matrix.
mattype (str): (``None``,``'adj'``) - if set to ``adj``, only
integers ``0`` and ``1`` are written. Note
that this only applied in ``txt`` mode.
"""
rank = _PETSc.Comm.Get_rank(_PETSc.COMM_WORLD)
createPath(filename)
if filetype == 'txt':
txtSave = _PETSc.Viewer().createASCII(filename, 'w',
format=_PETSc.Viewer.Format.ASCII_DENSE, comm=_PETSc.COMM_WORLD)
txtSave(mat)
txtSave.destroy()
if rank == 0:
for line in _fl.FileInput(filename,inplace=1):
if line[2] != 't':
if mattype == 'adj':
line = line.replace(" i","j")
line = line.replace(" -","-")
line = line.replace("+-","-")
line = line.replace("0000e+01+0.00000e+00j","")
line = line.replace(".00000e+00+0.00000e+00j","")
line = line.replace(".","")
line = line.replace(" -","\t-")
line = line.replace(" ","\t")
line = line.replace(" ","")
line = line.replace("\t"," ")
print line,
else:
line = line.replace(" i","j")
line = line.replace(" -","-")
line = line.replace("+-","-")
print line,
elif filetype == 'bin':
binSave = _PETSc.Viewer().createBinary(filename, 'w', comm=_PETSc.COMM_WORLD)
binSave(mat)
binSave.destroy()
mat.comm.barrier()
[docs]def loadMat(filename,filetype,delimiter=None):
""" Import a PETSc matrix from a file.
Args:
filename (str): path to input file.
filetype (str): the filetype.
* ``'txt'`` - a 2D matrix array in text format.
* ``'bin'`` - a PETSc matrix vector.
delimiter (str): this is passed to `numpy.genfromtxt\
<http://docs.scipy.org/doc/numpy/reference/generated/numpy.genfromtxt.html>`_
in the case of strange delimiters in an imported ``txt`` file.
"""
if filetype == 'txt':
try:
try:
if delimiter is None:
matArray = _np.genfromtxt(filename,dtype=_PETSc.ScalarType)
else:
matArray = _np.genfromtxt(filename,dtype=_PETSc.ScalarType,delimiter=delimiter)
except:
filefix = []
for line in _fl.FileInput(filename,inplace=0):
if line[2] != 't':
line = line.replace(" i","j")
line = line.replace(" -","-")
line = line.replace("+-","-")
filefix.append(line)
matArray = _np.genfromtxt(filefix,dtype=_PETSc.ScalarType)
return arrayToMat(matArray)
except:
print "\nERROR: input state space file " + filename\
+ " does not exist or is in an incorrect format"
_sys.exit()
elif filetype == 'bin':
binLoad = _PETSc.Viewer().createBinary(filename, 'r')
try:
return _PETSc.Mat().load(binLoad)
except:
print "\nERROR: input state space file " + filename\
+ " does not exist or is in an incorrect format"
_sys.exit()
binLoad.destroy()
[docs]def exportVecToMat(vec,filename,filetype):
""" Export a :math:`N^2` element PETSc vector as a :math:`N\\times N` matrix.
This is useful when wanting to view the full statespace of a 2 particle
quantum walk.
Args:
vec (petsc4py.PETSc.Vec): input :math:`N^2` element vector.
filename (str): path to desired output file.
filetype (str): the filetype of the exported vector.
* ``'txt'`` - an :math:`N\\times N` 2D matrix array in text format.
* ``'bin'`` - an :math:`N\\times N` PETSc binary matrix.
"""
rank = _PETSc.Comm.Get_rank(_PETSc.COMM_WORLD)
createPath(filename)
vecArray = vecToArray(vec)
matArray = vecArray.reshape([_np.sqrt(vecArray.size),_np.sqrt(vecArray.size)])
if filetype == 'txt':
#if rank == 0: _np.savetxt(filename,matArray)
txtSave = _PETSc.Viewer().createASCII(filename, 'w',
format=_PETSc.Viewer.Format.ASCII_DENSE, comm=_PETSc.COMM_WORLD)
txtSave(arrayToMat(matArray))
txtSave.destroy()
if rank == 0:
for line in _fl.FileInput(filename,inplace=1):
if line[2] != 't':
line = line.replace(" i","j")
line = line.replace(" -","-")
line = line.replace("+-","-")
print line,
elif filetype == 'bin':
binSave = _PETSc.Viewer().createBinary(filename, 'w', comm=_PETSc.COMM_WORLD)
binSave(arrayToMat(matArray))
binSave.destroy()
vec.comm.barrier()
[docs]def loadMatToVec(filename,filetype):
""" Load a :math:`N\\times N` matrix as a :math:`N^2` element PETSc vector.
This is useful when wanting to import the full statespace of a 2 particle
quantum walk to use for propagation.
Args:
filename (str): path to the input file.
filetype (str): the filetype
* ``'txt'`` - an :math:`N\\times N` 2D matrix array in text format.
* ``'bin'`` - **Not yet implemented! Please use a txt \
format for this type of import**.
"""
if filetype == 'txt':
try:
try:
matArray = _np.loadtxt(filename,dtype=_PETSc.ScalarType)
except:
filefix = []
for line in _fl.FileInput(filename,inplace=0):
if line[2] != 't':
line = line.replace(" i","j")
line = line.replace(" -","-")
line = line.replace("+-","-")
filefix.append(line)
matArray = _np.loadtxt(filefix,dtype=_PETSc.ScalarType)
vecArray = matArray.reshape(matArray.shape[0]**2)
return arrayToVec(vecArray)
except:
print "\nERROR: input state space file " + filename\
+ " does not exist or is in an incorrect format"
_sys.exit()
elif filetype == 'bin':
print '\nERROR: only works for txt storage!'
_sys.exit()