Files
lammps/lib/message/cslib/src/cslib.py
2018-07-24 17:22:15 -06:00

363 lines
12 KiB
Python

# ------------------------------------------------------------------------
# CSlib - Client/server library for code coupling
# http://cslib.sandia.gov, Sandia National Laboratories
# Steve Plimpton, sjplimp@sandia.gov
#
# Copyright 2018 National Technology & Engineering Solutions of
# Sandia, LLC (NTESS). Under the terms of Contract DE-NA0003525 with
# NTESS, the U.S. Government retains certain rights in this software.
# This software is distributed under the modified Berkeley Software
# Distribution (BSD) License.
#
# See the README file in the top-level CSlib directory.
# -------------------------------------------------------------------------
# Python wrapper on CSlib library via ctypes
# ctypes and Numpy data types:
# 32-bit int = c_int = np.intc = np.int32
# 64-bit int = c_longlong = np.int64
# 32-bit floating point = c_float = np.float32
# 64-bit floating point = c_double = np.float = np.float64
import sys,traceback
from ctypes import *
# Numpy and mpi4py packages may not exist
try:
import numpy as np
numpyflag = 1
except:
numpyflag = 0
try:
from mpi4py import MPI
mpi4pyflag = 1
except:
mpi4pyflag = 0
# wrapper class
class CSlib:
# instantiate CSlib thru its C-interface
def __init__(self,csflag,mode,ptr,comm):
# load libcslib.so
try:
if comm: self.lib = CDLL("libcsmpi.so",RTLD_GLOBAL)
else: self.lib = CDLL("libcsnompi.so",RTLD_GLOBAL)
except:
etype,value,tb = sys.exc_info()
traceback.print_exception(etype,value,tb)
raise OSError,"Could not load CSlib dynamic library"
# define ctypes API for each library method
self.lib.cslib_open.argtypes = [c_int,c_char_p,c_void_p,c_void_p,
POINTER(c_void_p)]
self.lib.cslib_open.restype = None
self.lib.cslib_close.argtypes = [c_void_p]
self.lib.cslib_close.restype = None
self.lib.cslib_send.argtypes = [c_void_p,c_int,c_int]
self.lib.cslib_send.restype = None
self.lib.cslib_pack_int.argtypes = [c_void_p,c_int,c_int]
self.lib.cslib_pack_int.restype = None
self.lib.cslib_pack_int64.argtypes = [c_void_p,c_int,c_longlong]
self.lib.cslib_pack_int64.restype = None
self.lib.cslib_pack_float.argtypes = [c_void_p,c_int,c_float]
self.lib.cslib_pack_float.restype = None
self.lib.cslib_pack_double.argtypes = [c_void_p,c_int,c_double]
self.lib.cslib_pack_double.restype = None
self.lib.cslib_pack_string.argtypes = [c_void_p,c_int,c_char_p]
self.lib.cslib_pack_string.restype = None
self.lib.cslib_pack.argtypes = [c_void_p,c_int,c_int,c_int,c_void_p]
self.lib.cslib_pack.restype = None
self.lib.cslib_pack_parallel.argtypes = [c_void_p,c_int,c_int,c_int,
POINTER(c_int),c_int,c_void_p]
self.lib.cslib_pack_parallel.restype = None
self.lib.cslib_recv.argtypes = [c_void_p,POINTER(c_int),
POINTER(POINTER(c_int)),
POINTER(POINTER(c_int)),
POINTER(POINTER(c_int))]
self.lib.cslib_recv.restype = c_int
self.lib.cslib_unpack_int.argtypes = [c_void_p,c_int]
self.lib.cslib_unpack_int.restype = c_int
self.lib.cslib_unpack_int64.argtypes = [c_void_p,c_int]
self.lib.cslib_unpack_int64.restype = c_longlong
self.lib.cslib_unpack_float.argtypes = [c_void_p,c_int]
self.lib.cslib_unpack_float.restype = c_float
self.lib.cslib_unpack_double.argtypes = [c_void_p,c_int]
self.lib.cslib_unpack_double.restype = c_double
self.lib.cslib_unpack_string.argtypes = [c_void_p,c_int]
self.lib.cslib_unpack_string.restype = c_char_p
# override return in unpack()
self.lib.cslib_unpack.argtypes = [c_void_p,c_int]
self.lib.cslib_unpack.restype = c_void_p
self.lib.cslib_unpack_data.argtypes = [c_void_p,c_int,c_void_p]
self.lib.cslib_unpack_data.restype = None
# override last arg in unpack_parallel()
self.lib.cslib_unpack_parallel.argtypes = [c_void_p,c_int,c_int,
POINTER(c_int),c_int,c_void_p]
self.lib.cslib_unpack_parallel.restype = None
self.lib.cslib_extract.argtypes = [c_void_p,c_int]
self.lib.cslib_extract.restype = c_int
# create an instance of CSlib with or w/out MPI communicator
self.cs = c_void_p()
if not comm:
self.lib.cslib_open(csflag,mode,ptr,None,byref(self.cs))
elif not mpi4pyflag:
print "Cannot pass MPI communicator to CSlib w/out mpi4py package"
sys.exit()
else:
address = MPI._addressof(comm)
comm_ptr = c_void_p(address)
if mode == "mpi/one":
address = MPI._addressof(ptr)
ptrcopy = c_void_p(address)
else: ptrcopy = ptr
self.lib.cslib_open(csflag,mode,ptrcopy,comm_ptr,byref(self.cs))
# destroy instance of CSlib
def __del__(self):
if self.cs: self.lib.cslib_close(self.cs)
def close(self):
self.lib.cslib_close(self.cs)
self.lib = None
# send a message
def send(self,msgID,nfield):
self.nfield = nfield
self.lib.cslib_send(self.cs,msgID,nfield)
# pack one field of message
def pack_int(self,id,value):
self.lib.cslib_pack_int(self.cs,id,value)
def pack_int64(self,id,value):
self.lib.cslib_pack_int64(self.cs,id,value)
def pack_float(self,id,value):
self.lib.cslib_pack_float(self.cs,id,value)
def pack_double(self,id,value):
self.lib.cslib_pack_double(self.cs,id,value)
def pack_string(self,id,value):
self.lib.cslib_pack_string(self.cs,id,value)
def pack(self,id,ftype,flen,data):
cdata = self.data_convert(ftype,flen,data)
self.lib.cslib_pack(self.cs,id,ftype,flen,cdata)
def pack_parallel(self,id,ftype,nlocal,ids,nper,data):
cids = self.data_convert(1,nlocal,ids)
cdata = self.data_convert(ftype,nper*nlocal,data)
self.lib.cslib_pack_parallel(self.cs,id,ftype,nlocal,cids,nper,cdata)
# convert input data to a ctypes vector to pass to CSlib
def data_convert(self,ftype,flen,data):
# tflag = type of data
# tflag = 1 if data is list or tuple
# tflag = 2 if data is Numpy array
# tflag = 3 if data is ctypes vector
# same usage of tflag as in unpack function
txttype = str(type(data))
if "numpy" in txttype: tflag = 2
elif "c_" in txttype: tflag = 3
else: tflag = 1
# create ctypes vector out of data to pass to lib
# cdata = ctypes vector to return
# NOTE: error check on ftype and tflag everywhere, also flen
if ftype == 1:
if tflag == 1: cdata = (flen * c_int)(*data)
elif tflag == 2: cdata = data.ctypes.data_as(POINTER(c_int))
elif tflag == 3: cdata = data
elif ftype == 2:
if tflag == 1: cdata = (flen * c_longlong)(*data)
elif tflag == 2: cdata = data.ctypes.data_as(POINTER(c_longlong))
elif tflag == 3: cdata = data
elif ftype == 3:
if tflag == 1: cdata = (flen * c_float)(*data)
elif tflag == 2: cdata = data.ctypes.data_as(POINTER(c_float))
elif tflag == 3: cdata = data
elif ftype == 4:
if tflag == 1: cdata = (flen * c_double)(*data)
elif tflag == 2: cdata = data.ctypes.data_as(POINTER(c_double))
elif tflag == 3: cdata = data
return cdata
# receive a message
def recv(self):
self.lib.cslib_recv.restype = c_int
nfield = c_int()
fieldID = POINTER(c_int)()
fieldtype = POINTER(c_int)()
fieldlen = POINTER(c_int)()
msgID = self.lib.cslib_recv(self.cs,byref(nfield),
byref(fieldID),byref(fieldtype),byref(fieldlen))
# copy returned C args to native Python int and lists
# store them in class so unpack() methods can access the info
self.nfield = nfield = nfield.value
self.fieldID = fieldID[:nfield]
self.fieldtype = fieldtype[:nfield]
self.fieldlen = fieldlen[:nfield]
return msgID,self.nfield,self.fieldID,self.fieldtype,self.fieldlen
# unpack one field of message
# tflag = type of data to return
# 3 = ctypes vector is default, since no conversion required
def unpack_int(self,id):
return self.lib.cslib_unpack_int(self.cs,id)
def unpack_int64(self,id):
return self.lib.cslib_unpack_int64(self.cs,id)
def unpack_float(self,id):
return self.lib.cslib_unpack_float(self.cs,id)
def unpack_double(self,id):
return self.lib.cslib_unpack_double(self.cs,id)
def unpack_string(self,id):
return self.lib.cslib_unpack_string(self.cs,id)
def unpack(self,id,tflag=3):
index = self.fieldID.index(id)
# reset data type of return so can morph by tflag
# cannot do this for the generic c_void_p returned by CSlib
if self.fieldtype[index] == 1:
self.lib.cslib_unpack.restype = POINTER(c_int)
elif self.fieldtype[index] == 2:
self.lib.cslib_unpack.restype = POINTER(c_longlong)
elif self.fieldtype[index] == 3:
self.lib.cslib_unpack.restype = POINTER(c_float)
elif self.fieldtype[index] == 4:
self.lib.cslib_unpack.restype = POINTER(c_double)
#elif self.fieldtype[index] == 5:
# self.lib.cslib_unpack.restype = POINTER(c_char)
cdata = self.lib.cslib_unpack(self.cs,id)
# tflag = user-requested type of data to return
# tflag = 1 to return data as list
# tflag = 2 to return data as Numpy array
# tflag = 3 to return data as ctypes vector
# same usage of tflag as in pack functions
# tflag = 2,3 should NOT perform a data copy
if tflag == 1:
data = cdata[:self.fieldlen[index]]
elif tflag == 2:
if numpyflag == 0:
print "Cannot return Numpy array w/out numpy package"
sys.exit()
data = np.ctypeslib.as_array(cdata,shape=(self.fieldlen[index],))
elif tflag == 3:
data = cdata
return data
# handle data array like pack() or unpack_parallel() ??
def unpack_data(self,id,tflag=3):
index = self.fieldID.index(id)
# unpack one field of message in parallel
# tflag = type of data to return
# 3 = ctypes vector is default, since no conversion required
# NOTE: allow direct use of user array (e.g. Numpy), if user provides data arg?
# as opposed to creating this cdata
# does that make any performance difference ?
# e.g. should we allow CSlib to populate an existing Numpy array's memory
def unpack_parallel(self,id,nlocal,ids,nper,tflag=3):
cids = self.data_convert(1,nlocal,ids)
# allocate memory for the returned data
# pass cdata ptr to the memory to CSlib unpack_parallel()
# this resets data type of last unpack_parallel() arg
index = self.fieldID.index(id)
if self.fieldtype[index] == 1: cdata = (nper*nlocal * c_int)()
elif self.fieldtype[index] == 2: cdata = (nlocal*nper * c_longlong)()
elif self.fieldtype[index] == 3: cdata = (nlocal*nper * c_float)()
elif self.fieldtype[index] == 4: cdata = (nlocal*nper * c_double)()
#elif self.fieldtype[index] == 5: cdata = (nlocal*nper * c_char)()
self.lib.cslib_unpack_parallel(self.cs,id,nlocal,cids,nper,cdata)
# tflag = user-requested type of data to return
# tflag = 1 to return data as list
# tflag = 2 to return data as Numpy array
# tflag = 3 to return data as ctypes vector
# same usage of tflag as in pack functions
if tflag == 1:
data = cdata[:nper*nlocal]
elif tflag == 2:
if numpyflag == 0:
print "Cannot return Numpy array w/out numpy package"
sys.exit()
# NOTE: next line gives ctypes warning for fieldtype = 2 = 64-bit int
# not sure why, reported as bug between ctypes and Numpy here:
# https://stackoverflow.com/questions/4964101/pep-3118-
# warning-when-using-ctypes-array-as-numpy-array
# but why not same warning when just using unpack() ??
# in Python these lines give same warning:
# >>> import ctypes,numpy
# >>> a = (10 * ctypes.c_longlong)()
# >>> b = numpy.ctypeslib.as_array(a)
data = np.ctypeslib.as_array(cdata,shape=(nlocal*nper,))
elif tflag == 3:
data = cdata
return data
# extract a library value
def extract(self,flag):
return self.lib.cslib_extract(self.cs,flag)