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 beloopy.Assignment
’s or tuples(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 or tuples(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/or tuples.
-
__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).-
__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 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.-
__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 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__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 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.
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
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.
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 attributesindices
andoffset
(viaindex_tuple
) by preprocessing expressions withindex_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 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
.
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 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 shape(3,)+shape
, 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 = 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 = 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 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, 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 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.-
__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 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.-
__init__
(sectors, **kwargs)[source]¶ - 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,))
.
-