def _make_accel_eval(self, equations, cache_nnps=False):
arrays = [self.pa]
kernel = CubicSpline(dim=self.dim)
a_eval = AccelerationEval(
particle_arrays=arrays, equations=equations, kernel=kernel
)
comp = SPHCompiler(a_eval, integrator=None)
comp.compile()
nnps = NNPS(dim=kernel.dim, particles=arrays, cache=cache_nnps)
nnps.update()
a_eval.set_nnps(nnps)
return a_eval
def _make_accel_eval(self, equations):
arrays = [self.pa]
kernel = CubicSpline(dim=self.dim)
a_eval = AccelerationEval(
particle_arrays=arrays, equations=equations, kernel=kernel
)
comp = SPHCompiler(a_eval, integrator=None)
comp.compile()
nnps = NNPS(dim=kernel.dim, particles=arrays)
nnps.update()
a_eval.set_nnps(nnps)
return a_eval
class SPHEvaluator(object):
def __init__(self,
arrays,
equations,
dim,
kernel=None,
domain_manager=None):
"""Constructor.
Parameters
----------
arrays: list(ParticleArray)
equations: list
dim: int
kernel: kernel instance.
domain_manager: DomainManager
"""
self.arrays = arrays
self.equations = equations
self.domain_manager = domain_manager
self.dim = dim
if kernel is None:
self.kernel = Gaussian(dim=dim)
else:
self.kernel = kernel
self.func_eval = AccelerationEval(arrays, equations, self.kernel)
compiler = SPHCompiler(self.func_eval, None)
compiler.compile()
self._create_nnps(arrays)
def evaluate(self, t=0.0, dt=0.1):
"""Evalute the SPH equations, dummy t and dt values can
be passed.
"""
self.func_eval.compute(t, dt)
def update_particle_arrays(self, arrays):
self._create_nnps(arrays)
self.func_eval.update_particle_arrays(arrays)
#### Private protocol ###################################################
def _create_nnps(self, arrays):
self.nnps = NNPS(dim=self.kernel.dim,
particles=arrays,
radius_scale=self.kernel.radius_scale,
domain=self.domain_manager,
cache=True)
self.nnps.update()
self.func_eval.set_nnps(self.nnps)
def test_cache_updates_with_changed_particles(self):
# Given
pa1 = self._make_random_parray('pa1', 5)
particles = [pa1]
nnps = LinkedListNNPS(dim=3, particles=particles)
cache = NeighborCache(nnps, dst_index=0, src_index=0)
cache.update()
# When
pa2 = self._make_random_parray('pa2', 2)
pa1.add_particles(x=pa2.x, y=pa2.y, z=pa2.z)
nnps.update()
cache.update()
nb_cached = UIntArray()
nb_direct = UIntArray()
for i in range(len(particles[0].x)):
nnps.get_nearest_particles_no_cache(0, 0, i, nb_direct, False)
cache.get_neighbors(0, i, nb_cached)
nb_e = nb_direct.get_npy_array()
nb_c = nb_cached.get_npy_array()
self.assertTrue(np.all(nb_e == nb_c))
class Interpolator(object):
"""Convenient class to interpolate particle properties onto a uniform grid
or given set of particles. This is particularly handy for visualization.
"""
def __init__(self,
particle_arrays,
num_points=125000,
kernel=None,
x=None,
y=None,
z=None,
domain_manager=None,
equations=None):
"""
The x, y, z coordinates need not be specified, and if they are not,
the bounds of the interpolated domain is automatically computed and
`num_points` number of points are used in this domain uniformly placed.
Parameters
----------
particle_arrays: list
A list of particle arrays.
num_points: int
the number of points to interpolate on to.
kernel: Kernel
the kernel to use for interpolation.
x: ndarray
the x-coordinate of points on which to interpolate.
y: ndarray
the y-coordinate of points on which to interpolate.
z: ndarray
the z-coordinate of points on which to interpolate.
domain_manager: DomainManager
An optional Domain manager for periodic domains.
equations: sequence
A sequence of equations or groups. Defaults to None. This is
used only if the default interpolation equations are inadequate.
"""
self._set_particle_arrays(particle_arrays)
bounds = get_bounding_box(self.particle_arrays)
shape = get_nx_ny_nz(num_points, bounds)
self.dim = 3 - list(shape).count(1)
if kernel is None:
self.kernel = Gaussian(dim=self.dim)
else:
self.kernel = kernel
self.pa = None
self.nnps = None
self.equations = equations
self.func_eval = None
self.domain_manager = domain_manager
if x is None and y is None and z is None:
self.set_domain(bounds, shape)
else:
self.set_interpolation_points(x=x, y=y, z=z)
#### Interpolator protocol ################################################
def set_interpolation_points(self, x=None, y=None, z=None):
"""Set the points on which we must interpolate the arrays.
If any of x, y, z is not passed it is assumed to be 0.0 and shaped
like the other non-None arrays.
Parameters
----------
x: ndarray
the x-coordinate of points on which to interpolate.
y: ndarray
the y-coordinate of points on which to interpolate.
z: ndarray
the z-coordinate of points on which to interpolate.
"""
tmp = None
for tmp in (x, y, z):
if tmp is not None:
break
if tmp is None:
raise RuntimeError('At least one non-None array must be given.')
def _get_array(_t):
return np.asarray(_t) if _t is not None else np.zeros_like(tmp)
x, y, z = _get_array(x), _get_array(y), _get_array(z)
self.shape = x.shape
self.pa = self._create_particle_array(x, y, z)
arrays = self.particle_arrays + [self.pa]
if self.func_eval is None:
self._compile_acceleration_eval(arrays)
self.update_particle_arrays(self.particle_arrays)
def set_domain(self, bounds, shape):
"""Set the domain to interpolate into.
Parameters
----------
bounds: tuple
(xmin, xmax, ymin, ymax, zmin, zmax)
shape: tuple
(nx, ny, nz)
"""
self.bounds = np.asarray(bounds)
self.shape = np.asarray(shape)
x, y, z = self._create_default_points(self.bounds, self.shape)
self.set_interpolation_points(x, y, z)
def interpolate(self, prop, gradient=False):
"""Interpolate given property.
Parameters
----------
prop: str
The name of the property to interpolate.
gradient: bool
Evaluate gradient and not function.
Returns
-------
A numpy array suitably shaped with the property interpolated.
"""
for array in self.particle_arrays:
data = array.get(prop, only_real_particles=False)
array.get('temp_prop', only_real_particles=False)[:] = data
self.func_eval.compute(0.0, 0.1) # These are junk arguments.
result = self.pa.prop.copy()
result.shape = self.shape
return result.squeeze()
def update_particle_arrays(self, particle_arrays):
"""Call this for a new set of particle arrays which have the
same properties as before.
For example, if you are reading the particle array data from files,
each time you load a new file a new particle array is read with the
same properties. Call this function to reset the arrays.
"""
self._set_particle_arrays(particle_arrays)
arrays = self.particle_arrays + [self.pa]
self._create_nnps(arrays)
self.func_eval.update_particle_arrays(arrays)
#### Private protocol #####################################################
def _create_nnps(self, arrays):
# create the neighbor locator object
self.nnps = NNPS(dim=self.kernel.dim,
particles=arrays,
radius_scale=self.kernel.radius_scale,
domain=self.domain_manager,
cache=True)
self.nnps.update()
self.func_eval.set_nnps(self.nnps)
def _create_default_points(self, bounds, shape):
b = bounds
n = shape
x, y, z = np.mgrid[b[0]:b[1]:n[0] * 1j, b[2]:b[3]:n[1] * 1j,
b[4]:b[5]:n[2] * 1j, ]
return x, y, z
def _create_particle_array(self, x, y, z):
xr = x.ravel()
yr = y.ravel()
zr = z.ravel()
self.x, self.y, self.z = x.squeeze(), y.squeeze(), z.squeeze()
hmax = self._get_max_h_in_arrays()
h = hmax * np.ones_like(xr)
prop = np.zeros_like(xr)
pa = get_particle_array(name='interpolate',
x=xr,
y=yr,
z=zr,
h=h,
number_density=np.zeros_like(xr),
prop=prop,
grad_x=np.zeros_like(xr),
grad_y=np.zeros_like(xr),
grad_z=np.zeros_like(xr))
return pa
def _compile_acceleration_eval(self, arrays):
names = [x.name for x in self.particle_arrays]
if self.equations is None:
equations = [
InterpolateFunction(dest='interpolate', sources=names)
]
else:
equations = self.equations
self.func_eval = AccelerationEval(arrays, equations, self.kernel)
compiler = SPHCompiler(self.func_eval, None)
compiler.compile()
def _get_max_h_in_arrays(self):
hmax = -1.0
for array in self.particle_arrays:
hmax = max(array.h.max(), hmax)
return hmax
def _set_particle_arrays(self, particle_arrays):
self.particle_arrays = particle_arrays
self._make_all_arrays_have_same_props(particle_arrays)
for array in self.particle_arrays:
if 'temp_prop' not in array.properties:
array.add_property('temp_prop')
def _make_all_arrays_have_same_props(self, particle_arrays):
"""Make sure all arrays have the same props.
"""
all_props = reduce(set.union,
[set(x.properties.keys()) for x in particle_arrays])
for array in particle_arrays:
all_props.update(array.properties.keys())
for array in particle_arrays:
array_props = set(array.properties.keys())
for prop in (all_props - array_props):
array.add_property(prop)
class Interpolator(object):
"""Convenient class to interpolate particle properties onto a uniform
grid. This is particularly handy for visualization.
"""
def __init__(self, particle_arrays, num_points=125000, kernel=None,
x=None, y=None, z=None):
"""
Parameters
----------
particle_arrays: A list of particle arrays.
num_points: the number of points to interpolate on to.
kernel: the kernel to use for interpolation.
x: ndarray: the x-coordinate of points on which to interpolate.
y: ndarray: the y-coordinate of points on which to interpolate.
z: ndarray: the z-coordinate of points on which to interpolate.
The x, y, z coordinates need not be specified, and if they are not,
the bounds of the interpolated domain is automatically computed and
`num_points` number of points are used in this domain uniformly placed.
"""
self._set_particle_arrays(particle_arrays)
bounds = get_bounding_box(self.particle_arrays)
shape = get_nx_ny_nz(num_points, bounds)
self.dim = 3 - list(shape).count(1)
if kernel is None:
self.kernel = CubicSpline(dim=self.dim)
else:
self.kernel = kernel
self.pa = None
self.nnps = None
self.func_eval = None
if x is None and y is None and z is None:
self.set_domain(bounds, shape)
else:
self.set_interpolation_points(x=x, y=y, z=z)
#### Interpolator protocol ################################################
def set_interpolation_points(self, x=None, y=None, z=None):
"""Set the points on which we must interpolate the arrays.
Parameters
-----------
x: ndarray: the x-coordinate of points on which to interpolate.
y: ndarray: the y-coordinate of points on which to interpolate.
z: ndarray: the z-coordinate of points on which to interpolate.
If any of x, y, z is not passed it is assumed to be 0.0 and shaped
like the other non-None arrays.
"""
tmp = None
for tmp in (x, y, z):
if tmp is not None:
break
if tmp is None:
raise RuntimeError('At least one non-None array must be given.')
def _get_array(_t):
return np.asarray(_t) if _t is not None else np.zeros_like(tmp)
x, y, z = _get_array(x), _get_array(y), _get_array(z)
self.shape = x.shape
self.pa = self._create_particle_array(x, y, z)
arrays = self.particle_arrays + [self.pa]
if self.func_eval is None:
self._compile_acceleration_eval(arrays)
self.update_particle_arrays(self.particle_arrays)
def set_domain(self, bounds, shape):
"""Set the domain to interpolate into.
Parameters:
-----------
bounds: (xmin, xmax, ymin, ymax, zmin, zmax)
shape: (nx, ny, nz)
"""
self.bounds = np.asarray(bounds)
self.shape = np.asarray(shape)
x, y, z = self._create_default_points(self.bounds, self.shape)
self.set_interpolation_points(x, y, z)
def interpolate(self, prop, gradient=False):
"""
:prop: The name of the property to interpolate.
:gradient: bool: Evaluate gradient and not function.
:return: A numpy array suitably shaped with the property
interpolated.
"""
for array in self.particle_arrays:
data = array.get(prop, only_real_particles=False)
array.get('temp_prop', only_real_particles=False)[:] = data
self.func_eval.compute(0.0, 0.1) # These are junk arguments.
result = self.pa.prop.copy()
result.shape = self.shape
return result.squeeze()
def update_particle_arrays(self, particle_arrays):
"""Call this for a new set of particle arrays which have the
same properties as before.
For example, if you are reading the particle array data from files,
each time you load a new file a new particle array is read with the
same properties. Call this function to reset the arrays.
"""
self._set_particle_arrays(particle_arrays)
arrays = self.particle_arrays + [self.pa]
self._create_nnps(arrays)
self.func_eval.update_particle_arrays(arrays)
#### Private protocol #####################################################
def _create_nnps(self, arrays):
# create the neighbor locator object
self.nnps = NNPS(dim=self.kernel.dim, particles=arrays,
radius_scale=self.kernel.radius_scale)
self.nnps.update()
self.func_eval.set_nnps(self.nnps)
def _create_default_points(self, bounds, shape):
b = bounds
n = shape
x, y, z = np.mgrid[b[0]:b[1]:n[0]*1j,
b[2]:b[3]:n[1]*1j,
b[4]:b[5]:n[2]*1j,
]
return x, y, z
def _create_particle_array(self, x, y, z):
xr = x.ravel()
yr = y.ravel()
zr = z.ravel()
self.x, self.y, self.z = x.squeeze(), y.squeeze(), z.squeeze()
hmax = self._get_max_h_in_arrays()
h = hmax*np.ones_like(xr)
prop = np.zeros_like(xr)
pa = get_particle_array(
name='interpolate',
x=xr, y=yr, z=zr, h=h,
number_density=np.zeros_like(xr),
prop=prop,
grad_x=np.zeros_like(xr),
grad_y=np.zeros_like(xr),
grad_z=np.zeros_like(xr)
)
return pa
def _compile_acceleration_eval(self, arrays):
names = [x.name for x in self.particle_arrays]
equations = [InterpolateFunction(dest='interpolate', sources=names)]
self.func_eval = AccelerationEval(arrays, equations, self.kernel)
compiler = SPHCompiler(self.func_eval, None)
compiler.compile()
def _get_max_h_in_arrays(self):
hmax = -1.0
for array in self.particle_arrays:
hmax = max(array.h.max(), hmax)
return hmax
def _set_particle_arrays(self, particle_arrays):
self.particle_arrays = particle_arrays
self._make_all_arrays_have_same_props(particle_arrays)
for array in self.particle_arrays:
if 'temp_prop' not in array.properties:
array.add_property('temp_prop')
def _make_all_arrays_have_same_props(self, particle_arrays):
"""Make sure all arrays have the same props.
"""
all_props = reduce(
set.union, [set(x.properties.keys()) for x in particle_arrays]
)
for array in particle_arrays:
all_props.update(array.properties.keys())
for array in particle_arrays:
array_props = set(array.properties.keys())
for prop in (all_props - array_props):
array.add_property(prop)
