pynbody.kdtree.KDTree#

class pynbody.kdtree.KDTree(pos, mass, leafsize=32, boxsize=None, num_threads=None, shared_mem=False)[source]#

Bases: object

KDTree can be used for smoothing, interpolating and geometrical queries.

You are unlikely to need to interact with the KDTree directly, instead using the higher-level SPH functionality such as

>>> f = pynbody.load("my_snapshot")
>>> f.gas['smooth'] # will build KDTree and calculate smoothing lengths for gas, if not available on disk
>>> f.gas['rho'] # will calculate densities for gas, if not available on disk
>>> f.dm['rho'] # will calculate densities for dark matter, if not available on disk

KDTree can also be used to accelerate geometrical queries on a snapshot, e.g.:

>>> f = pynbody.load('my_snapshot')
>>> f.build_kdtree()
>>> f[pynbody.filt.Sphere('10 kpc')]

See performance of filters for more information.

Most KDTree operations proceed in parallel, using the number of threads specified in the pynbody configuration.

Performance statistics can be tested using the performance_kdtree.py script in the tests folder.

Note that the KDTree takes into account periodicity of cosmological volumes if it can. However, as soon as a snapshot has been rotated, this may become impossible. In this case, the KDTree will issue a warning, and will not use periodicity. This is to avoid incorrect results due to the periodicity assumption, but it will lead to artefacts on the edge of the box. If you are working with a cosmological simulation, it is best to perform any KDTree operations (such as smoothing) before rotating the snapshot, e.g.:

>>> f = pynbody.load('my_snapshot')
>>> f['rho']; f['smooth'] # force KDTree to be built and density estimations made
>>> pynbody.analysis.faceon(f) # rotate the snapshot
>>> pynbody.plot.image(f.g, qty='rho', width='10 kpc')

Methods

all_nn([nn])

A list of neighbours information for all the particles in the snapshot.

array_name_to_id(name)

Convert the pynbody name of an array to the corresponding integer ID in the C++ code.

deserialize(pos, mass, serialized_state[, ...])

Reconstruct a KDTree from a serialized state.

get_array_ref(name)

Get the current array reference for a given name ('smooth', 'rho', 'mass', 'qty', or 'qty_sm').

nn([nn])

Generator of neighbour list.

particles_in_sphere(center, radius)

Find particles within a sphere.

populate(mode, nn)

Create the KDTree and perform the operation specified by mode.

serialize()

Produce a serialized description of the tree

set_array_ref(name, ar)

Give the C++ code access to a given array.

set_kernel([kernel])

Set the kernel for smoothing operations.

smooth_operation_to_id(name)

Convert the name of a smoothing operation to the corresponding integer ID in the C++ code.

sph_curl(array[, nsmooth])

Calculate the curl of a given array using SPH smoothing.

sph_dispersion(array[, nsmooth])

Calculate the SPH dispersion of a simulation array.

sph_divergence(array[, nsmooth])

Calculate the divergence of a given array using SPH smoothing.

sph_mean(array[, nsmooth])

Calculate the SPH mean of a simulation array.

__init__(pos, mass, leafsize=32, boxsize=None, num_threads=None, shared_mem=False)[source]#

Create a KDTree

Most users will not need to call this directly, but will instead use the higher-level functions, either via accessing SPH properties on a snapshot, or calling pynbody.snapshot.simsnap.SimSnap.build_tree().

Parameters:
  • pos (pynbody.array.SimArray) – Particles positions.

  • mass (pynbody.array.SimArray) – Particles mass.

  • leafsize (int, optional) – The number of particles in leaf nodes (default 32).

  • boxsize (float, optional) – Boxsize (default None)

  • num_threads (int, optional) – Number of threads to use when building tree (if None, use configured/detected number of processors).

  • shared_mem (bool, optional) – Whether to keep kdtree in shared memory so that it can be shared between processes (default False).

all_nn(nn=None)[source]#

A list of neighbours information for all the particles in the snapshot.

static array_name_to_id(name)[source]#

Convert the pynbody name of an array to the corresponding integer ID in the C++ code.

classmethod deserialize(pos, mass, serialized_state, num_threads=None, boxsize=None)[source]#

Reconstruct a KDTree from a serialized state.

get_array_ref(name)[source]#

Get the current array reference for a given name (‘smooth’, ‘rho’, ‘mass’, ‘qty’, or ‘qty_sm’).

nn(nn=None)[source]#

Generator of neighbour list.

Parameters:

nn (int, None) – number of neighbours. If None, default to 64.

Yields:

nbr_list (list[int, float, list[int], list[float]]) – Information about the neighbours of a particle.

List with four elements:

  • nbr_list[0] (int): the index of the particle in the snapshot’s arrays.

  • nbr_list[1] (float): the smoothing length of the particle

  • nbr_list[2] (list[int]): list of the indexes of the nn neighbouring particles.

  • nbr_list[3] (list[float]): list of distances squared of each neighbouring particles.

The lists of particle and of the relative distance squared are not sorted.

particles_in_sphere(center, radius)[source]#

Find particles within a sphere.

Most users will not need to call this directly, but will instead use pynbody.filt.Sphere.

Parameters:
  • center (array_like) – Center of the sphere.

  • radius (float) – Radius of the sphere.

Returns:

indices – Indices of the particles within the sphere.

Return type:

array_like

populate(mode, nn)[source]#

Create the KDTree and perform the operation specified by mode.

Parameters:
  • mode (str (see kdtree.smooth_operation_to_id)) – Specify operation to perform (compute smoothing lengths, density, SPH mean, or SPH dispersion).

  • nn (int) – Number of neighbours to be considered when smoothing.

serialize()[source]#

Produce a serialized description of the tree

set_array_ref(name, ar)[source]#

Give the C++ code access to a given array.

Parameters:
  • name (str) – Name of the array to set (can be ‘smooth’, ‘rho’, ‘mass’, ‘qty’, ‘qty_sm’)

  • ar (pynbody.array.SimArray) – Array that the C++ code will access. Note that INCREF is called on the array, so it will be kept alive as long as the KDTree object is alive (or until another set_array_ref call replaces it.)

set_kernel(kernel='CubicSpline')[source]#

Set the kernel for smoothing operations.

Parameters:

kernel (str) – Keyword to specify the smoothing kernel. Options: ‘CubicSpline’, ‘WendlandC2’

smooth_operation_to_id(name)[source]#

Convert the name of a smoothing operation to the corresponding integer ID in the C++ code.

Valid names are:

  • ‘hsm’ : compute smoothing lengths

  • ‘rho’ : compute density

  • ‘qty_mean’ : compute SPH mean of the ‘qty’ array

  • ‘qty_disp’ : compute SPH dispersion of the ‘qty’ array

  • ‘qty_div’ : compute divergence of the ‘qty’ array

  • ‘qty_curl’ : compute curl of the ‘qty’ array

sph_curl(array, nsmooth=64)[source]#

Calculate the curl of a given array using SPH smoothing.

Parameters:
  • array (pynbody.array.SimArray) – Quantity to compute curl of.

  • nsmooth (int) – Number of neighbours to use when smoothing.

Returns:

The curl of the input array.

Return type:

pynbody.array.SimArray

sph_dispersion(array, nsmooth=64)[source]#

Calculate the SPH dispersion of a simulation array.

Given a kernel W, this routine computes the smoothed quantity:

\[r_i = \sum_{j=0}^{N_{smooth}}\left( \frac{m_j}{\rho_j} q_i W(|x_i - x_j|, s_i) \right)\]

Then the squared root of the SPH smoothed variance:

\[d_i = \left\{ \sum_{j=0}^{N_{smooth}} \frac{m_j}{\rho_j} (q_i - r_i)^2 W(|x_i - x_j|, s_i) \right\}^{1/2}\]

where d_i is the calculated dispersion, m is the mass array, rho is the density array, q is the input array, s_i is the smoothing length, and W is the kernel function.

Actual computation is done in the C++-extension functions smooth.cpp::smDispQty[1,N]D.

Parameters:
  • array (pynbody.array.SimArray) – Quantity to compute dispersion of.

  • nsmooth (int) – Number of neighbours to use when smoothing.

Returns:

output – The dispersion of the input array.

Return type:

pynbody.array.SimArray

sph_divergence(array, nsmooth=64)[source]#

Calculate the divergence of a given array using SPH smoothing.

Parameters:
  • array (pynbody.array.SimArray) – Quantity to compute curl of.

  • nsmooth (int) – Number of neighbours to use when smoothing.

Returns:

The curl of the input array.

Return type:

pynbody.array.SimArray

sph_mean(array, nsmooth=64)[source]#

Calculate the SPH mean of a simulation array.

Given a kernel W, this calculates

\[r_i = \sum_{j=0}^{N_{smooth}}\left( \frac{m_j}{\rho_j} q_i W(|x_i - x_j|, s_i) \right)\]

where r is the result of the smoothing, m is the mass array, rho is the density array, q is the input array, s_i is the smoothing length, and W is the kernel function.

The actual computation is done in the C++-extension function smooth.cpp::smMeanQty[1,N]D.

Parameters:
  • array (pynbody.array.SimArray) – Quantity to smooth (compute SPH interpolation at particles position).

  • nsmooth – Number of neighbours to use when smoothing.

Returns:

output – The SPH mean of the input array.

Return type:

pynbody.array.SimArray