Merge pull request #332 from DedalusProject/cardinal

Cardinal basis

Author: kburns

dedalus/core/basis.py

@@ -26,7 +26,8 @@
 
 
 # Public interface
-__all__ = ['Jacobi',
+__all__ = ['CardinalBasis',
+           'Jacobi',
            'Legendre',
            'Ultraspherical',
            'Chebyshev',
@@ -334,10 +335,166 @@ def enum_indices(tensorsig):
         # self._ncc_matrices = [self._ncc_matrix_recursion(ncc.data[ind], ncc.domain.full_bases, operand.domain.full_bases, separability, **kw) for ind in np.ndindex(*tshape)]
 
 
-class IntervalBasis(Basis):
+class CardinalBasis(Basis):
+    """Cardinal basis."""
 
     dim = 1
+    group_shape = (1,)
+    subaxis_dependence = [False]
+
+    def __init__(self, coord, size):
+        self.coord = coord
+        self.coordsys = coord
+        self.size = size
+        self.shape = (size,)
+        self.dealias = (1,)
+        super().__init__(coord)
+
+    def __add__(self, other):
+        if other is None or other is self:
+            return self
+        return NotImplemented
+
+    def __mul__(self, other):
+        if other is None or other is self:
+            return self
+        return NotImplemented
+
+    def __rmatmul__(self, other):
+        # NCC (other) * operand (self)
+        if other is None or other is self:
+            return self
+        return NotImplemented
+
+    def elements_to_groups(self, grid_space, elements):
+        # No permutations
+        return elements
+
+    def valid_elements(self, tensorsig, grid_space, elements):
+        # No invalid modes
+        vshape = tuple(cs.dim for cs in tensorsig) + elements[0].shape
+        return np.ones(shape=vshape, dtype=bool)
+
+    def matrix_dependence(self, matrix_coupling):
+        return matrix_coupling
+
+    def global_grids(self, dist, scales):
+        """Global grids."""
+        return (self.global_grid(dist, scales[0]),)
+
+    def global_grid(self, dist, scale):
+        """Global grid."""
+        if scale != 1:
+            raise NotImplementedError("Cardinal basis only supports scale=1.")
+        return np.arange(self.size)
+
+    def local_grids(self, dist, scales):
+        """Local grids."""
+        return (self.local_grid(dist, scales[0]),)
+
+    def local_grid(self, dist, scale):
+        """Local grid."""
+        if scale != 1:
+            raise NotImplementedError("Cardinal basis only supports scale=1.")
+        local_elements = dist.grid_layout.local_elements(self.domain(dist), scales=scale)
+        return np.arange(self.size)[local_elements[dist.get_basis_axis(self)]]
+
+    def local_modes(self, dist):
+        """Local grid."""
+        local_elements = dist.coeff_layout.local_elements(self.domain(dist), scales=1)
+        return reshape_vector(local_elements[dist.get_basis_axis(self)], dim=dist.dim, axis=dist.get_basis_axis(self))
+
+    def global_shape(self, grid_space, scales):
+        return self.shape
+
+    def chunk_shape(self, grid_space):
+        return (1,)
+
+    def forward_transform(self, field, axis, gdata, cdata):
+        """Forward transform field data."""
+        np.copyto(cdata, gdata)
+
+    def backward_transform(self, field, axis, cdata, gdata):
+        """Backward transform field data."""
+        np.copyto(gdata, cdata)
+
+
+class ConvertConstantCardinal(operators.ConvertConstant, operators.SpectralOperator1D):
+    """Convert constant to Cardinal basis."""
+
+    output_basis_type = CardinalBasis
+    subaxis_dependence = [True]
+    subaxis_coupling = [True]
+
+    @staticmethod
+    def _full_matrix(input_basis, output_basis):
+        return np.ones((output_basis.size, 1))
+
+
+class InterpolateCardinal(operators.Interpolate, operators.SpectralOperator1D):
+    """Interpolate Cardinal basis."""
+
+    input_basis_type = CardinalBasis
+    basis_subaxis = 0
+    subaxis_dependence = [True]
+    subaxis_coupling = [True]
+
+    def __init__(self, operand, coord, position, out=None):
+        if not isinstance(position, (int, np.integer)):
+            raise TypeError("Cardinal interpolation position must be an integer")
+        super().__init__(operand, coord, position, out=out)
+
+    @staticmethod
+    def _output_basis(input_basis, position):
+        return None
+
+    @staticmethod
+    def _full_matrix(input_basis, output_basis, position):
+        interp_vector = np.zeros(input_basis.size)
+        interp_vector[position] = 1
+        return interp_vector[None, :]
+
+
+class IntegrateCardinal(operators.Integrate, operators.SpectralOperator1D):
+    """Cardinal basis integration."""
+
+    input_coord_type = Coordinate
+    input_basis_type = CardinalBasis
+    subaxis_dependence = [True]
+    subaxis_coupling = [True]
+
+    @staticmethod
+    def _output_basis(input_basis):
+        return None
+
+    @staticmethod
+    def _full_matrix(input_basis, output_basis):
+        integ_vector = np.ones(input_basis.size)
+        return integ_vector[None, :]
+
+
+class AverageCardinal(operators.Average, operators.SpectralOperator1D):
+    """Cardinal basis averaging."""
+
+    input_coord_type = Coordinate
+    input_basis_type = CardinalBasis
     subaxis_dependence = [True]
+    subaxis_coupling = [True]
+
+    @staticmethod
+    def _output_basis(input_basis):
+        return None
+
+    @staticmethod
+    def _full_matrix(input_basis, output_basis):
+        ave_vector = np.ones(input_basis.size) / input_basis.size
+        return ave_vector[None, :]
+
+
+class IntervalBasis(Basis):
+
+    dim = 1
+    subaxis_dependence = [False]
 
     def __init__(self, coord, size, bounds, dealias):
         self.coord = coord
@@ -6084,15 +6241,16 @@ def cfl_spacing(self):
         velocity = self.operand
         coordsys = velocity.tensorsig[0]
         spacing = []
-        for i, c in enumerate(coordsys.coords):
+        for c in coordsys.coords:
             basis = velocity.domain.get_basis(c)
             if basis:
                 dealias = basis.dealias[0]
                 axis_spacing = basis.local_grid_spacing(self.dist, dealias) * dealias
                 N = basis.grid_shape((dealias,))[0]
             if isinstance(basis, Jacobi) and basis.a == -1/2 and basis.b == -1/2:
                 #Special case for ChebyshevT (a=b=-1/2)
-                local_elements = self.dist.grid_layout.local_elements(basis.domain(self.dist), scales=dealias)[i]
+                axis = self.dist.get_basis_axis(basis)
+                local_elements = self.dist.grid_layout.local_elements(basis.domain(self.dist), scales=dealias)[axis]
                 i = np.arange(N)[local_elements].reshape(axis_spacing.shape)
                 theta = np.pi * (i + 1/2) / N
                 axis_spacing[:] = dealias * basis.COV.stretch * np.sin(theta) * np.pi / N

dedalus/tests/test_cardinal_operators.py

@@ -0,0 +1,89 @@
+"""Test Cardinal basis operators: ConvertConstant, Interpolate, Integrate, Average."""
+
+import pytest
+import numpy as np
+import dedalus.public as d3
+from dedalus.tools.cache import CachedMethod
+
+
+N_range = [5, 10]
+dtype_range = [np.float64, np.complex128]
+
+
+@CachedMethod
+def build_cardinal(N, dtype):
+    c = d3.Coordinate('n')
+    dist = d3.Distributor(c, dtype=dtype)
+    b = d3.CardinalBasis(c, size=N)
+    n = dist.local_grid(b, scale=1)
+    return c, dist, b, n
+
+
+@pytest.mark.parametrize('N', N_range)
+@pytest.mark.parametrize('dtype', dtype_range)
+@pytest.mark.parametrize('layout', ['g', 'c'])
+def test_cardinal_convert_constant(N, dtype, layout):
+    """Test conversion from constant to Cardinal basis (broadcasts scalar to all N entries)."""
+    c, dist, b, n = build_cardinal(N, dtype)
+    f = dist.Field()
+    f['g'] = 3
+    f.change_layout(layout)
+    g = d3.Convert(f, b).evaluate()
+    assert np.allclose(g['g'], 3 * np.ones(N))
+
+
+@pytest.mark.parametrize('N', N_range)
+@pytest.mark.parametrize('dtype', dtype_range)
+@pytest.mark.parametrize('index', [0, 2, -1])
+def test_cardinal_interpolate(N, dtype, index):
+    """Test Interpolate extracts the value at the given integer index."""
+    c, dist, b, n = build_cardinal(N, dtype)
+    f = dist.Field(bases=b)
+    f.fill_random('g')
+    g = d3.Interpolate(f, c, index).evaluate()
+    assert np.allclose(g['g'], f['g'][index])
+
+
+@pytest.mark.parametrize('N', N_range)
+@pytest.mark.parametrize('dtype', dtype_range)
+def test_cardinal_integrate(N, dtype):
+    """Test Integrate computes the discrete sum over all entries."""
+    c, dist, b, n = build_cardinal(N, dtype)
+    f = dist.Field(bases=b)
+    f.fill_random('g')
+    g = d3.Integrate(f, c).evaluate()
+    assert np.allclose(g['g'], f['g'].sum())
+
+
+@pytest.mark.parametrize('N', N_range)
+@pytest.mark.parametrize('dtype', dtype_range)
+def test_cardinal_integrate_constant(N, dtype):
+    """Test Integrate of a uniform field equals N * value."""
+    c, dist, b, n = build_cardinal(N, dtype)
+    f = dist.Field(bases=b)
+    f['g'] = 3
+    g = d3.Integrate(f, c).evaluate()
+    assert np.allclose(g['g'], 3 * N)
+
+
+@pytest.mark.parametrize('N', N_range)
+@pytest.mark.parametrize('dtype', dtype_range)
+def test_cardinal_average(N, dtype):
+    """Test Average computes the mean over all entries."""
+    c, dist, b, n = build_cardinal(N, dtype)
+    f = dist.Field(bases=b)
+    f.fill_random('g')
+    g = d3.Average(f, c).evaluate()
+    assert np.allclose(g['g'], f['g'].mean())
+
+
+@pytest.mark.parametrize('N', N_range)
+@pytest.mark.parametrize('dtype', dtype_range)
+def test_cardinal_average_constant(N, dtype):
+    """Test Average of a uniform field returns that value."""
+    c, dist, b, n = build_cardinal(N, dtype)
+    f = dist.Field(bases=b)
+    f['g'] = 7
+    g = d3.Average(f, c).evaluate()
+    assert np.allclose(g['g'], 7)
+

docs/notebooks/dedalus_tutorial_1.ipynb

@@ -246,8 +246,10 @@
     "\n",
     "* `RealFourier` for real periodic functions on an interval using cosine & sine modes.\n",
     "* `ComplexFourier` for complex periodic functions on an interval using complex exponentials.\n",
-    "* `Chebyshev` for functions on an interval.\n",
+    "* `Chebyshev` for functions on an interval with fast tranforms (DCTs).\n",
+    "* `Legendre` for functions on an interval with L2 orthogonality.\n",
     "* `Jacobi` for functions on an interval under a more general inner product (usually `Chebyshev` is best for performance).\n",
+    "* `CardinalBasis` for a finite set of discrete elements (e.g., ensemble members or parameter sweeps).\n",
     "* `DiskBasis` for functions on a full disk in polar coordinates.\n",
     "* `AnnulusBasis` for functions on an annulus in polar coordinates.\n",
     "* `SphereBasis` for functions on the 2-sphere in S2 or spherical coordinates.\n",