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 beloopy.Assignment
s ortuple
s(assignee, expression)
ofpymbolic
expressions, the latter of which can includeField
s. All entries will be processed withindex_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 beloopy.Assignment
s ortuple
s(assignee, expression)
ofpymbolic
expressions, the latter of which can includeField
s. The expressions will be processed withindex_fields()
. The statements produced fromtmp_instructions
will precede those ofmap_instructions
, andloopy.TemporaryVariable
arguments will be inferred as needed.args – A list of
loopy.KernelArgument
s to be specified toloopy.make_kernel()
. By default, all arguments (and their shapes) are inferred usingget_field_args()
, while any remaining (i.e., non-Field
) arguments are inferred byloopy.make_kernel()
. Any arguments passed viaargs
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 inputargs
. 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 parameterh
to, or atuple
, interpreted as values forhx
,hy
, andhz
. 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
andtmp_instructions
replacedmap_dict
andtmp_dict
and are allowed to belist
s with entries ofloopy.Assignment
s and/ortuple
s.- __call__(queue=None, filter_args=False, **kwargs)[source]#
Invokes the kernel,
knl
. All data arguments required byknl
must be passed by keyword.The following keyword arguments are recognized:
- Parameters:
queue –
The
pyopencl.CommandQueue
on which to enqueue the kernel. If None (the default),queue
is not passed (i.e., forloopy.ExecutableCTarget
).Note
For
loopy.PyOpenCLTarget
(the default), a validpyopencl.CommandQueue
is a required argument.filter_args – Whether to filter
kwargs
such that only arguments to theknl
are passed. Defaults to False.
In addition, any arguments recognized by
loopy.PyOpenCLKernelExecutor.__call__()
are also accepted.- Returns:
(evt, output)
whereevt
is thepyopencl.Event
associated with the kernel invocation andoutput
is any kernel output. Seeloopy
’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 parameterh
to, or atuple
, interpreted as values forhx
,hy
, andhz
. 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 thedict
’s values. The values may either be lists ofpymbolic
expressions or a lists oftuple
s(expr, op)
, whereexpr
is apymbolic
expression andop
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 fromSector.reducers
.a
list
ofSector
s. In this case, the input obtained from eachSector
(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 tolambda x: x
, i.e., doing nothing.
- __call__(queue=None, filter_args=False, **kwargs)[source]#
Computes reductions by calling
knl
andDomainDecomposition.allreduce()
.In addition to the arguments required by the actual kernel (passed by keyword only), the following keyword arguments are recognized:
- Parameters:
queue – The
pyopencl.CommandQueue
on which to enqueue the kernel. Defaults to None, in which casequeue
is not passed (i.e., forloopy.ExecutableCTarget
)filter_args – Whether to filter
kwargs
such that no unexpected arguments are passed to theknl
. Defaults to False.allocator – A
pyopencl
allocator used to allocate temporary arrays, i.e., most usefully apyopencl.tools.MemoryPool
. See the note in the documentation ofSpectralCollocator()
.
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 bygrid_size
. Ifgrid_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 arepymbolic
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 toround
.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
andDomainDecomposition.allreduce()
.In addition to the arguments required by the actual kernel (passed by keyword only), the following keyword arguments are recognized:
- Parameters:
queue – The
pyopencl.CommandQueue
on which to enqueue the kernel. Defaults to None, in which casequeue
is not passed (i.e., forloopy.ExecutableCTarget
)filter_args – Whether to filter
kwargs
such that no unexpected arguments are passed to theknl
. Defaults to False.allocator – A
pyopencl
allocator used to allocate temporary arrays, i.e., most usefully apyopencl.tools.MemoryPool
. See the note in the documentation ofSpectralCollocator()
.
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 attributesindices
andoffset
(viaindex_tuple
) by preprocessing expressions withindex_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 apymbolic.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 asindices
or a single value. In the latter case, the input is transformed into a tuple with the same length asindices
, each with the same value. Defaults to0
.
- 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 withindices
). Used byget_field_args()
. Defaults to an emptytuple
.
- 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 tooffset
, 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 ofshift_fields()
may specify offset array accesses by modifyingoffset
.
- dtype#
The datatype of the field. Defaults to None, in which case datatypes are inferred by
loopy
at kernel invocation.
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 associatedField
instances representing various derivatives of the baseField
.- dot#
A
Field
representing the time derivative of the baseField
. Defaults to aField
with named{self.child}dt
with the sameshape
,indices
, andoffset
, but may be specified via the argumentdot
.
- lap#
A
Field
representing the Laplacian of the baseField
. Defaults to aField
with namelap_{self.child}
with the sameshape
andindices
but with zerooffset
, but may be specified via the argumentlap
.
- pd#
A
Field
representing the spatial derivative(s) of the baseField
. Defaults to aField
with named{self.child}dx
with shapeshape+(3,)
, the sameindices
, and zerooffset
, but may be specified via the argumentpd
.
- d(*args)[source]#
Returns the (subscripted) derivative of the base
Field
, i.e., eitherdot
orpd
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
through3
denote spatial derivatives (whose array subscripts are0
through2
).>>> print(f.d(2)) # x^2 = y dfdx[1] >>> print(f.d(0, 1, 3)) # x^3 = z dfdx[0, 1, 2]
- pystella.index_fields(expr, prepend_with=None)[source]#
Appends subscripts to
Field
instances in an expression, turning them into ordinarypymbolic.primitives.Subscript
s. See the documentation ofField
for examples.- Parameters:
New in version 2020.1: Replaced
Indexer()
.
- pystella.shift_fields(expr, shift)[source]#
Returns an expression with all
Field
s shifted byshift
–i.e., withshift
added elementwise to eachField
’soffset
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. Ifx
is one oft
,x
,y
, orz
andf
is aDynamicField
, the corresponding derivativeField
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
Field
s fromexpressions
and returns a corresponding list ofloopy.ArrayArg
s, using theiroffset
andshape
attributes to determine their full shape.- Parameters:
expressions – The expressions from which to collect
Field
s.
The following keyword arguments are recognized:
- Parameters:
- Returns:
A
list
ofloopy.ArrayArg
s.
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.Expression
s intosympy
expressions and understandsField
s. The result can be converted back to apymbolic.primitives.Expression
with allField
s in place, accomplished via a subclass ofsympy.Symbol
which retains a copy of theField
.- Parameters:
expr – The
pymbolic
expression to be mapped.
Warning
Currently,
Field
s of the formField("f[0]")
will not be processed correctly.
- pystella.field.sympy.sympy_to_pymbolic(expr, *args, **kwargs)#
A mapper which converts
sympy
expressions intopymbolic.primitives.Expression
s and understands the customsympy
type used to representField
s bypymbolic_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 aField
. Usepymbolic.primitives.Subscript
instead (i.e., processField
s withindex_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 asympy
expression.
The following keyword arguments are recognized:
- Parameters:
sympy_out – A
bool
determining whether to return the simplifiedsympy
expression or to first convert it to apymbolic.primitives.Expression
. Defaults to False.- Returns:
A
pymbolic.primitives.Expression
containing the simplified form ofexpr
ifsympy_out
is True, else asympy
expression.
Sectors#
- class pystella.Sector[source]#
A unimplemented base class defining the methods and properties needed for code generation for, e.g., preheating simulations.
- rhs_dict#
An
@property
method returning adict
specifying the system of equations to be time-integrated. See the documentation ofStepper
.
- reducers#
An
@property
method returningdict
specifying the quantities to be computed, e.g., energy components forExpansion
and output. See the documentation ofReduction
.
- 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 byTensorPerturbationSector
, withdrop_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 toDynamicField("f", offset="h", shape=(nscalars,))
.potential – A
Callable
which takes as input apymbolic
expression or alist
thereof, returning the potential of the scalar fields. Defaults tolambda 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
Sector
s whosestress_tensor()
s source the tensor perturbations.
The following keyword-only arguments are recognized:
- Parameters:
hij – The
DynamicField
of tensor fields. Defaults toDynamicField("hij", offset="h", shape=(6,))
.