1 files changed, 55 insertions, 0 deletions
diff --git a/basis.py b/basis.py
new file mode 100644
index 0000000..759ebfe
--- /dev/null
+++ b/basis.py
@@ -0,0 +1,55 @@
+import abc
+import math
+import numpy as np
+
+class ExpansionBasis1D(object):
+    """
+    A family of one-dimensional functions that make up a basis.
+    """
+
+    __metaclass__ = abc.ABCMeta
+
+    @abc.abstractmethod
+    def eval(self, idx, coord, diff_order = 0):
+        """
+        Evaluate the diff_order-th derivative of the idx-th basis function at
+        the specified coordinates (where zeroth derivative means evaluating the
+        function itself).
+        """
+        pass
+
+    @abc.abstractmethod
+    def colloc_grid(self, order):
+        """
+        Get the coordinates of the optimal collocation grid of the specified
+        order (i.e. exactly order coordinates will be returned).
+        """
+        pass
+
+class CosBasis(ExpansionBasis1D):
+    PARITY_NONE = 0
+    PARITY_EVEN = 1
+    PARITY_ODD  = 2
+
+    _diff_fact = [(1, np.cos), (-1, np.sin), (-1, np.cos), (1, np.sin)]
+    _parity = None
+
+    def __init__(self, parity = PARITY_NONE):
+        self._parity = parity
+
+    def eval(self, idx, coord, diff_order = 0):
+        fact, f = self._diff_fact[diff_order % 4]
+
+        if self._parity == self.PARITY_EVEN:
+            idx *= 2
+        elif self._parity == self.PARITY_ODD:
+            idx = 2 * idx + 1
+
+        fact *= idx ** diff_order
+        return fact * f(idx * coord)
+
+    def colloc_grid(self, order):
+        if self._parity == self.PARITY_NONE:
+            return (np.array(range(0, order)) + 1) * 2 * np.pi / (order + 1)
+        else:
+            return (np.array(range(0, order)) + 1) * np.pi / (2 * (order + 1))