Parameter studies

arraycontext/__init__.py

@@ -88,7 +88,8 @@
     pytest_generate_tests_for_pyopencl_array_context,
 )
 from .transform_metadata import CommonSubexpressionTag, ElementwiseMapKernelTag
-
+from .parameter_study import pack_for_parameter_study, unpack_parameter_study
+from .parameter_study.transform import ParamStudyPytatoPyOpenCLArrayContext
 
 __all__ = (
     "Array",

arraycontext/impl/pytato/__init__.py

@@ -59,12 +59,13 @@
 import numpy as np
 
 from pytools import memoize_method
-from pytools.tag import Tag, ToTagSetConvertible, normalize_tags
+from pytools.tag import Tag, ToTagSetConvertible, normalize_tags, UniqueTag
 
 from arraycontext.container.traversal import rec_map_array_container, with_array_context
 from arraycontext.context import Array, ArrayContext, ArrayOrContainer, ScalarLike
 from arraycontext.metadata import NameHint
 
+from dataclasses import dataclass
 
 if TYPE_CHECKING:
     import pyopencl as cl
@@ -562,7 +563,10 @@ def _to_frozen(key: Tuple[Any, ...], ary) -> TaggableCLArray:
         evt, out_dict = pt_prg(self.queue,
                 allocator=self.allocator,
                 **bound_arguments)
-        evt.wait()
+        if isinstance(evt, list):
+            [_evt.wait() for _evt in evt]
+        else:
+            evt.wait()
         assert len(set(out_dict) & set(key_to_frozen_subary)) == 0
 
         key_to_frozen_subary = {
@@ -699,6 +703,8 @@ def clone(self):
 # }}}
 
 
+
+
 # {{{ PytatoJAXArrayContext
 
 class PytatoJAXArrayContext(_BasePytatoArrayContext):

arraycontext/impl/pytato/compile.py

@@ -646,7 +646,10 @@ def __call__(self, arg_id_to_arg) -> ArrayContainer:
         # FIXME Kernels (for now) allocate tons of memory in temporaries. If we
         # race too far ahead with enqueuing, there is a distinct risk of
         # running out of memory. This mitigates that risk a bit, for now.
-        evt.wait()
+        if isinstance(evt, list):
+            [_evt.wait() for _evt in evt]
+        else:
+            evt.wait()
 
         def to_output_template(keys, _):
             name_in_program = self.output_id_to_name_in_program[keys]

arraycontext/parameter_study/__init__.py

@@ -0,0 +1,146 @@
+"""
+.. currentmodule:: arraycontext
+
+A :mod:`pytato`-based array context defers the evaluation of an array until its
+frozen. The execution contexts for the evaluations are specific to an
+:class:`~arraycontext.ArrayContext` type. For ex.
+:class:`~arraycontext.ParamStudyPytatoPyOpenCLArrayContext`
+uses :mod:`pyopencl` to JIT-compile and execute the array expressions.
+
+Following :mod:`pytato`-based array context are provided:
+
+.. autoclass:: ParamStudyPytatoPyOpenCLArrayContext
+
+The compiled function is stored as.
+.. autoclass:: ParamStudyLazyPyOpenCLFunctionCaller
+
+
+Compiling a Python callable (Internal) for multiple distinct instances of
+execution.
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. automodule:: arraycontext.parameter_study
+"""
+__copyright__ = """
+Copyright (C) 2020-1 University of Illinois Board of Trustees
+"""
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+import sys
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Mapping,
+    Tuple,
+    Type,
+)
+
+import numpy as np
+
+import pytato as pt
+from pytato.array import Array
+
+from arraycontext.context import ArrayContext
+from arraycontext.parameter_study.transform import ParameterStudyAxisTag
+
+
+ParamStudyTagT = Type[ParameterStudyAxisTag]
+
+if TYPE_CHECKING:
+    import pyopencl as cl
+    import pytato
+
+if getattr(sys, "_BUILDING_SPHINX_DOCS", False):
+    import pyopencl as cl
+
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def pack_for_parameter_study(actx: ArrayContext,
+                             study_name_tag_type: ParamStudyTagT,
+                             newshape: Tuple[int, ...],
+                             *args: Array) -> Array:
+    """
+        Args is a list of variable names and the realized input data that needs
+        to be packed for a parameter study or uncertainty quantification.
+
+        Args needs to be in the format
+            [v0, v1, v2, ..., vN] where N is the total number of instances you want to
+            try. Note these may be across multiple parameter studies on the same inputs.
+    """
+
+    assert len(args) > 0
+    assert len(args) == np.prod(newshape)
+
+    orig_shape = args[0].shape
+    out = actx.np.stack(args, axis=len(args[0].shape))
+    outshape = tuple([*list(orig_shape), newshape])
+
+    # if len(newshape) > 1:
+    #    # Reshape the object
+    #    out = out.reshape(outshape)
+
+    for i in range(len(orig_shape), len(outshape)):
+        out = out.with_tagged_axis(i, [study_name_tag_type(i - len(orig_shape),
+                                                           newshape[i-len(orig_shape)])])
+    return out
+
+
+def unpack_parameter_study(data: Array,
+                           study_name_tag_type: ParamStudyTagT) -> Mapping[int,
+                                                                          List[Array]]:
+    """
+        Split the data array along the axes which vary according to
+        a ParameterStudyAxisTag whose name tag is an instance study_name_tag_type.
+
+        output[i] corresponds to the values associated with the ith parameter study that
+        uses the variable name :arg: `study_name_tag_type`.
+    """
+
+    ndim: int = len(data.shape)
+    out: Dict[int, List[Array]] = {}
+
+    study_count = 0
+    for i in range(ndim):
+        axis_tags = data.axes[i].tags_of_type(study_name_tag_type)
+        if axis_tags:
+            # Now we need to split this data.
+            breakpoint()
+            for j in range(data.shape[i]):
+                tmp: List[Any] = [slice(None)] * ndim
+                tmp[i] = j
+                the_slice = tuple(tmp)
+                # Needs to be a tuple of slices not list of slices.
+                if study_count in out.keys():
+                    out[study_count].append(data[the_slice])
+                else:
+                    out[study_count] = [data[the_slice]]
+            if study_count in out.keys():
+                study_count += 1
+                # yield data[the_slice]
+
+    return out