def __init__(self,
extent=None,
periods=(1, 1, 1),
comm=mpi.MPI.COMM_WORLD,
boundary_condition=None):
self._init_cells = False
self._init_decomp = False
self.version_id = 0
self._periods = periods
self._extent = data.ScalarArray(ncomp=3, dtype=ctypes.c_double)
self._extent_global = data.ScalarArray(ncomp=3, dtype=ctypes.c_double)
self._boundary_outer = None
self._boundary = None
self._init_extent = False
if extent is not None:
self.set_extent(extent)
self._cell_array = data.ScalarArray(np.array([1, 1, 1]),
dtype=ctypes.c_int)
self._cell_edge_lengths = data.ScalarArray(
np.array([1., 1., 1.], dtype=ctypes.c_double))
self._halos = True
self.boundary_condition = boundary_condition
#vars to return boudary cells
self._boundary_cell_version = -1
self._boundary_cells = None
self.comm = None
self._parent_comm = comm
def gather_data_on(self, rank=0):
#seen from MPI_COMM_WORLD
assert (rank > -1) and (rank < _MPISIZE), "Invalid mpi rank"
if _MPISIZE == 1:
return
else:
esize = ctypes.sizeof(self.idtype)
counts = _MPIWORLD.gather(self.npart_local, root=rank)
_ptr_new = 0
if _MPIRANK == rank:
_new_nloc = sum(counts)
_new = cuda_base._create_zeros(nrow=_new_nloc,
ncol=self.ncomp,
dtype=self.idtype)
_ptr_new = _new.ptr
disp = [0] + counts[:-1:]
disp = tuple(np.cumsum(self.ncomp * np.array(disp)))
counts = tuple([self.ncomp * c for c in counts])
ln = _MPISIZE
disp_ = data.ScalarArray(dtype=ctypes.c_int, ncomp=ln)
counts_ = data.ScalarArray(dtype=ctypes.c_int, ncomp=ln)
disp_[:] = esize * np.array(disp)
counts_[:] = esize * np.array(counts)
disp = disp_
counts = counts_
else:
disp = data.ScalarArray(dtype=ctypes.c_int, ncomp=_MPISIZE)
counts = data.ScalarArray(dtype=ctypes.c_int, ncomp=_MPISIZE)
send_count = ctypes.c_int(esize * self.npart_local * self.ncomp)
cuda_mpi.MPI_Gatherv(
_MPIWORLD, ctypes.cast(self.ctypes_data,
ctypes.c_void_p), send_count,
ctypes.cast(_ptr_new, ctypes.c_void_p), counts.ctypes_data,
disp.ctypes_data, ctypes.c_int(rank))
if _MPIRANK == rank:
self.npart_local = _new_nloc
self._ncol.value = self.ncomp
self._nrow.value = _new_nloc
self._dat = _new
def __init__(self, state, size=0, v0=None):
self._state = state
self._V0 = data.ParticleDat(self._state.npart_local, 3, name='v0')
self._VT = state.velocities
self._VO_SET = False
if v0 is not None:
self.set_v0(v0)
else:
self.set_v0(state=self._state)
self._VAF = data.ScalarArray(ncomp=1)
self._V = []
self._T = []
_headers = ['stdio.h']
_constants = None
_kernel_code = '''
VAF(0) += (v0(0)*VT(0) + v0(1)*VT(1) + v0(2)*VT(2))*Ni;
'''
_reduction = (kernel.Reduction('VAF', 'VAF[I]', '+'), )
_static_args = {'Ni': ctypes.c_double}
_kernel = kernel.Kernel('VelocityAutocorrelation', _kernel_code,
_constants, _headers, _reduction, _static_args)
self._datdict = {'VAF': self._VAF, 'v0': self._V0, 'VT': self._VT}
self._loop = loop.ParticleLoop(self._state.as_func('npart_local'),
None,
kernel=_kernel,
dat_dict=self._datdict)
def __init__(self,
velocities=None,
masses=None,
kinetic_energy_dat=None,
looping_method=None):
if looping_method is None:
looping_method = loop.ParticleLoop
if kinetic_energy_dat is None:
self.k = data.ScalarArray(ncomp=1, dtype=ctypes.c_double)
else:
self.k = kinetic_energy_dat
self._v = velocities
if looping_method is None:
looping_method = loop.ParticleLoop
_K_kernel_code = '''
k(0) += (V(0)*V(0) + V(1)*V(1) + V(2)*V(2))*0.5*M(0);
'''
_constants_K = []
_K_kernel = kernel.Kernel('K_kernel', _K_kernel_code, _constants_K)
self._kinetic_energy_lib = looping_method(kernel=_K_kernel,
dat_dict={
'V':
velocities(access.R),
'k': self.k(access.INC),
'M': masses(access.R)
})
self._ke_store = []
def _init_escape_lib(self):
''' Create a lookup table between xor map and linear index for direction '''
self._bin_to_lin = data.ScalarArray(ncomp=57, dtype=ctypes.c_int)
_lin_to_bin = np.zeros(26, dtype=ctypes.c_int)
'''linear to xor map'''
_lin_to_bin[0] = 1 ^ 2 ^ 4
_lin_to_bin[1] = 2 ^ 1
_lin_to_bin[2] = 32 ^ 2 ^ 1
_lin_to_bin[3] = 4 ^ 1
_lin_to_bin[4] = 1
_lin_to_bin[5] = 32 ^ 1
_lin_to_bin[6] = 4 ^ 1 ^ 16
_lin_to_bin[7] = 1 ^ 16
_lin_to_bin[8] = 32 ^ 16 ^ 1
_lin_to_bin[9] = 2 ^ 4
_lin_to_bin[10] = 2
_lin_to_bin[11] = 32 ^ 2
_lin_to_bin[12] = 4
_lin_to_bin[13] = 32
_lin_to_bin[14] = 4 ^ 16
_lin_to_bin[15] = 16
_lin_to_bin[16] = 32 ^ 16
_lin_to_bin[17] = 8 ^ 2 ^ 4
_lin_to_bin[18] = 2 ^ 8
_lin_to_bin[19] = 32 ^ 2 ^ 8
_lin_to_bin[20] = 4 ^ 8
_lin_to_bin[21] = 8
_lin_to_bin[22] = 32 ^ 8
_lin_to_bin[23] = 4 ^ 8 ^ 16
_lin_to_bin[24] = 8 ^ 16
_lin_to_bin[25] = 32 ^ 16 ^ 8
'''inverse map, probably not ideal'''
for ix in range(26):
self._bin_to_lin[_lin_to_bin[ix]] = ix
# Number of escaping particles in each direction
self._escape_count = host.Array(np.zeros(26), dtype=ctypes.c_int)
# Linked list to store the ids of escaping particles in a similar way
# to the cell list.
# | [0-25 escape directions, index of first in direction] [26-end
# current id and index of next id, (id, next_index) ]|
self._escape_linked_list = host.Array(
-1 * np.ones(26 + 2 * self.state.npart_local), dtype=ctypes.c_int)
dtype = self.state.get_position_dat().dtype
assert self.state.domain.boundary.dtype == dtype
self._escape_guard_lib = ppmd.lib.build.lib_from_file_source(
_LIB_SOURCES + 'EscapeGuard', 'EscapeGuard', {
'SUB_REAL': self.state.get_position_dat().ctype,
'SUB_INT': self._bin_to_lin.ctype
})['EscapeGuard']
def _distribute_domain(self):
_top = mpi.cartcomm_top_xyz(self.comm)
_dims = mpi.cartcomm_dims_xyz(self.comm)
opt.PROFILE[self.__class__.__name__ + ':mpi_dims'] = (_dims)
self._extent[0] = self._extent_global[0] / _dims[0]
self._extent[1] = self._extent_global[1] / _dims[1]
self._extent[2] = self._extent_global[2] / _dims[2]
_boundary = (-0.5 * self._extent_global[0] + _top[0] * self._extent[0],
-0.5 * self._extent_global[0] +
(_top[0] + 1.) * self._extent[0],
-0.5 * self._extent_global[1] + _top[1] * self._extent[1],
-0.5 * self._extent_global[1] +
(_top[1] + 1.) * self._extent[1],
-0.5 * self._extent_global[2] + _top[2] * self._extent[2],
-0.5 * self._extent_global[2] +
(_top[2] + 1.) * self._extent[2])
self._boundary = data.ScalarArray(_boundary, dtype=ctypes.c_double)
self._boundary_outer = data.ScalarArray(_boundary,
dtype=ctypes.c_double)
def __init__(self, epsilon=1.0, sigma=1.0, rc=None):
self._epsilon = epsilon
self._sigma = sigma
self._C_V = 4. * self._epsilon
self._C_F = -48 * self._epsilon / self._sigma**2
if rc is None:
self._rc = self._sigma * (5. / 2.)
else:
self._rc = rc
self._rn = 1. * self._rc
self._rc2 = self._rc**2
self._sigma2 = self._sigma**2
self._shift_internal = (self._sigma / self._rc)**6 - (self._sigma /
self._rc)**12
self._counter = data.ScalarArray([0],
dtype=ctypes.c_longlong,
name="counter")
self._counter.data[0] = 0
self._counter_outer = data.ScalarArray([0],
dtype=ctypes.c_longlong,
name="counter")
self._counter_outer.data[0] = 0
def cell_decompose(self, cell_width=None):
assert cell_width is not None, "ERROR: No cell size passed."
assert cell_width > 10.**-14, "ERROR: requested cell size stupidly small."
if not self._init_decomp:
print(
"WARNING: domain not spatial decomposed, see mpi_decompose()")
cell_width = float(cell_width)
self._cell_array[0] = int(self._extent[0] / cell_width)
self._cell_array[1] = int(self._extent[1] / cell_width)
self._cell_array[2] = int(self._extent[2] / cell_width)
assert self._cell_array[0] > 0, "Too many MPI ranks in dir 0"
assert self._cell_array[1] > 0, "Too many MPI ranks in dir 1"
assert self._cell_array[2] > 0, "Too many MPI ranks in dir 2"
self._cell_edge_lengths[0] = self._extent[0] / self._cell_array[0]
self._cell_edge_lengths[1] = self._extent[1] / self._cell_array[1]
self._cell_edge_lengths[2] = self._extent[2] / self._cell_array[2]
self._cell_array[0] += 2
self._cell_array[1] += 2
self._cell_array[2] += 2
_boundary = (self._boundary[0] - self._cell_edge_lengths[0],
self._boundary[1] + self._cell_edge_lengths[0],
self._boundary[2] - self._cell_edge_lengths[1],
self._boundary[3] + self._cell_edge_lengths[1],
self._boundary[4] - self._cell_edge_lengths[2],
self._boundary[5] + self._cell_edge_lengths[2])
self._boundary_outer = data.ScalarArray(_boundary,
dtype=ctypes.c_double)
self._init_cells = True
opt.PROFILE[self.__class__.__name__ +
':cell_array'] = (self.cell_array[:])
self.version_id += 1
return True
def __init__(self, positions, charges, domain, boundary_condition, r, l):
self.positions = positions
self.charges = charges
self.domain = domain
self.comm = self.domain.comm
self.boundary_condition = BCType(boundary_condition)
self.R = r
self.L = l
self.ncomp = (self.L ** 2) * 2
self.group = self.positions.group
self.subdivision = (2, 2, 2)
# interaction lists
self.il = fmm_interaction_lists.compute_interaction_lists(domain.extent, self.subdivision)
self.il_max_len = max(len(lx) for lx in self.il[0])
self.il_array = np.array(self.il[0], INT64)
self.il_scalararray = data.ScalarArray(ncomp=self.il_array.size, dtype=INT64)
self.il_scalararray[:] = self.il_array.ravel().copy()
self.il_earray = np.array(self.il[1], INT64)
# tree is stored in a GlobalArray
n = 0
for rx in range(self.R):
s = [sx ** rx for sx in self.subdivision]
n += s[0] * s[1] * s[2] * self.ncomp
self.tree = data.GlobalArray(ncomp=n, dtype=REAL)
self._init_dats()
s = self.subdivision
s = [sx ** (self.R - 1) for sx in s]
ncells_finest = s[0] * s[1] * s[2]
self.cell_occupation_ga = data.GlobalArray(ncomp=ncells_finest, dtype=INT64)
self.sph_gen = SphGen(l-1)
self.widths_x = data.ScalarArray(ncomp=self.R, dtype=REAL)
self.widths_y = data.ScalarArray(ncomp=self.R, dtype=REAL)
self.widths_z = data.ScalarArray(ncomp=self.R, dtype=REAL)
e = self.domain.extent
s = self.subdivision
self.widths_x[:] = [e[0] / (s[0] ** rx) for rx in range(self.R)]
self.widths_y[:] = [e[1] / (s[1] ** rx) for rx in range(self.R)]
self.widths_z[:] = [e[2] / (s[2] ** rx) for rx in range(self.R)]
self.ncells_x = data.ScalarArray(ncomp=self.R, dtype=INT64)
self.ncells_y = data.ScalarArray(ncomp=self.R, dtype=INT64)
self.ncells_z = data.ScalarArray(ncomp=self.R, dtype=INT64)
self.ncells_x[:] = [int(s[0] ** rx) for rx in range(self.R)]
self.ncells_y[:] = [int(s[1] ** rx) for rx in range(self.R)]
self.ncells_z[:] = [int(s[2] ** rx) for rx in range(self.R)]
self._extract_energy = data.GlobalArray(ncomp=1, dtype=REAL)
# find the max distance in the direct part using the exclusion lists
widths = [ex / (sx ** (self.R - 1)) for ex, sx in zip(self.domain.extent, self.subdivision)]
max_cell_width = max(
[(abs(ix[0]) + 1) * widths[0] for ix in self.il[1]] + \
[(abs(ix[1]) + 1) * widths[1] for ix in self.il[1]] + \
[(abs(ix[2]) + 1) * widths[2] for ix in self.il[1]]
)
self.max_cell_width = max_cell_width
assert max_cell_width <= np.min(self.domain.extent[:])
self.max_il_offset = np.max(np.abs(self.il[0]))
self.widths_x_str = ','.join([str(ix) for ix in self.widths_x])
self.widths_y_str = ','.join([str(ix) for ix in self.widths_y])
self.widths_z_str = ','.join([str(ix) for ix in self.widths_z])
self.ncells_x_str = ','.join([str(ix) for ix in self.ncells_x])
self.ncells_y_str = ','.join([str(ix) for ix in self.ncells_y])
self.ncells_z_str = ','.join([str(ix) for ix in self.ncells_z])
level_offsets = [0]
nx = 1
ny = 1
nz = 1
for level in range(1, self.R):
level_offsets.append(
level_offsets[-1] + nx * ny * nz * self.ncomp
)
nx *= self.subdivision[0]
ny *= self.subdivision[1]
nz *= self.subdivision[2]
self.level_offsets_str = ','.join([str(ix) for ix in level_offsets])
self._contrib_loop = None
self._init_contrib_loop()
self._extract_loop = None
self._init_extract_loop()
self.sh = pairloop.state_handler.StateHandler(state=None, shell_cutoff=max_cell_width)
self.cell_list = np.zeros(1000, INT64)
self.cell_occ = np.zeros(1000, INT64)
self.cell_remaps = np.zeros((1000, 3), INT64)
self._direct_lib = None
self._cell_remap_lib = None
self._init_direct_libs()
if self.boundary_condition == BCType.PBC:
self._init_pbc()
def __init__(self, positions, charges, domain, boundary_condition, r, l):
"""
This class implements protype code for the implementation. For production use see the MCFMM_LM class.
"""
self.positions = positions
self.charges = charges
self.domain = domain
self.comm = self.domain.comm
self.boundary_condition = BCType(boundary_condition)
self.R = r
self.L = l
self.ncomp = (self.L**2) * 2
self.subdivision = (2, 2, 2)
# tree
self.tree = octal.OctalTree(self.R, domain.comm)
self.tree_local = octal.OctalDataTree(self.tree, self.ncomp, 'plain',
REAL)
self.tree_local_ptrs = np.zeros(self.R, ctypes.c_void_p)
for rx in range(self.R):
self.tree_local_ptrs[rx] = self.tree_local[
rx].ctypes.get_as_parameter().value
self.tree_local_ga = data.GlobalArray(
ncomp=self.tree_local.num_cells() * self.ncomp, dtype=REAL)
self.tree_local_ga_offsets = data.ScalarArray(ncomp=self.R + 1,
dtype=INT64)
t = 0
self.tree_local_ga_offsets[0] = 0
for rx in range(self.R):
t += self.tree_local.num_data[rx]
self.tree_local_ga_offsets[rx + 1] = t
s = self.subdivision
s = [sx**(self.R - 1) for sx in s]
ncells_finest = s[0] * s[1] * s[2]
self.cell_occupation_ga = data.GlobalArray(ncomp=ncells_finest,
dtype=INT64)
# interaction lists
self.il = fmm_interaction_lists.compute_interaction_lists(
domain.extent, self.subdivision)
self.il_max_len = max(len(lx) for lx in self.il[0])
self.il_array = np.array(self.il[0], INT64)
self.il_scalararray = data.ScalarArray(ncomp=self.il_array.size,
dtype=INT64)
self.il_scalararray[:] = self.il_array.ravel().copy()
# expansion tools
self.lee = kmc_expansion_tools.LocalExpEval(self.L)
self.mc_lee = mc_expansion_tools.LocalExp(self.L)
self.group = self.positions.group
pd = type(self.charges)
g = self.group
l = self.il_max_len * self.R
# assume that ptr size and int64_t size may be different....
if ctypes.sizeof(ctypes.c_void_p) == ctypes.sizeof(INT64):
self.il_pd_ptr_stride = l
ptr_l = l
else:
ptr_l = int(
math.ceil((l * ctypes.sizeof(ctypes.c_void_p)) /
ctypes.sizeof(INT64)))
self.il_pd_ptr_stride = int(ptr_l * ctypes.sizeof(INT64) //
ctypes.sizeof(ctypes._c_void_p))
g._mc_fmm_cells = pd(ncomp=3, dtype=INT64)
g._mc_ids = pd(ncomp=1, dtype=INT64)
g._mc_nexp = pd(ncomp=1, dtype=INT64)
g._mc_charge = pd(ncomp=l, dtype=REAL)
g._mc_radius = pd(ncomp=l, dtype=REAL)
g._mc_theta = pd(ncomp=l, dtype=REAL)
g._mc_phi = pd(ncomp=l, dtype=REAL)
g._mc_du = pd(ncomp=l, dtype=REAL)
g._mc_level = pd(ncomp=l, dtype=INT64)
g._mc_cx = pd(ncomp=l, dtype=INT64)
g._mc_cy = pd(ncomp=l, dtype=INT64)
g._mc_cz = pd(ncomp=l, dtype=INT64)
g._mc_cl = pd(ncomp=l, dtype=INT64)
g._mc_ptrs = pd(ncomp=ptr_l, dtype=INT64)
tmp = 0
for datx in g.particle_dats:
tmp += getattr(g, datx)._dat.nbytes
self.dat_size = tmp
width = self.domain.extent[0] / (2**(self.R - 1))
self.sh = pairloop.state_handler.StateHandler(state=None,
shell_cutoff=width,
pair=False)
self.direct = DirectCommon(positions, charges, domain,
boundary_condition, r, self.subdivision,
g._mc_fmm_cells, g._mc_ids, self.sh)
self.energy = None
self._cell_bin_loop = None
self._init_bin_loop()
self._init_indirect_accept_lib()
self._init_indirect_propose_lib()
def __init__(self,
domain,
eps=10.**-6,
real_cutoff=None,
alpha=None,
recip_cutoff=None,
recip_nmax=None,
shared_memory=False,
shell_width=None,
work_ratio=1.0,
force_unit=1.0,
energy_unit=1.0):
self.domain = domain
self.eps = float(eps)
assert shared_memory in (False, 'omp', 'mpi')
ss = cmath.sqrt(scipy.special.lambertw(1. / eps)).real
if alpha is not None and real_cutoff is not None and recip_cutoff is not None:
pass
elif alpha is not None and real_cutoff is not None:
ss = real_cutoff * sqrt(alpha)
elif alpha is None:
alpha = (ss / real_cutoff)**2.
else:
real_cutoff = ss / sqrt(alpha)
assert alpha is not None, "no alpha deduced/passed"
assert real_cutoff is not None, "no real_cutoff deduced/passed"
self.real_cutoff = float(real_cutoff)
"""Real space cutoff"""
self.shell_width = shell_width
"""Real space padding width"""
self.alpha = float(alpha)
"""alpha"""
#self.real_cutoff = float(real_cutoff)
#alpha = 0.2
#print("alpha", alpha)
#print("r_c", self.real_cutoff)
# these parts are specific to the orthongonal box
extent = self.domain.extent
lx = (extent[0], 0., 0.)
ly = (0., extent[1], 0.)
lz = (0., 0., extent[2])
ivolume = 1. / np.dot(lx, np.cross(ly, lz))
gx = np.cross(ly, lz) * ivolume * 2. * pi
gy = np.cross(lz, lx) * ivolume * 2. * pi
gz = np.cross(lx, ly) * ivolume * 2. * pi
sqrtalpha = sqrt(alpha)
nmax_x = round(ss * extent[0] * sqrtalpha / pi)
nmax_y = round(ss * extent[1] * sqrtalpha / pi)
nmax_z = round(ss * extent[2] * sqrtalpha / pi)
#print gx, gy, gz
#print 'nmax:', nmax_x, nmax_y, nmax_z
#print "alpha", alpha, "sqrt(alpha)", sqrtalpha
gxl = np.linalg.norm(gx)
gyl = np.linalg.norm(gy)
gzl = np.linalg.norm(gz)
if recip_cutoff is None:
max_len = min(gxl * float(nmax_x), gyl * float(nmax_y),
gzl * float(nmax_z))
else:
max_len = recip_cutoff
if recip_nmax is None:
nmax_x = int(ceil(max_len / gxl))
nmax_y = int(ceil(max_len / gyl))
nmax_z = int(ceil(max_len / gzl))
else:
nmax_x = recip_nmax[0]
nmax_y = recip_nmax[1]
nmax_z = recip_nmax[2]
#print 'max reciprocal vector len:', max_len
nmax_t = max(nmax_x, nmax_y, nmax_z)
#print "nmax_t", nmax_t
self.last_real_energy = None
self.last_recip_energy = None
self.last_self_energy = None
self.kmax = (nmax_x, nmax_y, nmax_z)
"""Number of reciporcal vectors taken in each direction."""
#print("kmax", self.kmax)
self.recip_cutoff = max_len
"""Reciprocal space cutoff."""
self.recip_vectors = (gx, gy, gz)
"""Reciprocal lattice vectors"""
self.ivolume = ivolume
opt.PROFILE[self.__class__.__name__ +
':recip_vectors'] = (self.recip_vectors)
opt.PROFILE[self.__class__.__name__ +
':recip_cutoff'] = (self.recip_cutoff)
opt.PROFILE[self.__class__.__name__ + ':recip_kmax'] = (self.kmax)
opt.PROFILE[self.__class__.__name__ + ':alpha'] = (self.alpha)
opt.PROFILE[self.__class__.__name__ + ':tol'] = (eps)
opt.PROFILE[self.__class__.__name__ +
':real_cutoff'] = (self.real_cutoff)
# define persistent vars
self._vars = {}
self._vars['alpha'] = ctypes.c_double(alpha)
self._vars['max_recip'] = ctypes.c_double(max_len)
self._vars['nmax_vec'] = host.Array((nmax_x, nmax_y, nmax_z),
dtype=ctypes.c_int)
self._vars['recip_vec'] = host.Array(
np.zeros((3, 3), dtype=ctypes.c_double))
self._vars['recip_vec'][0, :] = gx
self._vars['recip_vec'][1, :] = gy
self._vars['recip_vec'][2, :] = gz
self._vars['ivolume'] = ivolume
self._vars['coeff_space_kernel'] = data.ScalarArray(
ncomp=((nmax_x + 1) * (nmax_y + 1) * (nmax_z + 1)),
dtype=ctypes.c_double)
self._vars['coeff_space'] = self._vars['coeff_space_kernel'].data.view(
).reshape(nmax_z + 1, nmax_y + 1, nmax_x + 1)
#self._vars['coeff_space'] = np.zeros((nmax_z+1, nmax_y+1, nmax_x+1), dtype=ctypes.c_double)
# pass stride in tmp space vector
self._vars['recip_axis_len'] = ctypes.c_int(nmax_t)
# |axis | planes | quads
reciplen = (nmax_t+1)*12 +\
8*nmax_x*nmax_y + \
8*nmax_y*nmax_z +\
8*nmax_z*nmax_x +\
16*nmax_x*nmax_y*nmax_z
self._vars['recip_space_kernel'] = data.GlobalArray(
size=reciplen, dtype=ctypes.c_double, shared_memory=shared_memory)
self._vars['recip_space_energy'] = data.GlobalArray(
size=1, dtype=ctypes.c_double, shared_memory=shared_memory)
self._vars['real_space_energy'] = data.GlobalArray(
size=1, dtype=ctypes.c_double, shared_memory=shared_memory)
self._vars['self_interaction_energy'] = data.GlobalArray(
size=1, dtype=ctypes.c_double, shared_memory=shared_memory)
self.shared_memory = shared_memory
#self._vars['recip_vec_kernel'] = data.ScalarArray(np.zeros(3, dtype=ctypes.c_double))
#self._vars['recip_vec_kernel'][0] = gx[0]
#self._vars['recip_vec_kernel'][1] = gy[1]
#self._vars['recip_vec_kernel'][2] = gz[2]
self._subvars = dict()
self._subvars['SUB_GX'] = str(gx[0])
self._subvars['SUB_GY'] = str(gy[1])
self._subvars['SUB_GZ'] = str(gz[2])
self._subvars['SUB_NKMAX'] = str(nmax_t)
self._subvars['SUB_NK'] = str(nmax_x)
self._subvars['SUB_NL'] = str(nmax_y)
self._subvars['SUB_NM'] = str(nmax_z)
self._subvars['SUB_NKAXIS'] = str(nmax_t)
self._subvars['SUB_LEN_QUAD'] = str(nmax_x * nmax_y * nmax_z)
self._subvars['SUB_MAX_RECIP'] = str(max_len)
self._subvars['SUB_MAX_RECIP_SQ'] = str(max_len**2.)
self._subvars['SUB_SQRT_ALPHA'] = str(sqrt(alpha))
self._subvars['SUB_REAL_CUTOFF_SQ'] = str(real_cutoff**2.)
self._subvars['SUB_REAL_CUTOFF'] = str(real_cutoff)
self._subvars['SUB_M_SQRT_ALPHA_O_PI'] = str(-1.0 * sqrt(alpha / pi))
self._subvars['SUB_M2_SQRT_ALPHAOPI'] = str(-2.0 * sqrt(alpha / pi))
self._subvars['SUB_MALPHA'] = str(-1.0 * alpha)
self._subvars['SUB_ENERGY_UNIT'] = str(energy_unit)
self._subvars['SUB_ENERGY_UNITO2'] = str(energy_unit * 0.5)
self._subvars['SUB_FORCE_UNIT'] = str(force_unit)
self._real_space_pairloop = None
self._init_libs()
self._init_coeff_space()
self._self_interaction_lib = None
def __init__(self, state, rmax=None, rsteps=100):
self._count = 0
self._state = state
self._extent = self._state.domain.extent
self._P = self._state.positions
self._N = self._state.npart_local
self._rmax = rmax
if self._rmax is None:
self._rmax = 0.5 * np.min(self._extent.data)
self._rsteps = rsteps
self._gr = data.ScalarArray(ncomp=self._rsteps, dtype=ctypes.c_int)
self._gr.scale(0.0)
_headers = ['math.h', 'stdio.h']
_kernel = '''
double R0 = P(1, 0) - P(0, 0);
double R1 = P(1, 1) - P(0, 1);
double R2 = P(1, 2) - P(0, 2);
if (abs_md(R0) > exto20 ) { R0 += isign(R0) * extent0 ; }
if (abs_md(R1) > exto21 ) { R1 += isign(R1) * extent1 ; }
if (abs_md(R2) > exto22 ) { R2 += isign(R2) * extent2 ; }
const double r2 = R0*R0 + R1*R1 + R2*R2;
if (r2 < rmax2){
double r20=0.0, r21 = r2;
r21 = sqrt(r2);
#pragma omp atomic
GR[(int) (abs_md(r21* rstepsoverrmax))]++;
}
'''
_constants = (kernel.Constant('rmaxoverrsteps',
0.2 * self._rmax / self._rsteps),
kernel.Constant('rstepsoverrmax',
self._rsteps / self._rmax),
kernel.Constant('rmax2', self._rmax**2),
kernel.Constant('extent0', self._extent[0]),
kernel.Constant('extent1', self._extent[1]),
kernel.Constant('extent2', self._extent[2]),
kernel.Constant('exto20', 0.5 * self._extent[0]),
kernel.Constant('exto21', 0.5 * self._extent[1]),
kernel.Constant('exto22', 0.5 * self._extent[2]))
_grkernel = kernel.Kernel('radial_distro_periodic_static',
_kernel,
_constants,
headers=_headers)
_datdict = {'P': self._P, 'GR': self._gr}
self._p = pairloop.DoubleAllParticleLoop(self._N,
kernel=_grkernel,
dat_dict=_datdict)
self.timer = ppmd.opt.Timer(runtime.TIMER, 0)
def plot(self, _plot=True):
"""
Plot the stored energy data.
:return:
"""
assert len(self._t) > 0, "EnergyStore error, no data to plot"
self._T_store = data.ScalarArray(self._t)
self._K_store = data.ScalarArray(self._k)
self._U_store = data.ScalarArray(self._u)
self._Q_store = data.ScalarArray(self._q)
'''REPLACE THIS WITH AN MPI4PY REDUCE CALL'''
if _MPISIZE > 1:
# data to collect
_d = [self._Q_store.data, self._U_store.data, self._K_store.data]
# make a temporary buffer.
if _MPIRANK == 0:
_buff = data.ScalarArray(ncomp=self._T_store.ncomp,
dtype=ctypes.c_double)
_T = self._T_store.data
_Q = data.ScalarArray(ncomp=self._T_store.ncomp,
dtype=ctypes.c_double)
_U = data.ScalarArray(ncomp=self._T_store.ncomp,
dtype=ctypes.c_double)
_K = data.ScalarArray(ncomp=self._T_store.ncomp,
dtype=ctypes.c_double)
_Q.data[::] += self._Q_store.data[::]
_U.data[::] += self._U_store.data[::]
_K.data[::] += self._K_store.data[::]
_dl = [_Q.data, _U.data, _K.data]
else:
_dl = [None, None, None]
for _di, _dj in zip(_d, _dl):
if _MPIRANK == 0:
_MS = mpi.Status()
for ix in range(1, _MPISIZE):
_MPIWORLD.Recv(_buff.data[::], ix, ix, _MS)
_dj[::] += _buff.data[::]
else:
_MPIWORLD.Send(_di[::], 0, _MPIRANK)
if _MPIRANK == 0:
_Q = _Q.data
_U = _U.data
_K = _K.data
else:
_T = self._T_store.data
_Q = self._Q_store.data
_U = self._U_store.data
_K = self._K_store.data
if (_MPIRANK == 0) and _GRAPHICS:
print("last total", _Q[-1])
print("last kinetic", _K[-1])
print("last potential", _U[-1])
print("=============================================")
print("first total", _Q[0])
print("first kinetic", _K[0])
print("first potential", _U[0])
if _plot:
plt.ion()
fig2 = plt.figure()
ax2 = fig2.add_subplot(111)
ax2.plot(_T, _Q, color='r', linewidth=2)
ax2.plot(_T, _U, color='g')
ax2.plot(_T, _K, color='b')
ax2.set_title(
'Red: Total energy, Green: Potential energy, Blue: kinetic energy'
)
ax2.set_xlabel('Time')
ax2.set_ylabel('Energy')
fig2.canvas.draw()
plt.show(block=False)
if _MPIRANK == 0:
if not os.path.exists(os.path.join(os.getcwd(), './output')):
os.system('mkdir ./output')
_fh = open('./output/energy.txt', 'w')
_fh.write("Time Kinetic Potential Total\n")
for ix in range(len(self._t)):
_fh.write("%(T)s %(K)s %(P)s %(Q)s\n" % {
'T': _T[ix],
'K': _K[ix],
'P': _U[ix],
'Q': _Q[ix]
})
_fh.close()
if (_MPIRANK == 0) and not _GRAPHICS:
print("last total", _Q[-1])
print("last kinetic", _K[-1])
print("last potential", _U[-1])
print("=============================================")
print("first total", _Q[0])
print("first kinetic", _K[0])
print("first potential", _U[0])
def draw(self):
"""
Update current plot, use for real time plotting.
"""
if _GRAPHICS:
self._N = self._state.npart_local
self._NT = self._state.npart
self._extents = self._state.domain.extent
'''Case where all particles are local'''
if _MPISIZE == 1:
self._pos = self._state.positions
self._gid = self._state.global_ids
else:
'''Need an mpi handle if not all particles are local'''
'''Allocate if needed'''
if self._Dat is None:
self._Dat = data.ParticleDat(self._NT, 3)
else:
self._Dat.resize(self._NT)
if self._gids is None:
self._gids = data.ScalarArray(ncomp=self._NT,
dtype=ctypes.c_int)
else:
self._gids.resize(self._NT)
_MS = mpi.Status()
if _MPIRANK == 0:
'''Copy the local data.'''
self._Dat.data[
0:self._N:, ::] = self._state.positions.data[
0:self._N:, ::]
self._gids.data[0:self._N:] = self._state.global_ids.data[
0:self._N:, 0]
_i = self._N # starting point pos
_ig = self._N # starting point gids
for ix in range(1, _MPISIZE):
_MPIWORLD.Recv(self._Dat.data[_i::, ::], ix, ix, _MS)
_i += _MS.Get_count(mpi.mpi_map[self._Dat.dtype]) // 3
_MPIWORLD.Recv(self._gids.data[_ig::], ix, ix, _MS)
_ig += _MS.Get_count(mpi.mpi_map[self._gids.dtype])
self._pos = self._Dat
self._gid = self._gids
else:
_MPIWORLD.Send(self._state.positions.data[0:self._N:, ::],
0, _MPIRANK)
_MPIWORLD.Send(self._state.global_ids.data[0:self._N:], 0,
_MPIRANK)
if _MPIRANK == 0:
plt.cla()
plt.ion()
for ix in range(self._pos.npart_local):
self._ax.scatter(self._pos.data[ix, 0],
self._pos.data[ix, 1],
self._pos.data[ix, 2],
color=self._key[self._gid[ix] % 2])
if _MPISIZE == 1:
self._ax.plot(
(self._pos.data[ix, 0], self._pos.data[ix, 0] +
self.norm_vec(self._state.forces.data[ix, 0])),
(self._pos.data[ix, 1], self._pos.data[ix, 1] +
self.norm_vec(self._state.forces.data[ix, 1])),
(self._pos.data[ix, 2], self._pos.data[ix, 2] +
self.norm_vec(self._state.forces.data[ix, 2])),
color=self._key[self._gid[ix] % 2],
linewidth=2)
self._ax.set_xlim(
[-0.5 * self._extents[0], 0.5 * self._extents[0]])
self._ax.set_ylim(
[-0.5 * self._extents[1], 0.5 * self._extents[1]])
self._ax.set_zlim(
[-0.5 * self._extents[2], 0.5 * self._extents[2]])
self._ax.set_xlabel('x')
self._ax.set_ylabel('y')
self._ax.set_zlabel('z')
plt.draw()
plt.show(block=False)