Fix 1D variable velocity example

Might have gone a bit overboard with moving stuff around in the examples, but the main fix is in RefStiffnessTOperator.

Depends on !57 (merged).

examples/advection/var-velocity.py

@@ -22,114 +22,199 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 """
 
-
+import os
 import numpy as np
-import pyopencl as cl  # noqa
-import pyopencl.array  # noqa
-import pyopencl.clmath  # noqa
 
-import pytest  # noqa
+import pyopencl as cl
+import pyopencl.array
+import pyopencl.clmath
 
-from pyopencl.tools import (  # noqa
-                pytest_generate_tests_for_pyopencl as pytest_generate_tests)
+from grudge import bind, sym
+from pytools.obj_array import join_fields
 
 import logging
 logger = logging.getLogger(__name__)
 
-from grudge import sym, bind, DGDiscretizationWithBoundaries
-from pytools.obj_array import join_fields
 
+# {{{ plotting (keep in sync with `weak.py`)
 
-FINAL_TIME = 5
+class Plotter:
+    def __init__(self, queue, discr, order, visualize=True, ylim=None):
+        self.queue = queue
+        self.dim = discr.ambient_dim
 
+        self.visualize = visualize
+        if not self.visualize:
+            return
 
-def main(write_output=True, order=4):
-    cl_ctx = cl.create_some_context()
-    queue = cl.CommandQueue(cl_ctx)
+        if self.dim == 1:
+            import matplotlib.pyplot as pt
+            self.fig = pt.figure(figsize=(8, 8), dpi=300)
+            self.ylim = ylim
 
-    dim = 2
+            volume_discr = discr.discr_from_dd(sym.DD_VOLUME)
+            self.x = volume_discr.nodes().get(self.queue)
+        else:
+            from grudge.shortcuts import make_visualizer
+            self.vis = make_visualizer(discr, vis_order=order)
 
-    resolution = 15
-    from meshmode.mesh.generation import generate_regular_rect_mesh
-    mesh = generate_regular_rect_mesh(a=(-0.5, -0.5), b=(0.5, 0.5),
-            n=(resolution, resolution), order=order)
+    def __call__(self, evt, basename, overwrite=True):
+        if not self.visualize:
+            return
 
-    dt_factor = 5
-    h = 1/resolution
+        if self.dim == 1:
+            u = evt.state_component.get(self.queue)
 
-    sym_x = sym.nodes(2)
+            filename = "%s.png" % basename
+            if not overwrite and os.path.exists(filename):
+                from meshmode import FileExistsError
+                raise FileExistsError("output file '%s' already exists" % filename)
 
-    advec_v = join_fields(-1*sym_x[1], sym_x[0]) / 2
+            ax = self.fig.gca()
+            ax.plot(self.x[0], u, '-')
+            ax.plot(self.x[0], u, 'k.')
+            if self.ylim is not None:
+                ax.set_ylim(self.ylim)
 
-    flux_type = "upwind"
+            ax.set_xlabel("$x$")
+            ax.set_ylabel("$u$")
+            ax.set_title("t = {:.2f}".format(evt.t))
+            self.fig.savefig(filename)
+            self.fig.clf()
+        else:
+            self.vis.write_vtk_file("%s.vtu" % basename, [
+                ("u", evt.state_component)
+                ], overwrite=overwrite)
 
-    source_center = np.array([0.1, 0.1])
-    source_width = 0.05
+# }}}
 
-    sym_x = sym.nodes(2)
-    sym_source_center_dist = sym_x - source_center
 
-    def f_gaussian(x):
-        return sym.exp(
-                    -np.dot(sym_source_center_dist, sym_source_center_dist)
-                    / source_width**2)
+def main(ctx_factory, dim=2, order=4, product_tag=None, visualize=True):
+    cl_ctx = ctx_factory()
+    queue = cl.CommandQueue(cl_ctx)
 
-    def f_step(x):
-        return sym.If(
-                sym.Comparison(
-                    np.dot(sym_source_center_dist, sym_source_center_dist),
-                    "<",
-                    (4*source_width)**2),
-                1, 0)
+    # {{{ parameters
 
-    def u_analytic(x):
-        return 0
+    # domain [0, d]^dim
+    d = 1.0
+    # number of points in each dimension
+    npoints = 25
+    # grid spacing
+    h = d / npoints
 
-    from grudge.models.advection import VariableCoefficientAdvectionOperator
-    from meshmode.discretization.poly_element import QuadratureSimplexGroupFactory  # noqa
+    # cfl
+    dt_factor = 1.0
+    # finale time
+    final_time = 0.5
+    # time steps
+    dt = dt_factor * h/order**2
+    nsteps = int(final_time // dt) + 1
+    dt = final_time/nsteps + 1.0e-15
+
+    # flux
+    flux_type = "upwind"
+    # velocity field
+    sym_x = sym.nodes(dim)
+    if dim == 1:
+        c = sym_x
+    else:
+        # solid body rotation
+        c = join_fields(
+                np.pi * (d/2 - sym_x[1]),
+                np.pi * (sym_x[0] - d/2),
+                0)[:dim]
+
+    # }}}
+
+    # {{{ discretization
 
+    from meshmode.mesh.generation import generate_regular_rect_mesh
+    mesh = generate_regular_rect_mesh(
+            a=(0,)*dim, b=(d,)*dim,
+            n=(npoints,)*dim,
+            order=order)
+
+    from meshmode.discretization.poly_element import \
+            QuadratureSimplexGroupFactory
+
+    if product_tag:
+        quad_tag_to_group_factory = {
+                product_tag: QuadratureSimplexGroupFactory(order=4*order)
+                }
+    else:
+        quad_tag_to_group_factory = {}
+
+    from grudge import DGDiscretizationWithBoundaries
     discr = DGDiscretizationWithBoundaries(cl_ctx, mesh, order=order,
-            quad_tag_to_group_factory={
-                #"product": None,
-                "product": QuadratureSimplexGroupFactory(order=4*order)
-                })
+            quad_tag_to_group_factory=quad_tag_to_group_factory)
+
+    # }}}
+
+    # {{{ symbolic operators
+
+    # gaussian parameters
+    source_center = np.array([0.5, 0.75, 0.0])[:dim]
+    source_width = 0.05
+    dist_squared = np.dot(sym_x - source_center, sym_x - source_center)
+
+    def f_gaussian(x):
+        return sym.exp(-dist_squared / source_width**2)
 
-    op = VariableCoefficientAdvectionOperator(2, advec_v,
-        u_analytic(sym.nodes(dim, sym.BTAG_ALL)), quad_tag="product",
-        flux_type=flux_type)
+    def f_step(x):
+        return sym.If(sym.Comparison(
+            dist_squared, "<", (4*source_width) ** 2),
+            1, 0)
 
-    bound_op = bind(
-            discr, op.sym_operator(),
-            #debug_flags=["dump_sym_operator_stages"]
-            )
+    def u_bc(x):
+        return 0.0
 
+    from grudge.models.advection import VariableCoefficientAdvectionOperator
+    op = VariableCoefficientAdvectionOperator(
+            c,
+            u_bc(sym.nodes(dim, sym.BTAG_ALL)),
+            quad_tag=product_tag,
+            flux_type=flux_type)
+
+    bound_op = bind(discr, op.sym_operator())
     u = bind(discr, f_gaussian(sym.nodes(dim)))(queue, t=0)
 
     def rhs(t, u):
         return bound_op(queue, t=t, u=u)
 
-    dt = dt_factor * h/order**2
-    nsteps = (FINAL_TIME // dt) + 1
-    dt = FINAL_TIME/nsteps + 1e-15
+    # }}}
+
+    # {{{ time stepping
 
     from grudge.shortcuts import set_up_rk4
     dt_stepper = set_up_rk4("u", dt, u, rhs)
-
-    from grudge.shortcuts import make_visualizer
-    vis = make_visualizer(discr, vis_order=2*order)
+    plot = Plotter(queue, discr, order, visualize=visualize,
+            ylim=[-0.1, 1.1])
 
     step = 0
+    norm = bind(discr, sym.norm(2, sym.var("u")))
+    for event in dt_stepper.run(t_end=final_time):
+        if not isinstance(event, dt_stepper.StateComputed):
+            continue
 
-    for event in dt_stepper.run(t_end=FINAL_TIME):
-        if isinstance(event, dt_stepper.StateComputed):
+        if step % 10 == 0:
+            norm_u = norm(queue, u=event.state_component)
+            plot(event, "fld-var-velocity-%04d" % step)
 
-            step += 1
-            if step % 30 == 0:
-                print(step)
+        step += 1
+        logger.info("[%04d] t = %.5f |u| = %.5e", step, event.t, norm_u)
 
-                vis.write_vtk_file("fld-var-velocity-%04d.vtu" % step,
-                        [("u", event.state_component)])
+    # }}}
 
 
 if __name__ == "__main__":
-    main()
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dim", default=2, type=int)
+    parser.add_argument("--qtag", default="product")
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO)
+    main(cl.create_some_context,
+            dim=args.dim,
+            product_tag=args.qtag)