pynbody.array

pynbody.array#

Defines arrays for simulation snapshots.

There are two distinct types of array used in pynbody.

  • SimArray defines a subclass of numpy.ndarray for extra functionality like unit tracking

  • IndexedSimArray defines a distinct type of array, which does not subclass numpy.ndarray but behaves like one. It points to a parent array and a set of indices, and can be used to represent a subset of elements within an array.

For most purposes, the differences between numpy.ndarray and SimArray are not important. However, when units are specified (by setting the SimArray.units attribute), the behaviour is slightly different. In particular,

  • it becomes impossible to add or subtract arrays with incompatible dimensions

    >>> SimArray([1,2], "Mpc") + SimArray([1,2], "Msol"))
ValueError

  • addition or subtraction causes auto-unit conversion. For example

    >>> SimArray([1,2], "Mpc") + SimArray([1,2], "kpc")
SimArray([1.001, 1.002], "Mpc")

  • Note that in this context the left value takes precedence in specifying the return units, so that reversing the order of the operation here would return results in kpc.

  • If only one of the arrays specifies a Unit, no checking occurs and the unit of the returned array is assumed to be the same as the one specified input unit.

  • Powers to single integer or rational powers will maintain unit tracking. Powers to float or other powers will not be able to do so.

    >>> SimArray([1,2],"Msol Mpc**-3")**2
SimArray([1, 4], 'Msol**2 Mpc**-6')
>>> SimArray([1,2],"Msol Mpc**-3")**(1,3)
SimArray([ 1.,1.26], 'Msol**1/3 Mpc**-1')

  • The syntax is similar to that used by pynbody.units.Unit, where a length-two tuple can represent a rational number, in this case one third. If a floating point number is used, the unit tracking is lost:

    >>> SimArray([1.,2], "Msol Mpc**-3")**0.333
SimArray([ 1.,1.26])  # Lost track of units

Getting the array in specified units#

Given an array, you can convert it in-place into units of your own chosing:

>>> x = SimArray([1,2], "Msol")
>>> x.convert_units('kg')
>>> print(x)
SimArray([  1.99e+30,   3.98e+30], 'kg')

Or you can leave the original array alone and get a copy in different units, correctly converted:

>>> x = SimArray([1,2], "Msol")
>>> print(x.in_units("kg"))
SimArray([  1.99e+30,   3.98e+30], 'kg')
>>> print(x)
SimArray([1,2], "Msol")

You can also supply context values for conversions (e.g. the cosmological scale factor a):

>>> x = SimArray([1,2], "kpc a")
>>> print(x.in_units("kpc", a=0.5))
SimArray([  0.5,   1. ], 'kpc')

Per-element (array-valued) context is supported and broadcasts over the leading axis:

>>> import numpy as np
>>> SimArray([1, 2, 3], "a kpc").in_units("kpc", a=np.array([0.2, 0.3, 0.4]))
SimArray([0.2, 0.6, 1.2], 'kpc')

If the SimArray (or IndexedSimArray) was created by a SimSnap (which is most likely), it has a pointer into the SimSnap’s properties so that the cosmological context is automatically fetched. For example, comoving -> physical conversions are correctly achieved:

>>> f = pynbody.load("fname")
>>> f['pos']
SimArray([[ 0.05419805, -0.0646539 , -0.15700017],
         [ 0.05169899, -0.06193341, -0.14475258],
         [ 0.05672406, -0.06384531, -0.15909944],
         ...,
         [ 0.0723075 , -0.07650762, -0.07657281],
         [ 0.07166634, -0.07453796, -0.08020873],
         [ 0.07165282, -0.07468577, -0.08020587]], '2.86e+04 kpc a')

>>> f['pos'].convert_units('kpc')
>>> f['pos']
SimArray([[ 1548.51403101, -1847.2525312 , -4485.71463308],
         [ 1477.1124212 , -1769.52421398, -4135.78377699],
         [ 1620.68592366, -1824.15000686, -4545.69387564],
         ...,
         [ 2065.9264273 , -2185.92982874, -2187.79225915],
         [ 2047.60759667, -2129.6537339 , -2291.6758134 ],
         [ 2047.2214441 , -2133.87693163, -2291.59406997]], 'kpc')

Specifying rules for ufuncs#

In general, it’s not possible to infer what the output units from a given ufunc should be. While numpy built-in ufuncs should be handled OK, other ufuncs will need their output units defined (otherwise a numpy.ndarray will be returned instead of our custom type.)

To do this, decorate a function with SimArray.ufunc_rule(). The function you define should take the same number of parameters as the ufunc. These will be the input parameters of the ufunc. You should return the correct units for the output, or raise units.UnitsException (in the latter case, the return array will be made into a numpy.ndarray.)

For example, here is the code for the standard numpy sqrt. (You don’t need this because it’s already built in, but it shows you how to do it.)

@SimArray.ufunc_rule(np.sqrt)
def sqrt_units(a):
    if a.units is not None:
        return a.units ** (1, 2)
    else:
        return None

Shared memory arrays#

The array package also provides mechanisms for creating arrays that back onto shared memory. These can be easily created using array_factory() with shared=True. For implementation details see pynbody.array.shared.

Functions

array_factory(dims, dtype, zeros, shared)

Create an array of dimensions dims with the given numpy dtype.

Classes

IndexedSimArray(array, ptr)

A view into a SimArray that allows for indexing and slicing.

SimArray(data[, units, sim])

A shallow wrapper around numpy.ndarray for extra functionality like unit-tracking.

Modules

shared

Support for numpy arrays in shared memory.
Contents