Basics¶

This fast-paced section reviews a wide range of standard operations. It assumes the following import declarations

import drjit as dr
from drjit.auto import Float, Array3f, UInt

Creating arrays¶

Recall that Dr.Jit array types are dynamically sized–for example, Float refers to a 1D array of single precision variables.

The simplest way to create such an array is to call its constructor with a list of explicit values:

a = Float(1, 2, 3, 4)
print(a) # [1, 2, 3, 4]

The constructor also accepts Sequence types (e.g. lists, tuples, NumPy arrays, PyTorch tensors, etc.):

x = Float([1, 2, 3, 4])

Nested array types store several variables—for example, Array3f is just a wrapper around 3 Float instances. They can be passed to the constructor explicitly, or via implicit conversion from constants, lists, etc.

a = Array3f([1, 2], 0, Float(10, 20))
print(a)
# Prints (with 'y' component broadcast to full size)
# [[1, 0, 10],
#  [2, 0, 20]]

Various functions can also create default-initialized arrays:

dr.zeros(): [0, 0, ...].
dr.ones(): [1, 1, ...].
dr.full(): [x, x, ...] given x.
dr.arange(): [0, 1, 2, ...].
dr.linspace(): linear interpolation of two endpoints.
dr.empty(): allocate uninitialized memory.

These always take the desired output type as first argument. You can optionally request a given size along the dynamic axis, e.g.:

b = dr.zeros(Array3f)
print(b.shape) # Prints: (3, 1)

b = dr.zeros(Array3f, shape=(3, 1000))
print(b.shape) # Prints: (3, 1000)

Element access¶

Use the default array[index] syntax to read/write array entries. Nested static 1-4D arrays further expose equivalent .x / .y / .z / .w members:

a = Array3f(1, 2, 3)
a.x += a.z + a[1]

Static 1-4D arrays also support swizzling, which arbitrarily reorders elements:

a.xy = a.xx + a.yx

Arithmetic operations¶

Except for a few special cases (e.g., matrix multiplication), arithmetic operations transform arrays element-wise. If needed, the system will implicitly broadcast the operands and promote types.

>>> a = abs(Float(-1.25, 2) + UInt32(1))
>>> type(a)
<class 'drjit.cuda.Float'>
>>> a
[0.25, 3]

In the above example, broadcasting automatically extended the size of the scalar (size-1) array, and the UInt32 type was promoted to Float. Type promotion follows the rules of the C programming language.

Besides built-in arithmetic operators, the following standard functions are available:

dr.abs(): Absolute value.
dr.fma(): Fused multiply-add.
dr.minimum(), dr.maximum(): Element-wise minimum/maximum of two arrays.
dr.ceil(), dr.floor(), dr.round(), dr.trunc(): Round up, down, to nearest, or to zero.
dr.sqrt(), dr.cbrt(): Square and cube root.
dr.rcp(): Reciprocal.
dr.rsqrt(): Reciprocal square root.
dr.sign(): Extract the sign.
dr.copysign(): Copy sign from one value to another.
dr.clip(): Clip a value to an interval.
dr.lerp(): Linearly interpolate.

The library implements common transcendental functions:

dr.sin(), dr.cos(), dr.tan(): Trigonometric functions.
dr.asin(), dr.acos(), dr.atan(), dr.atan2(): .. and their inverses.
dr.sinh(), dr.cosh(), dr.tanh(): Hyperbolic trigonometric functions.
dr.asinh(), dr.acosh(), dr.atanh(): .. and their inverses.
dr.sincos(), dr.sincosh(): Fast combined evaluation.
dr.erf(), dr.erfinv(): Error function.
dr.exp(), dr.log(), dr.exp2(), dr.log2(): Exponentials and logarithms.
dr.power(): Power function.
dr.lgamma(): Gamma function.

Most of these support real and complex-valued inputs. A subset accepts quaternions (see the section on array types for details). Integer arrays further support the following bit-level operations

dr.lzcnt(), dr.tzcnt(): Leading/trailing zero count.
dr.popcnt(): Population count.
dr.brev(): Bit reverse.

Mask operations¶

Equality and inequality comparisons produce masks (i.e., boolean-valued arrays) with support for binary arithmetic. The dr.select() function blends results from two arrays based on a mask analogous to the ternary (”?”) operator in C/C++.

>>> a = dr.arange(Float, 5) - 3
>>> mask = (a < 0) | (a == 2)
>>> mask
[True, True, False, False, True]
>>> dr.select(mask, -1, a)        # select(mask, true_value, false_value)
[-1, -1, 0, 1, -1]

Masks can also be applied to arrays in order to zero out the False indices by using the & operator.

>>> a = Float([1, 2, 3])
>>> mask = Bool([True, False, True])
>>> a & mask
[1, 0, 3]

Reductions¶

Reductions use a given operation (e.g., addition) to combine values along one or several dimensions.

dr.sum(), dr.prod(): Sum and product reduction.
dr.min(), dr.max(): Minimum/maximum reduction.
dr.all(), dr.any(), dr.none(): Boolean reductions for mask arrays.
dr.reduce(): Generalized reduction operator.

By default, they reduce arrays along the leading array dimension. For example, the following reduction is equivalent to a.x + a.y + a.z. By reducing this value once more or specifying axis=None, we can sum over all entries.

>>> a = Array3f([1, 2], [10, 20], [100, 200])
>>> dr.sum(a)
[111, 222]
>>> dr.sum(a, axis=None)
[333]

Accessing memory: gather/scatter¶

The function dr.gather() fetches values from a 1D array with positions specified by an index array. For example:

>>> buf = Float(10, 20, 30, 40, 50, 60)
>>> index = UInt32(1, 0)
>>> dr.gather(Float, buf, index)
[20, 10]

Note how the operation takes the desired output type as first argument. We can also gather nested arrays (assumed to be flattened in the source 1D array using C-style order) by requesting a different result type.

>>> dr.gather(Array3f, buf, index)
[[40, 50, 60],
 [10, 20, 30]]

Whereas gather reads memory, dr.scatter() realizes the corresponding write operation.

>>> dr.scatter(buf, Array3f(0, 1, 2), UInt32(1))
>>> buf
[10, 20, 30, 0, 1, 2]

Finally, dr.scatter_add() (and the more general dr.scatter_reduce()) atomically accumulates values into an array.

>>> dr.scatter_add(buf, Array3f(100), UInt32(1))
>>> buf
[10, 20, 30, 100, 101, 102]

Random number generation¶

Dr.Jit was originally developed for Monte Carlo simulation, and programs in that domain require a source of (pseudo-) randomness. For this, the system provides a member of the PCG Family of random number generators by Melissa O’Neill. To try it, import the class drjit.*.PCG32 from your backend of choice and initialize with the desired output array size.

>>> from drjit.auto import PCG32
>>> rng = PCG32(10000)
>>> rng.next_float32()
[0.108379, 0.533909, 0.00684452, .. 9994 skipped .., 0.511698, 0.600626, 0.219648]
>>> rng.next_uint32_bounded(4)
[1, 0, 0, .. 9994 skipped .., 0, 3, 3]

Please see the documentation of this class for a review of its features.