[Refactor] Move feature extractor methods to utils

resspect/__init__.py

@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .bazin import *
 from .build_plasticc_canonical import *
 from .build_plasticc_metadata import *
 from .build_snpcc_canonical import *
@@ -41,16 +40,13 @@
 from .time_domain_loop import *
 from .batch_functions import *
 from .query_budget_strategies import *
-from .bump import *
 
 __all__ = ['accuracy',
            'assign_cosmo',
-           'bazin',
            'build_canonical',
            'build_plasticc_canonical',
            'build_plasticc_metadata',
            'build_snpcc_canonical',
-           'bump',
            'calculate_SNR',
            'Canonical',
            'CanonicalPLAsTiCC',

resspect/feature_extractors/bazin.py

@@ -5,9 +5,9 @@
 import matplotlib.pylab as plt
 import numpy as np
 
-from resspect import bazin
-from resspect.bazin import fit_scipy
 from resspect.feature_extractors.light_curve import LightCurve
+from resspect.utils.bazin_utils import bazin
+from resspect.utils.bazin_utils import fit_scipy
 
 __all__ = ['BazinFeatureExtractor']
 

resspect/feature_extractors/bump.py

@@ -5,9 +5,9 @@
 import matplotlib.pylab as plt
 import numpy as np
 
-from resspect import bump
-from resspect.bump import fit_bump
 from resspect.feature_extractors.light_curve import LightCurve
+from resspect.utils.bump_utils import bump
+from resspect.utils.bump_utils import fit_bump
 
 
 class BumpFeatureExtractor(LightCurve):

resspect/tests/test_bazin.py

@@ -16,7 +16,7 @@ def test_bazin():
     Test the Bazin function evaluation.
     """
 
-    from resspect import bazin
+    from resspect.utils.bazin_utils import bazin
 
     time = 3
     a = 1
@@ -35,7 +35,8 @@ def test_errfunc():
     Test the error between calculates and observed error.
     """
 
-    from resspect import bazin, errfunc
+    from resspect.utils.bazin_utils import bazin
+    from resspect.utils.bazin_utils import errfunc
 
     # input for bazin
     time = np.arange(0, 50, 3.5)

resspect/tests/test_bump.py

@@ -16,7 +16,7 @@ def test_bump():
     Test the Bump function evaluation.
     """
 
-    from resspect import bump
+    from resspect.utils.bump_utils import bump
 
     time = np.array([0])
     p1 = 0.225

resspect/bazin.py → resspect/utils/bazin_utils.py

@@ -52,6 +52,7 @@ def bazin(time, a, b, t0, tfall, trise):
         X = np.exp(-(time - t0) / tfall) / (1 + np.exp(-(time - t0) / trise))
         return a * X + b
 
+
 def bazinr(time, a, b, t0, tfall, r):
     """
     A wrapper function for bazin() which replaces trise by r = tfall/trise.
@@ -78,14 +79,15 @@ def bazinr(time, a, b, t0, tfall, r):
 
     """
 
-    trise = tfall/r    
+    trise = tfall / r
     res = bazin(time, a, b, t0, tfall, trise)
-    
+
     if max(res) < 10e10:
         return res
     else:
         return np.array([item if item < 10e10 else 10e10 for item in res])
 
+
 def errfunc(params, time, flux, fluxerr):
     """
     Absolute difference between theoretical and measured flux.
@@ -133,51 +135,50 @@ def fit_scipy(time, flux, fluxerr):
     flux = np.asarray(flux)
     imax = flux.argmax()
     flux_max = flux[imax]
-    
+
     # Parameter bounds
     a_bounds = [1.e-3, 10e10]
     b_bounds = [-10e10, 10e10]
-    t0_bounds = [-0.5*time.max(), 1.5*time.max()]
+    t0_bounds = [-0.5 * time.max(), 1.5 * time.max()]
     tfall_bounds = [1.e-3, 10e10]
     r_bounds = [1, 10e10]
 
     # Parameter guess
-    a_guess = 2*flux_max
+    a_guess = 2 * flux_max
     b_guess = 0
     t0_guess = time[imax]
-    
-    tfall_guess = time[imax-2:imax+2].std()/2
+
+    tfall_guess = time[imax - 2:imax + 2].std() / 2
     if np.isnan(tfall_guess):
-        tfall_guess = time[imax-1:imax+1].std()/2
+        tfall_guess = time[imax - 1:imax + 1].std() / 2
         if np.isnan(tfall_guess):
-            tfall_guess=50
-    if tfall_guess<1:
-        tfall_guess=50
+            tfall_guess = 50
+    if tfall_guess < 1:
+        tfall_guess = 50
 
     r_guess = 2
 
     # Clip guesses to stay in bound
-    a_guess = np.clip(a=a_guess,a_min=a_bounds[0],a_max=a_bounds[1])
-    b_guess = np.clip(a=b_guess,a_min=b_bounds[0],a_max=b_bounds[1])
-    t0_guess = np.clip(a=t0_guess,a_min=t0_bounds[0],a_max=t0_bounds[1])
-    tfall_guess = np.clip(a=tfall_guess,a_min=tfall_bounds[0],a_max=tfall_bounds[1])
-    r_guess = np.clip(a=r_guess,a_min=r_bounds[0],a_max=r_bounds[1])
-
-
-    guess = [a_guess,b_guess,t0_guess,tfall_guess,r_guess]
+    a_guess = np.clip(a=a_guess, a_min=a_bounds[0], a_max=a_bounds[1])
+    b_guess = np.clip(a=b_guess, a_min=b_bounds[0], a_max=b_bounds[1])
+    t0_guess = np.clip(a=t0_guess, a_min=t0_bounds[0], a_max=t0_bounds[1])
+    tfall_guess = np.clip(a=tfall_guess, a_min=tfall_bounds[0], a_max=tfall_bounds[1])
+    r_guess = np.clip(a=r_guess, a_min=r_bounds[0], a_max=r_bounds[1])
 
+    guess = [a_guess, b_guess, t0_guess, tfall_guess, r_guess]
 
     bounds = [[a_bounds[0], b_bounds[0], t0_bounds[0], tfall_bounds[0], r_bounds[0]],
               [a_bounds[1], b_bounds[1], t0_bounds[1], tfall_bounds[1], r_bounds[1]]]
-    
-    result = least_squares(errfunc, guess, args=(time, flux, fluxerr), method='trf', loss='linear',bounds=bounds)
-    
-    a_fit,b_fit,t0_fit,tfall_fit,r_fit = result.x
-    trise_fit = tfall_fit/r_fit
-    final_result = np.array([a_fit,b_fit,t0_fit,tfall_fit,trise_fit])
-    
+
+    result = least_squares(errfunc, guess, args=(time, flux, fluxerr), method='trf', loss='linear', bounds=bounds)
+
+    a_fit, b_fit, t0_fit, tfall_fit, r_fit = result.x
+    trise_fit = tfall_fit / r_fit
+    final_result = np.array([a_fit, b_fit, t0_fit, tfall_fit, trise_fit])
+
     return final_result
 
+
 def main():
     return None
 

resspect/bump.py → resspect/utils/bump_utils.py

@@ -1,6 +1,6 @@
 """
 # Author: Etienne Russeil and Emille E. O. Ishida
-#    
+#
 # created on 2 July 2022
 #
 # Licensed GNU General Public License v3.0;
@@ -22,6 +22,7 @@
 
 __all__ = ['protected_exponent', 'protected_sig', 'bump', 'fit_bump']
 
+
 def protected_exponent(x):
     """
     Exponential function : cannot exceed e**10
@@ -34,38 +35,36 @@ def protected_sig(x):
     """
     Sigmoid function using the protected exponential function
     """
-    return 1/(1+protected_exponent(-x))
+    return 1 / (1 + protected_exponent(-x))
 
 
 def bump(x, p1, p2, p3):
     """ Parametric function, fit transient behavior
         Need to fit normalised light curves (divided by maximum flux)
-    
+
     Parameters
     ----------
-    x : np.array 
+    x : np.array
         Array of mjd translated to 0
     p1,p2,p3 : floats
         Parameters of the function
-        
+
     Returns
     -------
     np.array
         Fitted flux array
     """
-    
+
     # The function is by construction meant to fit light curve centered on 40
     x = x + 40
-
-    return protected_sig(p1*x + p2 - protected_exponent(p3*x))
 
+    return protected_sig(p1 * x + p2 - protected_exponent(p3 * x))
 
 
 def fit_bump(time, flux, fluxerr):
-
     """
     Find best-fit parameters using iminuit least squares.
-    
+
     Parameters
     ----------
     time : array_like
@@ -76,26 +75,26 @@ def fit_bump(time, flux, fluxerr):
         error in response variable (flux)
     Returns
     -------
-    output : np.ndarray of floats 
+    output : np.ndarray of floats
         Array is [p1, p2, p3, time_shift, max_flux]
-        
+
          p1, p2, p3 are best fit parameter values
          time_shift is time at maximum flux
          max_flux is the maximum flux
-         
+
     """
-    
+
     # Center the maxflux at 0
-    time_shift = -time[np.argmax(flux)] 
+    time_shift = -time[np.argmax(flux)]
     time = time + time_shift
-    
+
     # The function is by construction meant to fit light curve with flux normalised
     max_flux = np.max(flux)
     flux = flux / max_flux
     fluxerr = fluxerr / max_flux
-    
+
     # Initial guess of the fit
-    parameters_dict = {'p1':0.225, 'p2':-2.5, 'p3':0.038}
+    parameters_dict = {'p1': 0.225, 'p2': -2.5, 'p3': 0.038}
 
     least_squares = LeastSquares(time, flux, fluxerr, bump)
     fit = Minuit(least_squares, **parameters_dict)
@@ -106,10 +105,10 @@ def fit_bump(time, flux, fluxerr):
     parameters = []
     for fit_values in range(len(fit.values)):
         parameters.append(fit.values[fit_values])
-        
+
     parameters.append(time_shift)
     parameters.append(max_flux)
-    
+
     return parameters