port examples to arraycontext

inducer · Jun 29, 2022 · 4b7e4ae · 4b7e4ae
1 parent 5f89ebf
commit 4b7e4ae
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 41 deletions.
diff --git a/examples/cost_model.py b/examples/cost_model.py
@@ -1,29 +1,16 @@
+import os
+
 import numpy as np
 import pyopencl as cl
-import sys
 
 import logging
-import os
-
-# Configure the root logger
 logging.basicConfig(level=os.environ.get("LOGLEVEL", "WARNING"))
-
 logger = logging.getLogger(__name__)
-
-# Set the logger level of this module to INFO so that logging outputs of this module
-# are shown
 logger.setLevel(logging.INFO)
 
-# `process_elapsed` in `ProcessTimer` is only supported for Python >= 3.3
-SUPPORTS_PROCESS_TIME = (sys.version_info >= (3, 3))
-
 
 def demo_cost_model():
-    if not SUPPORTS_PROCESS_TIME:
-        raise NotImplementedError(
-            "Currently this script uses process time which only works on Python>=3.3"
-        )
-
+    from boxtree.array_context import PyOpenCLArrayContext, to_numpy
     from boxtree.pyfmmlib_integration import (
             Kernel,
             FMMLibTreeIndependentDataForWrangler,
@@ -36,6 +23,7 @@ def demo_cost_model():
 
     ctx = cl.create_some_context()
     queue = cl.CommandQueue(ctx)
+    actx = PyOpenCLArrayContext(queue, force_device_scalars=True)
 
     traversals = []
     traversals_dev = []
@@ -49,30 +37,27 @@ def fmm_level_to_order(tree, ilevel):
         # {{{ Generate sources, targets and target_radii
 
         from boxtree.tools import make_normal_particle_array as p_normal
-        sources = p_normal(queue, nsources, dims, dtype, seed=15)
-        targets = p_normal(queue, ntargets, dims, dtype, seed=18)
+        sources = p_normal(actx, nsources, dims, dtype, seed=15)
+        targets = p_normal(actx, ntargets, dims, dtype, seed=18)
 
-        from pyopencl.clrandom import PhiloxGenerator
-        rng = PhiloxGenerator(queue.context, seed=22)
-        target_radii = rng.uniform(
-            queue, ntargets, a=0, b=0.05, dtype=dtype
-        ).get()
+        rng = np.random.default_rng(seed=22)
+        target_radii = rng.uniform(low=0.0, high=0.05, size=ntargets)
 
         # }}}
 
         # {{{ Generate tree and traversal
 
         from boxtree import TreeBuilder
-        tb = TreeBuilder(ctx)
+        tb = TreeBuilder(actx)
         tree, _ = tb(
-            queue, sources, targets=targets, target_radii=target_radii,
+            actx, sources, targets=targets, target_radii=target_radii,
             stick_out_factor=0.15, max_particles_in_box=30, debug=True
         )
 
         from boxtree.traversal import FMMTraversalBuilder
-        tg = FMMTraversalBuilder(ctx, well_sep_is_n_away=2)
-        trav_dev, _ = tg(queue, tree, debug=True)
-        trav = trav_dev.get(queue=queue)
+        tg = FMMTraversalBuilder(actx, well_sep_is_n_away=2)
+        trav_dev, _ = tg(actx, tree, debug=True)
+        trav = to_numpy(trav_dev, actx)
 
         traversals.append(trav)
         traversals_dev.append(trav_dev)
@@ -103,7 +88,7 @@ def fmm_level_to_order(tree, ilevel):
         traversal = traversals_dev[icase]
         model_results.append(
             cost_model.cost_per_stage(
-                queue, traversal, level_orders_list[icase],
+                actx, traversal, level_orders_list[icase],
                 FMMCostModel.get_unit_calibration_params(),
             )
         )
@@ -114,7 +99,7 @@ def fmm_level_to_order(tree, ilevel):
     )
 
     predicted_time = cost_model.cost_per_stage(
-        queue, traversals_dev[-1], level_orders_list[-1], params,
+        actx, traversals_dev[-1], level_orders_list[-1], params,
     )
     queue.finish()
 

diff --git a/examples/demo.py b/examples/demo.py
@@ -4,46 +4,51 @@
 import logging
 logging.basicConfig(level="INFO")
 
+from boxtree.array_context import PyOpenCLArrayContext, to_numpy
+
 ctx = cl.create_some_context()
 queue = cl.CommandQueue(ctx)
+actx = PyOpenCLArrayContext(queue, force_device_scalars=True)
 
 dims = 2
 nparticles = 500
 
 # -----------------------------------------------------------------------------
 # generate some random particle positions
 # -----------------------------------------------------------------------------
-from pyopencl.clrandom import PhiloxGenerator
-rng = PhiloxGenerator(ctx, seed=15)
-
 from pytools.obj_array import make_obj_array
+rng = np.random.default_rng(seed=15)
+
 particles = make_obj_array([
-    rng.normal(queue, nparticles, dtype=np.float64)
+    actx.from_numpy(rng.normal(size=nparticles))
     for i in range(dims)])
 
 # -----------------------------------------------------------------------------
 # build tree and traversals (lists)
 # -----------------------------------------------------------------------------
 from boxtree import TreeBuilder
-tb = TreeBuilder(ctx)
-tree, _ = tb(queue, particles, max_particles_in_box=5)
+tb = TreeBuilder(actx)
+tree, _ = tb(actx, particles, max_particles_in_box=5)
 
 from boxtree.traversal import FMMTraversalBuilder
-tg = FMMTraversalBuilder(ctx)
-trav, _ = tg(queue, tree)
+tg = FMMTraversalBuilder(actx)
+trav, _ = tg(actx, tree)
 
 # ENDEXAMPLE
 
 # -----------------------------------------------------------------------------
 # plot the tree
 # -----------------------------------------------------------------------------
 
+particles = to_numpy(particles, actx)
+tree = to_numpy(tree, actx)
+
 import matplotlib.pyplot as pt
+from boxtree.visualization import TreePlotter
 
-pt.plot(particles[0].get(), particles[1].get(), "+")
+pt.plot(particles[0], particles[1], "+")
+plotter = TreePlotter(tree)
 
-from boxtree.visualization import TreePlotter
-plotter = TreePlotter(tree.get(queue=queue))
 plotter.draw_tree(fill=False, edgecolor="black")
 #plotter.draw_box_numbers()
 plotter.set_bounding_box()