basis: add more bases

1 files changed, 143 insertions, 0 deletions

diff --git a/basis.py b/basis.py
index 759ebfe..ea26172 100644
--- a/basis.py
+++ b/basis.py
@@ -50,6 +50,149 @@ class CosBasis(ExpansionBasis1D):
 
     def colloc_grid(self, order):
         if self._parity == self.PARITY_NONE:
+            return (np.array(range(0, order)) + 1) * np.pi / (order + 1)
+            return (np.array(range(0, order))) * np.pi / (order - 1)
+        else:
+            return (np.array(range(0, order)) + 1) * np.pi / (2 * (order + 1))
+
+class ChebBasis(ExpansionBasis1D):
+
+    _scale = None
+
+    def __init__(self, scale = 1.0):
+        self._scale = scale
+
+    def eval(self, idx, coord, diff_order = 0):
+        x = 2.0 * coord / self._scale - 1.0
+        t = np.arccos(x)
+        st = np.sin(t)
+
+        #idx += 1
+
+        if diff_order == 0:
+            return np.cos(idx * t);
+        elif diff_order == 1:
+            return (2.0 / self._scale) * np.where(np.fabs(np.fabs(x) - 1.0) < 1e-10, idx * idx, idx * np.sin(idx * t) / st)
+        elif diff_order == 2:
+            return ((2.0 / self._scale) ** 2) * np.where(np.fabs(np.fabs(x) - 1.0) < 1e-10, idx * idx * (idx * idx - 1.0) / 3., -(idx ** 2) * np.cos(idx * t) / (st * st) + idx * np.cos(t) * np.sin(idx * t) / (st * st * st))
+        else:
+            raise NotImplementedError
+
+    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))
+
+class ChebEvenBasis(ExpansionBasis1D):
+
+    _scale = None
+
+    def __init__(self, scale = 1.0):
+        self._scale = scale
+
+    def eval(self, idx, coord, diff_order = 0):
+        x = coord / self._scale
+        t = np.arccos(x)
+        st = np.sin(t)
+
+        idx *= 2
+
+        if diff_order == 0:
+            return np.cos(idx * t);
+        elif diff_order == 1:
+            return (1.0 / self._scale) * np.where(np.fabs(np.fabs(x) - 1.0) < 1e-10, idx * idx, idx * np.sin(idx * t) / st)
+        elif diff_order == 2:
+            return ((1.0 / self._scale) ** 2) * np.where(np.fabs(np.fabs(x) - 1.0) < 1e-10, idx * idx * (idx * idx - 1.0) / 3., -(idx ** 2) * np.cos(idx * t) / (st * st) + idx * np.cos(t) * np.sin(idx * t) / (st * st * st))
+        else:
+            raise NotImplementedError
+
+    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))
+
+
+class ChebTLBasis(ExpansionBasis1D):
+
+    _scale = None
+
+    def __init__(self, scale = 1.0):
+        self._scale = scale
+
+    def eval(self, idx, coord, diff_order = 0):
+        y   = np.abs(coord)
+        y12 = np.sqrt(y)
+        t = 2 * np.arctan2(np.sqrt(self._scale), y12)
+        #idx += 1
+
+        if diff_order == 0:
+            return np.cos(idx * t)
+        elif diff_order == 1:
+            return np.where(np.abs(y) < 1e-10,
+                    -2.0 * idx * idx * ((-1.0)**idx) / self._scale,
+                    idx * np.sqrt(self._scale) * np.sin(idx * t) / (y12 * (y12 * y12 + self._scale)))
+        elif diff_order == 2:
+            return np.where(np.abs(y) < 1e-10,
+                    4.0 * idx * idx * (idx * idx + 2) * ((-1.0) ** idx) / (3 * (self._scale**2)),
+                    -((self._scale * idx * idx) * np.cos(idx * t) / (y * ((y + self._scale)**2)) + np.sqrt(self._scale) * idx * np.sin(idx * t) / (y12 * ((y + self._scale)**2)) + np.sqrt(self._scale) * idx * np.sin(idx * t) / (2 * y * y12 * (y +
+                        self._scale))))
+        else:
+            raise NotImplementedError
+
+    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))
+
+class ChebSBBasis(ExpansionBasis1D):
+
+    _scale = None
+
+    def __init__(self, scale = 1.0):
+        self._scale = scale
+
+    def eval(self, idx, coord, diff_order = 0):
+        idx *= 2
+        val = np.arctan2(self._scale, np.abs(coord))
+
+        if diff_order == 0:
+            return np.sin((idx + 1) * val)
+        elif diff_order == 1:
+            return -self._scale * (idx + 1) * np.sign(coord) * np.cos((idx + 1) * val) / ((self._scale ** 2) + (coord ** 2))
+        elif diff_order == 2:
+            return self._scale * (idx + 1) * (2 * np.abs(coord) * np.cos((1 + idx) * val) - self._scale * (1 + idx) * np.sin((1 + idx) * val)) / (((self._scale ** 2) + (coord ** 2)) ** 2)
+        else:
+            raise NotImplementedError
+
+    def colloc_grid(self, order):
+        idx = np.array(range(0, order))
+        t = (idx + 2) * np.pi / (2 * order + 2)
+        return self._scale / np.tan(t)
+
+class ChebTBBasis(ExpansionBasis1D):
+
+    _scale = None
+
+    def __init__(self, scale = 1.0):
+        self._scale = scale
+
+    def eval(self, idx, coord, diff_order = 0):
+        idx *= 2
+        val = np.arctan2(self._scale, np.abs(coord))
+
+        if diff_order == 0:
+            return np.cos(idx * val)
+        elif diff_order == 1:
+            return -self._scale * (idx + 1) * np.sign(coord) * np.cos((idx + 1) * val) / ((self._scale ** 2) + (coord ** 2))
+        elif diff_order == 2:
+            return self._scale * (idx + 1) * (2 * np.abs(coord) * np.cos((1 + idx) * val) - self._scale * (1 + idx) * np.sin((1 + idx) * val)) / (((self._scale ** 2) + (coord ** 2)) ** 2)
+        else:
+            raise NotImplementedError
+
+    def colloc_grid(self, order):
+        idx = np.array(range(0, order))
+        t = (idx + 2) * np.pi / (2 * order + 4)
+        return self._scale / np.tan(t)