Clean up integration interface

Author: kburns

dedalus/core/operators.py

@@ -1085,20 +1085,40 @@ def integrate(arg, spaces=None):
         arg = Integrate(arg, space)
     return arg
 
-
+@alias("integ")
 class Integrate(LinearOperator, metaclass=MultiClass):
     """
-    Integration along one dimension.
+    Definite integration over operand bases.
 
     Parameters
     ----------
     operand : number or Operand object
-    space : Space object
-
+    coords : Coordinate or CoordinateSystem object, or list of these
     """
 
     name = "Integrate"
 
+    @classmethod
+    def _preprocess_args(cls, operand, coord=None):
+        # Handle zeros
+        if operand == 0:
+            raise SkipDispatchException(output=0)
+        # Integrate over all operand bases by default
+        if coord is None:
+            coord = [basis.coordsystem for basis in operand.domain.bases]
+        # Recurse over multiple coordinates
+        if isinstance(coord, (tuple, list)):
+            if len(coord) > 1:
+                operand = Integrate(operand, coord[:-1])
+            coord = coord[-1]
+        # Resolve strings to coordinates
+        if isinstance(coord, str):
+            coord = operand.domain.get_coord(coord)
+        # Check coordinate type
+        if not isinstance(coord, (coords.Coordinate, coords.CoordinateSystem)):
+            raise ValueError("coords must be Coordinate or str")
+        return (operand, coord), {}
+
     @classmethod
     def _check_args(cls, operand, coords):
         # Dispatch by operand basis
@@ -1109,18 +1129,6 @@ def _check_args(cls, operand, coords):
                     return True
         return False
 
-    @classmethod
-    def _preprocess_args(cls, operand, coord):
-        if isinstance(operand, Number):
-            raise SkipDispatchException(output=operand)
-        if isinstance(coord, (coords.Coordinate, coords.CoordinateSystem)):
-            pass
-        elif isinstance(coord, str):
-            coord = operand.domain.get_coord(coord)
-        else:
-            raise ValueError("coord must be Coordinate or str")
-        return (operand, coord), {}
-
     def __init__(self, operand, coord):
         SpectralOperator.__init__(self, operand)
         # Require integrand is a scalar

examples/ivp_2d_rayleigh_benard/rayleigh_benard.py

@@ -28,8 +28,6 @@
 import logging
 logger = logging.getLogger(__name__)
 
-# TODO: cleanup integ shortcuts
-
 
 # Parameters
 Lx, Lz = 4, 1
@@ -63,7 +61,6 @@
 nu = (Rayleigh / Prandtl)**(-1/2)
 x, z = dist.local_grids(xbasis, zbasis)
 ex, ez = coords.unit_vector_fields(dist)
-integ = lambda A: d3.Integrate(d3.Integrate(A, 'x'), 'z')
 lift_basis = zbasis.clone_with(a=1/2, b=1/2) # First derivative basis
 lift = lambda A: d3.Lift(A, lift_basis, -1)
 grad_u = d3.grad(u) + ez*lift(tau_u1) # First-order reduction

examples/ivp_2d_shear_flow/shear_flow.py

@@ -24,8 +24,6 @@
 import logging
 logger = logging.getLogger(__name__)
 
-# TODO: cleanup integ shortcuts
-
 
 # Parameters
 Lx, Lz = 1, 2
@@ -55,7 +53,6 @@
 D = nu / Schmidt
 x, z = dist.local_grids(xbasis, zbasis)
 ex, ez = coords.unit_vector_fields(dist)
-integ = lambda A: d3.Integrate(d3.Integrate(A, 'x'), 'z')
 
 # Problem
 problem = d3.IVP([u, s, p, tau_p], namespace=locals())

examples/ivp_ball_internally_heated_convection/internally_heated_convection.py

@@ -75,8 +75,6 @@
 lift = lambda A: d3.Lift(A, basis, -1)
 strain_rate = d3.grad(u) + d3.trans(d3.grad(u))
 shear_stress = d3.angular(d3.radial(strain_rate(r=1), index=1))
-integ = lambda A: d3.Integrate(A, coords)
-rad = d3.RadialComponent
 
 # Problem
 problem = d3.IVP([p, u, T, tau_p, tau_u, tau_T], namespace=locals())
@@ -85,7 +83,7 @@
 problem.add_equation("dt(T) - kappa*lap(T) + lift(tau_T) = - dot(u,grad(T)) + kappa*T_source")
 problem.add_equation("shear_stress = 0")  # Stress free
 problem.add_equation("radial(u(r=1)) = 0")  # No penetration
-problem.add_equation("rad(grad(T)(r=1)) = -2")
+problem.add_equation("radial(grad(T)(r=1)) = -2")
 problem.add_equation("integ(p) = 0")  # Pressure gauge
 
 # Solver

examples/ivp_disk_libration/libration.py

@@ -53,7 +53,6 @@
 # Substitutions
 phi, r = dist.local_grids(basis)
 nu = Ekman
-integ = lambda A: d3.Integrate(A, coords)
 lift = lambda A: d3.Lift(A, basis, -1)
 
 # Background librating flow
@@ -83,7 +82,7 @@
 snapshots = solver.evaluator.add_file_handler('snapshots', sim_dt=0.1, max_writes=20)
 snapshots.add_task(-d3.div(d3.skew(u)), scales=(4, 1), name='vorticity')
 scalars = solver.evaluator.add_file_handler('scalars', sim_dt=0.01)
-scalars.add_task(integ(0.5*d3.dot(u,u)), name='KE')
+scalars.add_task(d3.integ(0.5*d3.dot(u,u)), name='KE')
 
 # Flow properties
 flow = d3.GlobalFlowProperty(solver, cadence=100)

examples/ivp_shell_convection/shell_convection.py

@@ -69,7 +69,6 @@
 lift = lambda A: d3.Lift(A, lift_basis, -1)
 grad_u = d3.grad(u) + rvec*lift(tau_u1) # First-order reduction
 grad_b = d3.grad(b) + rvec*lift(tau_b1) # First-order reduction
-integ = lambda A: d3.Integrate(A, coords)
 
 # Problem
 problem = d3.IVP([p, b, u, tau_p, tau_b1, tau_b2, tau_u1, tau_u2], namespace=locals())