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.

Parameters:

map_instructions – A list of instructions which write to global arrays. Entries may be loopy.Assignments or tuples (assignee, expression) of pymbolic expressions, the latter of which can include Fields. 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.Assignments or tuples (assignee, expression) of pymbolic expressions, the latter of which can include Fields. 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.KernelArguments 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.KernelArguments 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 lists with entries of loopy.Assignments 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).

In addition to the parameters to ElementWiseMap(), 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.

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 tuples (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 Sectors. 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 initialization, it is inferred (at a slight performance penalty).

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

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

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.

  • dtype – The datatype of the resulting histogram.

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

New in version 2020.1.

Changed in version 2020.2: Positional argument rank_shape no longer required.

__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.

Fields#

class pystella.Field(child, offset=0, shape=(), indices=('i', 'j', 'k'), ignore_prepends=False, base_offset=None, dtype=None, dim_tags=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 = ps.Field("f", offset="h")
>>> print(ps.index_fields(f))
f[i + h, j + h, k + h]
>>> print(ps.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.

dtype#

The datatype of the field. Defaults to None, in which case datatypes are inferred by loopy at kernel invocation.

index_tuple#

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

Changed in version 2020.2: Added dtype.

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, dtype=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 shape+(3,), 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 = ps.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.Subscripts. 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 Fields 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 Fields 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 Fields. If x is one of t, x, y, or z and f is a DynamicField, the corresponding derivative Field is returned.

Examples:

>>> f = ps.DynamicField("f")
>>> print(ps.diff(f**3, f))
3*f**2
>>> print(ps.diff(f**3, f, f))
3*2*f
>>> print(ps.diff(f**3, "t"))
3*f**2*dfdt
>>> print(ps.diff(f**3, f, "t"))
3*2*f*dfdt
>>> print(ps.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 Fields from expressions and returns a corresponding list of loopy.ArrayArgs, using their offset and shape attributes to determine their full shape.

Parameters:

expressions – The expressions from which to collect Fields.

The following keyword arguments are recognized:

Parameters:

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

  • prepend_with – A tuple to prepend to the shape of any Fields 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.ArrayArgs.

Example:

>>> f = ps.Field("f", offset="h")
>>> g = ps.Field("g", shape=(3, "a"), offset=1)
>>> ps.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.Expressions into sympy expressions and understands Fields. The result can be converted back to a pymbolic.primitives.Expression with all Fields 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, Fields 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.Expressions and understands the custom sympy type used to represent Fields 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 Fields 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.

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.

Parameters:

sectors – The Sectors 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.