Reference: Code Generation

Kernel creation

class pystella.ElementWiseMap(map_instructions, **kwargs)[source]

An interface to loopy.make_kernel(), which creates a kernel with parallelization suitable for operations which are “local”—namely, element-wise maps where each workitem (thread) only accesses one element of global arrays.

__init__(map_instructions, **kwargs)[source]
Parameters

map_instructions – A list of instructions which write to global arrays. Entries may be loopy.Assignment’s or tuples (assignee, expression) of pymbolic expressions, the latter of which can include Field’s. All entries will be processed with index_fields().

The following keyword-only arguments are recognized:

Parameters

  • tmp_instructions – A list of instructions which write to temporary variables (i.e., local or private memory). Entries may be loopy.Assignment’s or tuples (assignee, expression) of pymbolic expressions, the latter of which can include Field’s. The expressions will be processed with index_fields(). The statements produced from tmp_instructions will precede those of map_instructions, and loopy.TemporaryVariable arguments will be inferred as needed.

  • args – A list of loopy.KernelArgument’s to be specified to loopy.make_kernel(). By default, all arguments (and their shapes) are inferred using get_field_args(), while any remaining (i.e., non-Field) arguments are inferred by loopy.make_kernel(). Any arguments passed via args override those inferred by either of the above options.

  • dtype – The default datatype of arrays to assume. Will only be applied to all loopy.KernelArgument’s whose datatypes were not already specified by any input args. Defaults to None.

  • lsize – The size of local parallel workgroups. Defaults to (16, 4, 1), which should come close to saturating memory bandwidth in many cases.

  • rank_shape – A 3-tuple specifying the shape of looped-over arrays. Defaults to None, in which case these values are not fixed (and will be inferred when the kernel is called at a slight performance penalty).

  • halo_shape – The number of halo layers on (both sides of) each axis of the computational grid. May either be an int, interpreted as a value to fix the parameter h to, or a tuple, interpreted as values for hx, hy, and hz. Defaults to None, in which case no such values are fixed at kernel creation.

Any remaining keyword arguments are passed to loopy.make_kernel().

Changed in version 2020.1: Arguments map_instructions and tmp_instructions replaced map_dict and tmp_dict and are allowed to be list’s with entries of loopy.Assignment’s and/or tuples.

__call__(queue=None, filter_args=False, **kwargs)[source]

Invokes the kernel, knl. All data arguments required by knl must be passed by keyword.

The following keyword arguments are recognized:

Parameters

In addition, any arguments recognized by loopy.PyOpenCLKernelExecutor.__call__() are also accepted.

Returns

(evt, output) where evt is the pyopencl.Event associated with the kernel invocation and output is any kernel output. See loopy’s tutorial for details.

knl

The generated loopy.LoopKernel.

class pystella.Stencil(map_instructions, halo_shape, **kwargs)[source]

A subclass of ElementWiseMap, which creates a kernel with parallelization suitable for stencil-type operations which are “non-local”—namely, computations which combine multiple neighboring values from a global array into a single output (per workitem/thread).

__init__(map_instructions, halo_shape, **kwargs)[source]

In addition to the parameters to ElementWiseMap.__init__(), the following arguments are required:

Parameters

halo_shape – The number of halo layers on (both sides of) each axis of the computational grid. May either be an int, interpreted as a value to fix the parameter h to, or a tuple, interpreted as values for hx, hy, and hz. Defaults to None, in which case no such values are fixed at kernel creation.

The following keyword-only arguments are recognized:

Parameters

prefetch_args – A list of arrays (namely, their name as a string) which should be prefetched into local memory. Defaults to an empty list.

class pystella.Reduction(decomp, input, **kwargs)[source]

A subclass of ElementWiseMap which computes (an arbitrary number of) reductions.

__init__(decomp, input, **kwargs)[source]
Parameters

  • decomp – An instance of DomainDecomposition.

  • input –

    May be one of the following:

    • a dict. The output of __call__() will be a dictionary with the same keys whose values are the corresponding reductions of the dict’s values. The values may either be lists of pymbolic expressions or a lists of tuple’s (expr, op), where expr is a pymbolic expression and op is the reduction operation to perform. Valid options are 'avg' (default), 'sum', 'prod', 'max', and 'min'.

    • a Sector. In this case, the reduction dictionary will be obtained from Sector.reducers.

    • a list of Sector’s. In this case, the input obtained from each Sector (as described above) will be combined.

The following keyword-only arguments are recognized (in addition to those accepted by ElementWiseMap):

Parameters

  • grid_size – The total number of gridpoints on the entire computational grid. Defaults to None, in which case it will be inferred at __call__() (if averages are being performed).

  • callback – A callable used to process the reduction results before __call__() returns. Defaults to lambda x: x, i.e., doing nothing.

__call__(queue=None, filter_args=False, **kwargs)[source]

Computes reductions by calling knl and DomainDecomposition.allreduce().

In addition to the arguments required by the actual kernel (passed by keyword only), the following keyword arguments are recognized:

Parameters

Any remaining keyword arguments are passed to knl.

Returns

A dict with the same keys as (interpreted from) input whose values are the corresponding (lists of) reduced values. Averages are obtained by dividing by grid_size. If grid_size was not supplied at __init__(), it is inferred (at a slight performance penalty).

class pystella.Histogrammer(decomp, histograms, num_bins, rank_shape, dtype, **kwargs)[source]

A subclass of ElementWiseMap which computes (an arbitrary number of) histograms.

__init__(decomp, histograms, num_bins, rank_shape, dtype, **kwargs)[source]
Parameters

  • decomp – An instance of DomainDecomposition.

  • histograms –

    A dict with values of the form (bin_expr, weight_expr), which are pymbolic expressions whose result determines the bin number and the associated weight contributed to that bin’s count, respectively. The output of __call__() will be a dictionary with the same keys whose values are the specified histograms.

    Note

    The values computed by bin_expr will by default be cast to integers by truncation. To instead round to the nearest integer, wrap the expression in a call to round.

  • num_bins – The number of bins of the computed histograms.

  • rank_shape – A 3-tuple specifying the shape of the computational sub-grid on the calling process.

  • dtype – The datatype of the resulting histogram.

In addition, any keyword-only arguments accepted by ElementWiseMap are also recognized.

__call__(queue=None, filter_args=False, **kwargs)[source]

Computes histograms by calling knl and DomainDecomposition.allreduce().

In addition to the arguments required by the actual kernel (passed by keyword only), the following keyword arguments are recognized:

Parameters

Any remaining keyword arguments are passed to knl.

Returns

A dict with the same keys as the input whose values are the corresponding histograms.

New in version 2020.1.

Fields

class pystella.Field(child, offset=0, shape=(), indices=('i', 'j', 'k'), ignore_prepends=False, base_offset=None)[source]

A pymbolic.primitives.Expression designed to mimic an array by carrying information about indexing. Kernel generators (Reduction, ElementWiseMap, and subclasses) automatically append indexing specified by the attributes indices and offset (via index_tuple) by preprocessing expressions with index_fields().

Examples:

>>> f = Field('f', offset='h')
>>> print(index_fields(f))
f[i + h, j + h, k + h]
>>> print(index_fields(f[0]))
f[0, i + h, j + h, k + h]

See test_field.py for more examples of the intended functionality.

child

The child expression representing the unsubscripted field. May be a string, a pymbolic.primitives.Variable, or a pymbolic.primitives.Subscript.

offset

The amount of padding by which to offset the array axes corresponding to the elements of indices. May be a tuple with the same length as indices or a single value. In the latter case, the input is transformed into a tuple with the same length as indices, each with the same value. Defaults to 0.

shape

The shape of axes preceding those indexed by indices. For example, Field('f', shape=(3, 'n')) would correspond to an array with shape (3, n, Nx, Ny, Nz) (using (Nx, Ny, Nz) as the shape along the final three axes indexed with indices). Used by get_field_args(). Defaults to an empty tuple.

indices

A tuple of (symbolic) array indices that will subscript the array. Each entry may be a pymbolic.primitives.Variable or a string which parses to one. Defaults to ('i', 'j', 'k')

ignore_prepends

Whether to ignore array subscripts prepended when processed with index_fields(). Useful for timestepping kernels (e.g., RungeKuttaStepper) which prepend array indices corresponding to extra storage axes (to specify that an array does not have this axis). Defaults to False.

base_offset

The amount of padding by which to offset the array axes corresponding to the elements of indices. In contrast to offset, denotes the offset of an “unshifted” array access, so that this attribute is used in determining the fully-padded shape of the underlying array, while use of shift_fields() may specify offset array accesses by modifying offset.

index_tuple

The fully-expanded subscript (i.e., indices offset by offset.)

Changed in version 2020.1: Added shape.

class pystella.DynamicField(child, offset='0', shape=(), indices=('i', 'j', 'k'), base_offset=None, dot=None, lap=None, pd=None)[source]

A subclass of Field which also contains associated Field instances representing various derivatives of the base Field.

dot

A Field representing the time derivative of the base Field. Defaults to a Field with name d{self.child}dt with the same shape, indices, and offset, but may be specified via the argument dot.

lap

A Field representing the Laplacian of the base Field. Defaults to a Field with name lap_{self.child} with the same shape and indices but with zero offset, but may be specified via the argument lap.

pd

A Field representing the spatial derivative(s) of the base Field. Defaults to a Field with name d{self.child}dx with shape (3,)+shape, the same indices, and zero offset, but may be specified via the argument pd.

d(*args)[source]

Returns the (subscripted) derivative of the base Field, i.e., either dot or pd with the appropriate index.

For example, the “time” derivative of a field would be

>>> f = DynamicField('f')
>>> print(f.d(0))  # x^0 = "time"
dfdt

Additional arguments are interpreted as subscripts to the resulting array; the final argument corresponds to the coordinate being differentiated with respect to.

>>> print(f.d(1, 2, 0))
dfdt[1, 2]

Spatial indices 1 through 3 denote spatial derivatives (whose array subscripts are 0 through 2).

>>> print(f.d(2))  # x^2 = y
dfdx[1]
>>> print(f.d(0, 1, 3))  # x^3 = z
dfdx[0, 1, 2]

Changed in version 2020.1: Specifying the names of dot, lap, and pd was replaced by passing actual Field instances.

pystella.index_fields(expr, prepend_with=None)[source]

Appends subscripts to Field instances in an expression, turning them into ordinary pymbolic.primitives.Subscript’s. See the documentation of Field for examples.

Parameters

  • expr – The expression(s) to be mapped.

  • prepend_with – A tuple of indices to prepend to the subscript of any Field’s in expr (unless a given Field has ignore_prepends set to False. Passed by keyword. Defaults to an empty tuple.

New in version 2020.1: Replaced Indexer().

pystella.shift_fields(expr, shift)[source]

Returns an expression with all Field’s shifted by shift–i.e., with shift added elementwise to each Field’s offset attribute.

Parameters

  • expr – The expression(s) to be mapped.

  • shift – A tuple.

New in version 2020.1.

pystella.diff(f, *x, allowed_nonsmoothness='discontinuous')[source]

A differentiator which computes \(\partial f / \partial x\) and understands Field’s. If x is one of t, x, y, or z and f is a DynamicField, the corresponding derivative Field is returned.

Examples:

>>> f = DynamicField('f')
>>> print(diff(f**3, f))
3*f**2
>>> print(diff(f**3, f, f))
3*2*f
>>> print(diff(f**3, 't'))
3*f**2*dfdt
>>> print(diff(f**3, f, 't'))
3*2*f*dfdt
>>> print(diff(f + 2, 'x'))
dfdx[0]
Parameters

  • f – A pymbolic expression to be differentiated.

  • x – A pymbolic.primitives.Expression or a string to be parsed (or multiple thereof). If multiple positional arguments are provided, derivatives are taken with respect to each in order. (See the examples above.)

pystella.get_field_args(expressions, unpadded_shape=None, prepend_with=None)[source]

Collects all Field’s from expressions and returns a corresponding list of loopy.ArrayArg’s, using their offset and shape attributes to determine their full shape.

Parameters

expressions – The expressions from which to collect Field’s.

The following keyword arguments are recognized:

Parameters

  • unpadded_shape – The shape of Field’s in expressions (sans padding). Defaults to (Nx, Ny, Nz).

  • prepend_with – A tuple to prepend to the shape of any Field’s in expressions (unless a given Field has ignore_prepends set to False. Passed by keyword. Defaults to an empty tuple.

Returns

A list of loopy.ArrayArg’s.

Example:

>>> f, g = Field('f', offset='h'), Field('g', shape=(3, 'a'), offset=1)
>>> get_field_args({f: g + 1})

would return the equivalent of:

>>> [lp.GlobalArg('f', shape='(Nx+2*h, Ny+2*h, Nz+2*h)', offset=lp.auto),
...  lp.GlobalArg('g', shape='(3, a, Nx+2, Ny+2, Nz+2)', offset=lp.auto)]

Changed in version 2020.1: Uses Field.shape to determine the full array shape.

Sympy interoperability

pystella.field.sympy.pymbolic_to_sympy(expr, *args, **kwargs)

A mapper which converts pymbolic.primitives.Expression’s into sympy expressions and understands Field’s. The result can be converted back to a pymbolic.primitives.Expression with all Field’s in place, accomplished via a subclass of sympy.Symbol which retains a copy of the Field.

Parameters

expr – The pymbolic expression to be mapped.

Warning

Currently, Field’s of the form Field('f[0]') will not be processed correctly.

pystella.field.sympy.sympy_to_pymbolic(expr, *args, **kwargs)

A mapper which converts sympy expressions into pymbolic.primitives.Expression’s and understands the custom sympy type used to represent Field’s by pymbolic_to_sympy().

Parameters

expr – The sympy expression to be mapped.

Warning

Currently, any modifications to the indices of a SympyField will not be reflected when mapped back to a Field. Use pymbolic.primitives.Subscript instead (i.e., process Field’s with index_fields() first).

pystella.field.sympy.simplify(expr, sympy_out=False)[source]

A wrapper to sympy.simplify().

Parameters

expr – The expression to be simplified. May either be a pymbolic.primitives.Expression or a sympy expression.

The following keyword arguments are recognized:

Parameters

sympy_out – A bool determining whether to return the simplified sympy expression or to first convert it to a pymbolic.primitives.Expression. Defaults to False.

Returns

A pymbolic.primitives.Expression containing the simplified form of expr if sympy_out is True, else a sympy expression.

Sectors

class pystella.Sector[source]

A unimplemented base class defining the methods and properties needed for code generation for, e.g., preheating simulations.

__init__()[source]

Processes input needed to specify a model for the particular Sector.

rhs_dict

An @property method returning a dict specifying the system of equations to be time-integrated. See the documentation of Stepper.

reducers

An @property method returning dict specifying the quantities to be computed, e.g., energy components for Expansion and output. See the documentation of Reduction.

stress_tensor(mu, nu, drop_trace=True)[source]
Parameters

drop_trace – Whether to drop the term \(g_{\mu\nu} \mathcal{L}\). Defauls to False.

Returns

The component \(T_{\mu\nu}\) of the stress-energy tensor of the particular Sector. Used by TensorPerturbationSector, with drop_trace=True.

class pystella.ScalarSector(nscalars, **kwargs)[source]

A Sector of scalar fields.

__init__(nscalars, **kwargs)[source]
Parameters

nscalars – The total number of scalar fields.

The following keyword-only arguments are recognized:

Parameters

  • f – The DynamicField of scalar fields. Defaults to DynamicField('f', offset='h', shape=(nscalars,)).

  • potential – A callable which takes as input a pymbolic expression or a list thereof, returning the potential of the scalar fields. Defaults to lambda x: 0.

Raises

ValueError – if a particular field is coupled to its own kinetic term.

class pystella.TensorPerturbationSector(sectors, **kwargs)[source]

A Sector of tensor perturbations.

__init__(sectors, **kwargs)[source]
Parameters

sectors – The Sector’s whose stress_tensor()’s source the tensor perturbations.

The following keyword-only arguments are recognized:

Parameters

hij – The DynamicField of tensor fields. Defaults to DynamicField('hij', offset='h', shape=(6,)).

pystella.sectors.get_rho_and_p(energy)[source]

Convenience callback for energy reductions which computes \(\rho\) and \(P\).

Parameters

energy – A dictionary of energy components as returned by Reduction.