class MPISharedParticlesProxy(object):
def __init__(self, particles, shared_id, concurrent_processes):
self.shared_id = shared_id
self.particles = particles
self.concurrent_processes = concurrent_processes
def __getattr__(self, name):
return self.particles.__getattr__(name)
def distribute(self):
self.concurrent_processes.call(
self.distribute_on_root,
self.distribute_not_on_root
)
def distribute_on_root(self):
attribute_names = self.particles.get_attribute_names_defined_in_store()
values = self.particles.get_values_in_store(
self.particles.get_all_indices_in_store(),
attribute_names
)
units = [x.unit for x in values]
units_dump = pickle.dumps(units)
attributes_dump = pickle.dumps(attribute_names)
units_dump = numpy.fromstring(units_dump,dtype='uint8')
attributes_dump = numpy.fromstring(attributes_dump,dtype='uint8')
sendbuffer = numpy.zeros(4, dtype='int64')
sendbuffer[0] = self.shared_id
sendbuffer[1] = len(self.particles)
sendbuffer[2] = len(units_dump)
sendbuffer[3] = len(attributes_dump)
self.concurrent_processes.mpi_comm.Bcast([sendbuffer, MPI.INTEGER8], root=self.concurrent_processes.ROOT)
sendbuffer = self.particles.key
self.concurrent_processes.mpi_comm.Bcast([sendbuffer, MPI.INTEGER8], root=self.concurrent_processes.ROOT)
attribute_names = self.particles.get_attribute_names_defined_in_store()
self.concurrent_processes.mpi_comm.Bcast([units_dump, MPI.CHARACTER], root=self.concurrent_processes.ROOT)
self.concurrent_processes.mpi_comm.Bcast([attributes_dump, MPI.CHARACTER], root=self.concurrent_processes.ROOT)
for x, unit in zip(values, units):
value = x.value_in(unit)
self.concurrent_processes.mpi_comm.Bcast([value, MPI.DOUBLE], root=self.concurrent_processes.ROOT)
def distribute_not_on_root(self):
sendbuffer = numpy.zeros(4, dtype='int64')
self.concurrent_processes.mpi_comm.Bcast([sendbuffer, MPI.INTEGER8], root=self.concurrent_processes.ROOT)
shared_id = sendbuffer[0]
number_of_particles = sendbuffer[1]
units_dump_len = sendbuffer[2]
attributes_dump_len = sendbuffer[3]
sendbuffer = numpy.zeros(number_of_particles, dtype='int64')
self.concurrent_processes.mpi_comm.Bcast([sendbuffer, MPI.INTEGER8], root=self.concurrent_processes.ROOT)
units_dump = numpy.zeros(units_dump_len, dtype='uint8')
self.concurrent_processes.mpi_comm.Bcast([units_dump, MPI.CHARACTER], root=self.concurrent_processes.ROOT)
attributes_dump = numpy.zeros(attributes_dump_len, dtype='uint8')
self.concurrent_processes.mpi_comm.Bcast([attributes_dump, MPI.CHARACTER], root=self.concurrent_processes.ROOT)
units = pickle.loads(units_dump.tostring())
attributes = pickle.loads(attributes_dump.tostring())
values = []
for x in units:
value = numpy.zeros(number_of_particles, dtype='float64')
self.concurrent_processes.mpi_comm.Bcast([value, MPI.DOUBLE], root=self.concurrent_processes.ROOT)
values.append(x.new_quantity(value))
self.particles = Particles(keys = sendbuffer)
self.particles.set_values_in_store(self.particles.get_all_indices_in_store(), attributes, values)
def potential_energy(self, smoothing_length_squared = zero, G = constants.G):
mpi_comm = self.concurrent_processes.mpi_comm
mass = self.mass
x_vector = self.x
y_vector = self.y
z_vector = self.z
sum_of_energies = zero
number_of_particles = len(self)
block_size = (number_of_particles - 1) / mpi_comm.size
start = mpi_comm.rank * block_size
if mpi_comm.rank == (mpi_comm.size - 1):
block_size = (number_of_particles - 1) - start
for i in range(start, start + block_size):
x = x_vector[i]
y = y_vector[i]
z = z_vector[i]
dx = x - x_vector[i+1:]
dy = y - y_vector[i+1:]
dz = z - z_vector[i+1:]
dr_squared = (dx * dx) + (dy * dy) + (dz * dz)
dr = (dr_squared+smoothing_length_squared).sqrt()
m_m = mass[i] * mass[i+1:]
energy_of_this_particle = (m_m / dr).sum()
sum_of_energies -= energy_of_this_particle
value = sum_of_energies.value_in(sum_of_energies.unit)
# for not assume unit is the same accross processes,
# so units are not distributed!
input = numpy.zeros(1, dtype='float64')
output = numpy.zeros(1, dtype='float64')
input[0] = value
mpi_comm.Reduce(
[input, MPI.DOUBLE],
[output, MPI.DOUBLE],
op=MPI.SUM,
root=0
)
return G * sum_of_energies.unit.new_quantity(output[0])
def __len__(self):
return len(self.particles)
