amplitude base added

few/amplitude/interp2dcubicspline.py

@@ -2,7 +2,7 @@
 import os
 import h5py
 
-from few.utils.baseclasses import SchwarzschildEccentric
+from few.utils.baseclasses import SchwarzschildEccentric, AmplitudeBase
 
 from pyInterp2DAmplitude import pyAmplitudeGenerator
 
@@ -11,7 +11,7 @@
 dir_path = os.path.dirname(os.path.realpath(__file__))
 
 
-class Interp2DAmplitude(SchwarzschildEccentric):
+class Interp2DAmplitude(SchwarzschildEccentric, AmplitudeBase):
     """Calculate Teukolsky amplitudes by 2D Cubic Spline interpolation.
 
     Please see the documentations for
@@ -32,6 +32,7 @@ class Interp2DAmplitude(SchwarzschildEccentric):
     def __init__(self, **kwargs):
 
         SchwarzschildEccentric.__init__(self, **kwargs)
+        AmplitudeBase.__init__(self, **kwargs)
 
         few_dir = dir_path + "/../../"
 
@@ -45,7 +46,7 @@ def __init__(self, **kwargs):
 
         self.amplitude_generator = pyAmplitudeGenerator(self.lmax, self.nmax, few_dir)
 
-    def __call__(self, p, e, *args, specific_modes=None, **kwargs):
+    def get_amplitudes(self, p, e, *args, specific_modes=None, **kwargs):
         """Calculate Teukolsky amplitudes for Schwarzschild eccentric.
 
         This function takes the inputs the trajectory in :math:`(p,e)` as arrays

few/amplitude/romannet.py

@@ -11,7 +11,7 @@
 
 import pymatmul_cpu
 
-from few.utils.baseclasses import SchwarzschildEccentric
+from few.utils.baseclasses import SchwarzschildEccentric, AmplitudeBase
 
 try:
     import cupy as xp
@@ -44,7 +44,7 @@
 """
 
 
-class RomanAmplitude(SchwarzschildEccentric):
+class RomanAmplitude(SchwarzschildEccentric, AmplitudeBase):
     """Calculate Teukolsky amplitudes with a ROMAN.
 
     ROMAN stands for reduced-order models with artificial neurons. Please see
@@ -81,6 +81,7 @@ class RomanAmplitude(SchwarzschildEccentric):
     def __init__(self, max_input_len=1000, **kwargs):
 
         SchwarzschildEccentric.__init__(self, **kwargs)
+        AmplitudeBase.__init__(self, **kwargs)
 
         self.few_dir = dir_path + "/../../"
 
@@ -177,7 +178,7 @@ def _p_to_y(self, p, e):
 
         return self.xp.log(-(21 / 10) - 2 * e + p)
 
-    def __call__(self, p, e, *args, specific_modes=None, **kwargs):
+    def get_amplitudes(self, p, e, *args, specific_modes=None, **kwargs):
         """Calculate Teukolsky amplitudes for Schwarzschild eccentric.
 
         This function takes the inputs the trajectory in :math:`(p,e)` as arrays

few/utils/baseclasses.py

@@ -296,7 +296,7 @@ def get_inspiral(self, *args, **kwargs):
         returns:
             2D double np.ndarray: t, p, e, Phi_phi, Phi_r, flux with shape:
                 (params, traj length).
-                
+
         raises:
             NotImplementedError: The child class does not have this method.
 
@@ -501,3 +501,60 @@ def __call__(self, t, teuk_modes, ylms, dt, T, *args, **kwargs):
         self.sum(*args_in)
 
         return self.waveform
+
+
+class AmplitudeBase(ABC):
+    """Base class used for amplitude modules.
+
+    This class provides a common flexible interface to various amplitude
+    implementations. Specific arguments to each amplitude module can be found
+    with each associated module discussed below.
+
+    args:
+        pad_output (bool, optional): Add zero padding to the waveform for time
+            between plunge and observation time. Default is False.
+
+    """
+
+    def __init__(self, **kwargs):
+        pass
+
+    @classmethod
+    def get_amplitudes(self, *args, **kwargs):
+        """Amplitude Generator
+
+        @classmethod that requires a child class to have a get_amplitudes method.
+
+        raises:
+            NotImplementedError: The child class does not have this method.
+
+        """
+        raise NotImplementedError
+
+    def __call__(self, p, e, *args, specific_modes=None, **kwargs):
+        """Common call for Teukolsky amplitudes
+
+        This function takes the inputs the trajectory in :math:`(p,e)` as arrays
+        and returns the complex amplitude of all modes to adiabatic order at
+        each step of the trajectory.
+
+        args:
+            p (1D double numpy.ndarray): Array containing the trajectory for values of
+                the semi-latus rectum.
+            e (1D double numpy.ndarray): Array containing the trajectory for values of
+                the eccentricity.
+            *args (tuple, placeholder): Added to create future flexibility when calling different
+                amplitude modules.
+            specific_modes (list, optional): List of tuples for (l, m, n) values
+                desired modes. Default is None.
+            **kwargs (dict, placeholder): Added to create flexibility when calling different
+                amplitude modules. It is not used.
+
+        returns:
+            2D array (double): If specific_modes is None, Teukolsky modes in shape (number of trajectory points, number of modes)
+            dict: Dictionary with requested modes.
+
+
+        """
+
+        return self.get_amplitudes(p, e, *args, specific_modes=specific_modes, **kwargs)

few/waveform.py

@@ -44,7 +44,7 @@
 # TODO: shared memory based on CUDA_ARCH
 # TODO: deal with file locations and removing files from git history
 # TODO: add tutorials to documentation
-# TODO: Amplitude and general waveform base class
+# TODO: general waveform base class
 
 # Shorter Term
 # TODO: Zenodo versioning