Add normalization method to fields. Remove all old references to SphericalShell and use Shell instead.

Author: kburns

dedalus/core/basis.py

@@ -313,6 +313,7 @@ def __init__(self, coord, size, bounds, dealias):
         self.size = size
         self.shape = (size,)
         self.bounds = bounds
+        self.volume = bounds[1] - bounds[0]
         if isinstance(dealias, tuple):
             self.dealias = dealias
         else:
@@ -1960,6 +1961,7 @@ def __init__(self, coordsystem, shape, dtype, radii=(1,2), k=0, alpha=(-0.5,-0.5
                 radius_library = "matrix"
         self.radius_library = radius_library
         self.radii = tuple(radii)
+        self.volume = np.pi * (radii[1]**2 - radii[0]**2)
         self.dR = radii[1] - radii[0]
         self.rho = (radii[1] + radii[0])/self.dR
         self.alpha = tuple(alpha)
@@ -2207,6 +2209,7 @@ def __init__(self, coordsystem, shape, dtype, radius=1, k=0, alpha=0, dealias=(1
             radius_library = "matrix"
         self.radius_library = radius_library
         self.radius = radius
+        self.volume = np.pi * radius**2
         self.alpha = alpha
         self.radial_COV = AffineCOV((0, 1), (0, radius))
         if self.mmax > 2*self.Nmax:
@@ -2538,6 +2541,7 @@ def __init__(self, coordsystem, shape, dtype, radius=1, dealias=(1,1), colatitud
         if colatitude_library is None:
             colatitude_library = "matrix"
         self.radius = radius
+        self.volume = 4 * np.pi * radius**2
         self.colatitude_library = colatitude_library
         # Set Lmax for optimal load balancing
         if self.dtype == np.float64:
@@ -3153,7 +3157,7 @@ def symbol(spinindex_in, spinindex_out, ell, radius):
         return k_lap / radius**2
 
 
-# These are common for BallRadialBasis and SphericalShellRadialBasis
+# These are common for BallRadialBasis and ShellRadialBasis
 class RegularityBasis(SpinRecombinationBasis, MultidimensionalBasis):
 
     dim = 3
@@ -3414,7 +3418,7 @@ def n_slice(self, ell):
         return slice(nmin, nmax+1)
 
 
-class SphericalShellRadialBasis(RegularityBasis, metaclass=CachedClass):
+class ShellRadialBasis(RegularityBasis, metaclass=CachedClass):
 
     def __init__(self, coordsystem, radial_size, dtype, radii=(1,2), alpha=(-0.5,-0.5), dealias=(1,), k=0, radius_library=None):
         super().__init__(coordsystem, radial_size, k=k, dealias=dealias, dtype=dtype)
@@ -3426,6 +3430,7 @@ def __init__(self, coordsystem, radial_size, dtype, radii=(1,2), alpha=(-0.5,-0.
             else:
                 radius_library = "matrix"
         self.radii = radii
+        self.volume = 4 / 3 * np.pi * (radii[1]**3 - radii[0]**3)
         self.dR = self.radii[1] - self.radii[0]
         self.rho = (self.radii[1] + self.radii[0])/self.dR
         self.alpha = alpha
@@ -3439,7 +3444,7 @@ def __init__(self, coordsystem, radial_size, dtype, radii=(1,2), alpha=(-0.5,-0.
                                     self.backward_transform_radius]
 
     def __eq__(self, other):
-        if isinstance(other, SphericalShellRadialBasis):
+        if isinstance(other, ShellRadialBasis):
             if self.coordsystem == other.coordsystem:
                 if self.grid_params == other.grid_params:
                     if self.k == other.k:
@@ -3454,7 +3459,7 @@ def __add__(self, other):
             return self
         if other is self:
             return self
-        if isinstance(other, SphericalShellRadialBasis):
+        if isinstance(other, ShellRadialBasis):
             if self.grid_params == other.grid_params:
                 radial_size = max(self.shape[2], other.shape[2])
                 k = max(self.k, other.k)
@@ -3464,7 +3469,7 @@ def __add__(self, other):
     def __mul__(self, other):
         if other is None:
             return self
-        if isinstance(other, SphericalShellRadialBasis):
+        if isinstance(other, ShellRadialBasis):
             if self.grid_params == other.grid_params:
                 radial_size = max(self.shape[2], other.shape[2])
                 k = self.k + other.k
@@ -3482,7 +3487,7 @@ def __mul__(self, other):
             args['azimuth_library'] = other.azimuth_library
             args['colatitude_library'] = other.colatitude_library
             args['radius_library'] = self.radius_library
-            return SphericalShellBasis(**args)
+            return ShellBasis(**args)
         return NotImplemented
 
     def __matmul__(self, other):
@@ -3491,7 +3496,7 @@ def __matmul__(self, other):
     def __rmatmul__(self, other):
         if other is None:
             return self
-        if isinstance(other, SphericalShellRadialBasis):
+        if isinstance(other, ShellRadialBasis):
             if self.grid_params == other.grid_params:
                 radial_size = max(self.shape[2], other.shape[2])
                 k = self.k + other.k
@@ -3650,6 +3655,7 @@ def __init__(self, coordsystem, radial_size, dtype, radius=1, k=0, alpha=0, deal
         if radius_library is None:
             radius_library = "matrix"
         self.radius = radius
+        self.volume = 4 / 3 * np.pi * radius**3
         self.alpha = alpha
         self.radial_COV = AffineCOV((0, 1), (0, radius))
         self.radius_library = radius_library
@@ -4052,15 +4058,16 @@ def multiplication_matrix(self, subproblem, arg_basis, coeffs, *args, **kw):
             raise ValueError("Cannot build NCCs of non-radial fields.")
 
 
-class SphericalShellBasis(Spherical3DBasis, metaclass=CachedClass):
+class ShellBasis(Spherical3DBasis, metaclass=CachedClass):
 
     def __init__(self, coordsystem, shape, dtype, radii=(1,2), alpha=(-0.5,-0.5), dealias=(1,1,1), k=0, azimuth_library=None, colatitude_library=None, radius_library=None):
         if np.isscalar(dealias):
             dealias = (dealias, dealias, dealias)
         self.alpha = alpha
         self.radii = radii
+        self.volume = 4 / 3 * np.pi * (radii[1]**3 - radii[0]**3)
         self.radius_library = radius_library
-        self.radial_basis = SphericalShellRadialBasis(coordsystem, shape[2], radii=radii, alpha=alpha, dealias=(dealias[2],), k=k, dtype=dtype, radius_library=radius_library)
+        self.radial_basis = ShellRadialBasis(coordsystem, shape[2], radii=radii, alpha=alpha, dealias=(dealias[2],), k=k, dtype=dtype, radius_library=radius_library)
         Spherical3DBasis.__init__(self, coordsystem, shape[:2], dealias[:2], self.radial_basis, dtype=dtype, azimuth_library=azimuth_library, colatitude_library=colatitude_library)
         self.grid_params = (coordsystem, radii, alpha, dealias)
 #        self.forward_transform_radius = self.radial_basis.forward_transform
@@ -4073,7 +4080,7 @@ def __init__(self, coordsystem, shape, dtype, radii=(1,2), alpha=(-0.5,-0.5), de
                                     self.backward_transform_radius]
 
     def __eq__(self, other):
-        if isinstance(other, SphericalShellBasis):
+        if isinstance(other, ShellBasis):
             if self.coordsystem == other.coordsystem:
                 if self.grid_params == other.grid_params:
                     if self.k == other.k:
@@ -4088,40 +4095,40 @@ def __add__(self, other):
             return self
         if other is self:
             return self
-        if isinstance(other, SphericalShellBasis):
+        if isinstance(other, ShellBasis):
             if self.grid_params == other.grid_params:
                 shape = tuple(np.maximum(self.shape, other.shape))
                 k = max(self.k, other.k)
-                return SphericalShellBasis(self.coordsystem, shape, radii=self.radial_basis.radii, alpha=self.radial_basis.alpha, dealias=self.dealias, k=k,
+                return ShellBasis(self.coordsystem, shape, radii=self.radial_basis.radii, alpha=self.radial_basis.alpha, dealias=self.dealias, k=k,
                                            dtype=self.dtype, azimuth_library=self.azimuth_library, colatitude_library=self.colatitude_library,
                                            radius_library=self.radial_basis.radius_library)
         return NotImplemented
 
     def __mul__(self, other):
         if other is None:
             return self
-        if isinstance(other, SphericalShellBasis):
+        if isinstance(other, ShellBasis):
             if self.grid_params == other.grid_params:
                 shape = tuple(np.maximum(self.shape, other.shape))
                 k = 0
-                return SphericalShellBasis(self.coordsystem, shape, radii=self.radial_basis.radii, alpha=self.radial_basis.alpha, dealias=self.dealias, k=k,
+                return ShellBasis(self.coordsystem, shape, radii=self.radial_basis.radii, alpha=self.radial_basis.alpha, dealias=self.dealias, k=k,
                                            dtype=self.dtype, azimuth_library=self.azimuth_library, colatitude_library=self.colatitude_library,
                                            radius_library=self.radial_basis.radius_library)
-        if isinstance(other, SphericalShellRadialBasis):
+        if isinstance(other, ShellRadialBasis):
             radial_basis = other * self.radial_basis
             return self._new_k(radial_basis.k)
         return NotImplemented
 
     def __rmatmul__(self, other):
         if other is None:
             return self
-        if isinstance(other, SphericalShellRadialBasis):
+        if isinstance(other, ShellRadialBasis):
             radial_basis = other @ self.radial_basis
             return self._new_k(radial_basis.k)
         return NotImplemented
 
     def _new_k(self, k):
-        return SphericalShellBasis(self.coordsystem, self.shape, radii=self.radial_basis.radii, alpha=self.radial_basis.alpha, dealias=self.dealias, k=k,
+        return ShellBasis(self.coordsystem, self.shape, radii=self.radial_basis.radii, alpha=self.radial_basis.alpha, dealias=self.dealias, k=k,
                                    dtype=self.dtype, azimuth_library=self.azimuth_library, colatitude_library=self.colatitude_library,
                                    radius_library=self.radial_basis.radius_library)
 
@@ -4162,9 +4169,6 @@ def backward_transform_radius(self, field, axis, cdata, gdata):
             gdata *= radial_basis.radial_transform_factor(field.scales[axis], data_axis, self.k)
 
 
-ShellBasis = SphericalShellBasis
-
-
 class BallBasis(Spherical3DBasis, metaclass=CachedClass):
 
     transforms = {}
@@ -4174,6 +4178,7 @@ def __init__(self, coordsystem, shape, dtype, radius=1, k=0, alpha=0, dealias=(1
             dealias = (dealias, dealias, dealias)
         self.alpha = alpha
         self.radius = radius
+        self.volume = 4 / 3 * np.pi * radius**3
         self.radius_library = radius_library
         self.radial_basis = BallRadialBasis(coordsystem, shape[2], radius=radius, k=k, alpha=alpha, dealias=(dealias[2],), dtype=dtype, radius_library=radius_library)
         Spherical3DBasis.__init__(self, coordsystem, shape[:2], dealias[:2], self.radial_basis, dtype=dtype, azimuth_library=azimuth_library, colatitude_library=colatitude_library)
@@ -4347,7 +4352,7 @@ def radial_matrix(self, regindex_in, regindex_out, ell):
 
 class ConvertConstantShell(operators.ConvertConstant, operators.SphericalEllOperator):
 
-    output_basis_type = (ShellBasis, SphericalShellRadialBasis)
+    output_basis_type = (ShellBasis, ShellRadialBasis)
     subaxis_dependence = [False, True, True]
     subaxis_coupling = [False, False, True]
 
@@ -4657,7 +4662,7 @@ def radial_matrix(self, regindex_in, regindex_out, m):
 class LiftShell(operators.Lift, operators.SphericalEllOperator):
 
     input_basis_type = SphereBasis
-    output_basis_type = SphericalShellBasis
+    output_basis_type = ShellBasis
 
     def regindex_out(self, regindex_in):
         return (regindex_in,)
@@ -5252,9 +5257,9 @@ def _radial_matrix(basis, ell, regtotal, position):
         return reshape_vector(basis.interpolation(ell, regtotal, position), dim=2, axis=1)
 
 
-class SphericalShellRadialInterpolate(operators.Interpolate, operators.SphericalEllOperator):
+class ShellRadialInterpolate(operators.Interpolate, operators.SphericalEllOperator):
 
-    basis_type = SphericalShellBasis
+    basis_type = ShellBasis
     basis_subaxis = 2
 
     def __init__(self, operand, coord, position, out=None):

dedalus/core/domain.py

@@ -6,6 +6,7 @@
 import numpy as np
 from collections import OrderedDict
 
+from ..tools.array import prod
 from ..tools.cache import CachedMethod, CachedClass, CachedAttribute
 from ..tools.general import unify_attributes, unify, OrderedSet
 from .coords import Coordinate, CartesianCoordinates
@@ -45,6 +46,10 @@ def __init__(self, dist, bases):
         self.bases = bases  # Preprocessed to remove Nones and duplicates
         self.dim = sum(basis.dim for basis in self.bases)
 
+    @CachedAttribute
+    def volume(self):
+        return prod([basis.volume for basis in self.bases])
+
     @CachedAttribute
     def bases_by_axis(self):
         bases_by_axis = OrderedDict()

dedalus/core/field.py

@@ -731,6 +731,37 @@ def allreduce_data_norm(self, layout=None, order=2):
     def allreduce_data_max(self, layout=None):
         return self.allreduce_data_norm(layout=layout, order=np.inf)
 
+    def allreduce_L2_norm(self, normalize_volume=True):
+        from . import arithmetic
+        from . import operators
+        # Compute local self inner product
+        rank = len(self.tensorsig)
+        if rank == 0:
+            self_inner_product = np.conj(self) * self
+        elif rank == 1:
+            self_inner_product = arithmetic.dot(np.conj(self), self)
+        elif rank == 2:
+            self_inner_product = arithmetic.Trace(arithmetic.Dot(operators.Transpose(np.conj(self)), self))
+        else:
+            raise ValueError("Norms only implemented up to rank-2 tensors.")
+        # Compute L2 norm
+        norm_sq = operators.Integrate(self_inner_product).evaluate().allreduce_data_max()
+        if normalize_volume:
+            norm_sq /= self.domain.volume
+        return norm_sq ** 0.5
+
+    def normalize(self, normalize_volume=True):
+        """
+        Normalize field inplace using L2 norm.
+
+        Parameters
+        ----------
+        normalize_volume : bool, optional
+            Normalize inner product by domain volume. Default: True.
+        """
+        norm = self.allreduce_L2_norm(normalize_volume=normalize_volume)
+        self.data /= norm
+
     def broadcast_ghosts(self, output_nonconst_dims):
         """Copy data over constant distributed dimensions for arithmetic broadcasting."""
         # Determine deployment dimensions

dedalus/tests/test_grid_operators.py

@@ -43,7 +43,7 @@ def build_ball(Nphi, Ntheta, Nr, dtype, dealias, radius=1):
 def build_shell(Nphi, Ntheta, Nr, dtype, dealias, radii=(0.5,1)):
     c = coords.SphericalCoordinates('phi', 'theta', 'r')
     d = distributor.Distributor((c,))
-    b = basis.SphericalShellBasis(c, (Nphi, Ntheta, Nr), radii=radii, dtype=dtype, dealias=(dealias, dealias, dealias))
+    b = basis.ShellBasis(c, (Nphi, Ntheta, Nr), radii=radii, dtype=dtype, dealias=(dealias, dealias, dealias))
     phi, theta, r = b.local_grids()
     x, y, z = c.cartesian(phi, theta, r)
     return c, d, b, phi, theta, r, x, y, z

dedalus/tests/test_ivp.py

@@ -25,7 +25,7 @@ def build_shell(Nphi, Ntheta, Nr, radii_shell, dealias, dtype, grid_scale=1):
     c = coords.SphericalCoordinates('phi', 'theta', 'r')
     d = distributor.Distributor((c,))
     dealias_tuple = (dealias, dealias, dealias)
-    b = basis.SphericalShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, dealias=dealias_tuple, dtype=dtype)
+    b = basis.ShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, dealias=dealias_tuple, dtype=dtype)
     grid_scale_tuple = (grid_scale, grid_scale, grid_scale)
     phi, theta, r = b.local_grids(grid_scale_tuple)
     x, y, z = c.cartesian(phi, theta, r)

dedalus/tests/test_lbvp.py

@@ -283,7 +283,7 @@ def test_heat_ball_cart(Nmax, Lmax, dtype):
 def build_shell(Nphi, Ntheta, Nr, dealias, dtype):
     c = coords.SphericalCoordinates('phi', 'theta', 'r')
     d = distributor.Distributor((c,))
-    b = basis.SphericalShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, dealias=(dealias, dealias, dealias), dtype=dtype)
+    b = basis.ShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, dealias=(dealias, dealias, dealias), dtype=dtype)
     phi, theta, r = b.local_grids()
     x, y, z = c.cartesian(phi, theta, r)
     return c, d, b, phi, theta, r, x, y, z

dedalus/tests/test_output.py

@@ -71,7 +71,7 @@ def build_shell(Nphi, Ntheta, Nr, k, dealias, dtype):
     radii_shell = (0.5, 1.5)
     c = coords.SphericalCoordinates('phi', 'theta', 'r')
     d = distributor.Distributor((c,))
-    b = basis.SphericalShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, k=k, dealias=(dealias, dealias, dealias), dtype=dtype)
+    b = basis.ShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, k=k, dealias=(dealias, dealias, dealias), dtype=dtype)
     phi, theta, r = b.local_grids(b.domain.dealias)
     x, y, z = c.cartesian(phi, theta, r)
     return c, d, b, phi, theta, r, x, y, z

dedalus/tests/test_spherical3D_arithmetic.py

@@ -28,7 +28,7 @@ def build_ball(Nphi, Ntheta, Nr, dealias, dtype):
 def build_shell(Nphi, Ntheta, Nr, dealias, dtype):
     c = coords.SphericalCoordinates('phi', 'theta', 'r')
     d = distributor.Distributor((c,))
-    b = basis.SphericalShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, dealias=(dealias, dealias, dealias), dtype=dtype)
+    b = basis.ShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, dealias=(dealias, dealias, dealias), dtype=dtype)
     phi, theta, r = b.local_grids()
     x, y, z = c.cartesian(phi, theta, r)
     return c, d, b, phi, theta, r, x, y, z

dedalus/tests/test_spherical3D_calculus.py

@@ -28,7 +28,7 @@ def build_ball(Nphi, Ntheta, Nr, dealias, dtype):
 def build_shell(Nphi, Ntheta, Nr, dealias, dtype):
     c = coords.SphericalCoordinates('phi', 'theta', 'r')
     d = distributor.Distributor((c,))
-    b = basis.SphericalShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, dealias=(dealias, dealias, dealias), dtype=dtype)
+    b = basis.ShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, dealias=(dealias, dealias, dealias), dtype=dtype)
     phi, theta, r = b.local_grids(b.domain.dealias)
     x, y, z = c.cartesian(phi, theta, r)
     return c, d, b, phi, theta, r, x, y, z

dedalus/tests/test_spherical3D_operators.py

@@ -30,7 +30,7 @@ def build_ball(Nphi, Ntheta, Nr, k, dealias, dtype):
 def build_shell(Nphi, Ntheta, Nr, k, dealias, dtype):
     c = coords.SphericalCoordinates('phi', 'theta', 'r')
     d = distributor.Distributor((c,))
-    b = basis.SphericalShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, k=k, dealias=(dealias, dealias, dealias), dtype=dtype)
+    b = basis.ShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, k=k, dealias=(dealias, dealias, dealias), dtype=dtype)
     phi, theta, r = b.local_grids(b.domain.dealias)
     x, y, z = c.cartesian(phi, theta, r)
     return c, d, b, phi, theta, r, x, y, z