forked from lijiext/lammps
194 lines
6.9 KiB
Python
194 lines
6.9 KiB
Python
# ----------------------------------------------------------------------
|
|
# LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
|
|
# http://lammps.sandia.gov, Sandia National Laboratories
|
|
# Steve Plimpton, sjplimp@sandia.gov
|
|
#
|
|
# Copyright (2003) Sandia Corporation. Under the terms of Contract
|
|
# DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
|
|
# certain rights in this software. This software is distributed under
|
|
# the GNU General Public License.
|
|
#
|
|
# See the README file in the top-level LAMMPS directory.
|
|
# -------------------------------------------------------------------------
|
|
|
|
# Python wrapper on LAMMPS library via ctypes
|
|
|
|
import sys,traceback,types
|
|
from ctypes import *
|
|
from os.path import dirname,abspath,join
|
|
from inspect import getsourcefile
|
|
|
|
class lammps:
|
|
def __init__(self,name="",cmdargs=None,ptr=None):
|
|
|
|
# determine module location
|
|
modpath = dirname(abspath(getsourcefile(lambda:0)))
|
|
|
|
# load liblammps.so by default
|
|
# if name = "g++", load liblammps_g++.so
|
|
|
|
try:
|
|
if not name: self.lib = CDLL(join(modpath,"liblammps.so"),RTLD_GLOBAL)
|
|
else: self.lib = CDLL(join(modpath,"liblammps_%s.so" % name),RTLD_GLOBAL)
|
|
except:
|
|
type,value,tb = sys.exc_info()
|
|
traceback.print_exception(type,value,tb)
|
|
raise OSError,"Could not load LAMMPS dynamic library from %s" % modpath
|
|
|
|
# if no ptr provided, create an instance of LAMMPS
|
|
# don't know how to pass an MPI communicator from PyPar
|
|
# no_mpi call lets LAMMPS use MPI_COMM_WORLD
|
|
# cargs = array of C strings from args
|
|
# if ptr, then are embedding Python in LAMMPS input script
|
|
# ptr is the desired instance of LAMMPS
|
|
# just convert it to ctypes ptr and store in self.lmp
|
|
|
|
if not ptr:
|
|
self.opened = 1
|
|
if cmdargs:
|
|
cmdargs.insert(0,"lammps.py")
|
|
narg = len(cmdargs)
|
|
cargs = (c_char_p*narg)(*cmdargs)
|
|
self.lmp = c_void_p()
|
|
self.lib.lammps_open_no_mpi(narg,cargs,byref(self.lmp))
|
|
else:
|
|
self.lmp = c_void_p()
|
|
self.lib.lammps_open_no_mpi(0,None,byref(self.lmp))
|
|
# could use just this if LAMMPS lib interface supported it
|
|
# self.lmp = self.lib.lammps_open_no_mpi(0,None)
|
|
else:
|
|
self.opened = 0
|
|
# magic to convert ptr to ctypes ptr
|
|
pythonapi.PyCObject_AsVoidPtr.restype = c_void_p
|
|
pythonapi.PyCObject_AsVoidPtr.argtypes = [py_object]
|
|
self.lmp = c_void_p(pythonapi.PyCObject_AsVoidPtr(ptr))
|
|
|
|
def __del__(self):
|
|
if self.lmp and self.opened: self.lib.lammps_close(self.lmp)
|
|
|
|
def close(self):
|
|
if self.opened: self.lib.lammps_close(self.lmp)
|
|
self.lmp = None
|
|
|
|
def version(self):
|
|
return self.lib.lammps_version(self.lmp)
|
|
|
|
def file(self,file):
|
|
self.lib.lammps_file(self.lmp,file)
|
|
|
|
def command(self,cmd):
|
|
self.lib.lammps_command(self.lmp,cmd)
|
|
|
|
def extract_global(self,name,type):
|
|
if type == 0:
|
|
self.lib.lammps_extract_global.restype = POINTER(c_int)
|
|
elif type == 1:
|
|
self.lib.lammps_extract_global.restype = POINTER(c_double)
|
|
else: return None
|
|
ptr = self.lib.lammps_extract_global(self.lmp,name)
|
|
return ptr[0]
|
|
|
|
def extract_atom(self,name,type):
|
|
if type == 0:
|
|
self.lib.lammps_extract_atom.restype = POINTER(c_int)
|
|
elif type == 1:
|
|
self.lib.lammps_extract_atom.restype = POINTER(POINTER(c_int))
|
|
elif type == 2:
|
|
self.lib.lammps_extract_atom.restype = POINTER(c_double)
|
|
elif type == 3:
|
|
self.lib.lammps_extract_atom.restype = POINTER(POINTER(c_double))
|
|
else: return None
|
|
ptr = self.lib.lammps_extract_atom(self.lmp,name)
|
|
return ptr
|
|
|
|
def extract_compute(self,id,style,type):
|
|
if type == 0:
|
|
if style > 0: return None
|
|
self.lib.lammps_extract_compute.restype = POINTER(c_double)
|
|
ptr = self.lib.lammps_extract_compute(self.lmp,id,style,type)
|
|
return ptr[0]
|
|
if type == 1:
|
|
self.lib.lammps_extract_compute.restype = POINTER(c_double)
|
|
ptr = self.lib.lammps_extract_compute(self.lmp,id,style,type)
|
|
return ptr
|
|
if type == 2:
|
|
self.lib.lammps_extract_compute.restype = POINTER(POINTER(c_double))
|
|
ptr = self.lib.lammps_extract_compute(self.lmp,id,style,type)
|
|
return ptr
|
|
return None
|
|
|
|
# in case of global datum, free memory for 1 double via lammps_free()
|
|
# double was allocated by library interface function
|
|
|
|
def extract_fix(self,id,style,type,i=0,j=0):
|
|
if type == 0:
|
|
if style > 0: return None
|
|
self.lib.lammps_extract_fix.restype = POINTER(c_double)
|
|
ptr = self.lib.lammps_extract_fix(self.lmp,id,style,type,i,j)
|
|
result = ptr[0]
|
|
self.lib.lammps_free(ptr)
|
|
return result
|
|
if type == 1:
|
|
self.lib.lammps_extract_fix.restype = POINTER(c_double)
|
|
ptr = self.lib.lammps_extract_fix(self.lmp,id,style,type,i,j)
|
|
return ptr
|
|
if type == 2:
|
|
self.lib.lammps_extract_fix.restype = POINTER(POINTER(c_double))
|
|
ptr = self.lib.lammps_extract_fix(self.lmp,id,style,type,i,j)
|
|
return ptr
|
|
return None
|
|
|
|
# free memory for 1 double or 1 vector of doubles via lammps_free()
|
|
# for vector, must copy nlocal returned values to local c_double vector
|
|
# memory was allocated by library interface function
|
|
|
|
def extract_variable(self,name,group,type):
|
|
if type == 0:
|
|
self.lib.lammps_extract_variable.restype = POINTER(c_double)
|
|
ptr = self.lib.lammps_extract_variable(self.lmp,name,group)
|
|
result = ptr[0]
|
|
self.lib.lammps_free(ptr)
|
|
return result
|
|
if type == 1:
|
|
self.lib.lammps_extract_global.restype = POINTER(c_int)
|
|
nlocalptr = self.lib.lammps_extract_global(self.lmp,"nlocal")
|
|
nlocal = nlocalptr[0]
|
|
result = (c_double*nlocal)()
|
|
self.lib.lammps_extract_variable.restype = POINTER(c_double)
|
|
ptr = self.lib.lammps_extract_variable(self.lmp,name,group)
|
|
for i in xrange(nlocal): result[i] = ptr[i]
|
|
self.lib.lammps_free(ptr)
|
|
return result
|
|
return None
|
|
|
|
# set variable value
|
|
# value is converted to string
|
|
# returns 0 for success, -1 if failed
|
|
|
|
def set_variable(self,name,value):
|
|
return self.lib.lammps_set_variable(self.lmp,name,str(value))
|
|
|
|
# return total number of atoms in system
|
|
|
|
def get_natoms(self):
|
|
return self.lib.lammps_get_natoms(self.lmp)
|
|
|
|
# return vector of atom properties gathered across procs, ordered by atom ID
|
|
|
|
def gather_atoms(self,name,type,count):
|
|
natoms = self.lib.lammps_get_natoms(self.lmp)
|
|
if type == 0:
|
|
data = ((count*natoms)*c_int)()
|
|
self.lib.lammps_gather_atoms(self.lmp,name,type,count,data)
|
|
elif type == 1:
|
|
data = ((count*natoms)*c_double)()
|
|
self.lib.lammps_gather_atoms(self.lmp,name,type,count,data)
|
|
else: return None
|
|
return data
|
|
|
|
# scatter vector of atom properties across procs, ordered by atom ID
|
|
# assume vector is of correct type and length, as created by gather_atoms()
|
|
|
|
def scatter_atoms(self,name,type,count,data):
|
|
self.lib.lammps_scatter_atoms(self.lmp,name,type,count,data)
|