Repository: COINtoolbox/RESSPECT (language: python)

Issue:
[In Progress] Modify LightCurve class to allow creation of feature extractor classes
New feature extractor classes can be added using `LightCurve` base class 
Feature extractor class should implement necessary methods. Please check `BazinFeatureExtractor` or `BumpFeatureExtractor` for reference
```
class NewFeatureExtractor(LightCurve):
    def __init__(self):
        super().__init__()

    def evaluate(self):
        """
        Evaluate the Bazin function given parameter values.
        """

    def fit(self):
        """
        Extracts features for one filter.
        """

    def fit_all(self):
        """
        Performs fit for all filters independently and concatenate results.
        """

    def plot_fit(self):
        """
        Plots data with fitted function.
        """
```
Todo:
- [ ] Bump extractor headers are missing
- [ ] Verify the implementation with different datasets 



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 254bfc08..b1fc6b3f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -19,4 +19,7 @@ resspect/.ipynb_checkpoints/
 __pycache__/
 
 # PyCharm Files
-.idea
\ No newline at end of file
+.idea
+
+auxiliary_files/cosmo
+auxiliary_files/cosmo.hpp
diff --git a/resspect/build_snpcc_canonical.py b/resspect/build_snpcc_canonical.py
index d2713474..a08ab1ba 100644
--- a/resspect/build_snpcc_canonical.py
+++ b/resspect/build_snpcc_canonical.py
@@ -25,7 +25,7 @@
 from sklearn.neighbors import NearestNeighbors, KernelDensity
 
 from resspect.database import DataBase
-from resspect.fit_lightcurves import LightCurve
+from resspect.feature_extractors.light_curve import LightCurve
 from resspect.lightcurves_utils import SNPCC_CANONICAL_FEATURES
 from resspect.lightcurves_utils import SNPCC_META_HEADER
 
diff --git a/resspect/database.py b/resspect/database.py
index 95109923..4435bdf8 100644
--- a/resspect/database.py
+++ b/resspect/database.py
@@ -16,21 +16,27 @@
 # limitations under the License.
 
 import io
-import numpy as np
 import os
 import pandas as pd
 import tarfile
 
 from resspect.classifiers import *
-
+from resspect.feature_extractors.bazin import BazinFeatureExtractor
+from resspect.feature_extractors.bump import BumpFeatureExtractor
 from resspect.query_strategies import *
 from resspect.query_budget_strategies import *
-from resspect.metrics import efficiency, purity, fom, accuracy, get_snpcc_metric
+from resspect.metrics import get_snpcc_metric
 
 
 __all__ = ['DataBase']
 
 
+FEATURE_EXTRACTOR_MAPPING = {
+    "bazin": BazinFeatureExtractor,
+    "bump": BumpFeatureExtractor
+}
+
+
 class DataBase:
     """DataBase object, upon which the active learning loop is performed.
 
@@ -142,7 +148,7 @@ class DataBase:
 
     Initiate the DataBase object and load the data.
     >>> data = DataBase()
-    >>> data.load_extracted_features(path_to_bazin_file, method='Bazin')
+    >>> data.load_features(path_to_bazin_file, method='bazin')
 
     Separate training and test samples and classify
 
@@ -212,8 +218,8 @@ def __init__(self):
         self.validation_metadata = pd.DataFrame()
         self.validation_prob = np.array([])
 
-    def load_extracted_features(self, path_to_features_file: str, screen=False,
-                            survey='DES', sample=None, function='bazin'):
+    def load_features(self, path_to_features_file: str, screen=False,
+                      survey='DES', sample=None, feature_extractor='bazin'):
         """Load features from file.
 
         Populate properties: features, feature_names, metadata
@@ -233,7 +239,7 @@ def load_extracted_features(self, path_to_features_file: str, screen=False,
             If None, sample is given by a column within the given file.
             else, read independent files for 'train' and 'test'.
             Default is None.
-        function: str (optional)
+        feature_extractor: str (optional)
             Function used for feature extraction. Options are "bazin" or 
             "bump". Default is "bump".
         """
@@ -257,12 +263,12 @@ def load_extracted_features(self, path_to_features_file: str, screen=False,
 
         # list of features to use
         if survey == 'DES':
-            if function == "bazin":
+            if feature_extractor == "bazin":
                 self.features_names = ['gA', 'gB', 'gt0', 'gtfall', 'gtrise', 'rA',
                                        'rB', 'rt0', 'rtfall', 'rtrise', 'iA', 'iB',
                                        'it0', 'itfall', 'itrise', 'zA', 'zB', 'zt0',
                                        'ztfall', 'ztrise']
-            elif function == 'bump':
+            elif feature_extractor == 'bump':
                 self.features_names = ['gp1', 'gp2', 'gp3', 'gmax_flux', 
                                        'rp1', 'rp2', 'rp3', 'rmax_flux', 
                                        'ip1', 'ip2', 'ip3', 'imax_flux', 
@@ -343,7 +349,6 @@ def load_extracted_features(self, path_to_features_file: str, screen=False,
                 print('Loaded ', self.pool_metadata.shape[0], ' ' + \
                       sample +  ' samples!')
 
-
     def load_photometry_features(self, path_to_photometry_file: str,
                                  screen=False, sample=None):
         """Load photometry features from file.
@@ -413,8 +418,8 @@ def load_photometry_features(self, path_to_photometry_file: str,
                 print('\n Loaded ', self.test_metadata.shape[0],
                       ' samples! \n')
 
-    def load_features(self, path_to_file: str, method='Bazin', screen=False,
-                      survey='DES', sample=None ):
+    def load_features(self, path_to_file: str, feature_extractor: str ='bazin',
+                      screen=False, survey='DES', sample=None ):
         """Load features according to the chosen feature extraction method.
 
         Populates properties: data, features, feature_list, header
@@ -424,10 +429,10 @@ def load_features(self, path_to_file: str, method='Bazin', screen=False,
         ----------
         path_to_file: str
             Complete path to features file.
-        method: str (optional)
+        feature_extractor: str (optional)
             Feature extraction method. The current implementation only
-            accepts method=='Bazin', 'bump' or 'photometry'.
-            Default is 'Bazin'.
+            accepts =='bazin', 'bump' or 'photometry'.
+            Default is 'bazin'.
         screen: bool (optional)
             If True, print on screen number of light curves processed.
             Default is False.
@@ -440,16 +445,15 @@ def load_features(self, path_to_file: str, method='Bazin', screen=False,
             else, read independent files for 'train' and 'test'.
             Default is None.
         """
-
-        if method == 'Bazin' or method == 'bump':
-            self.load_extracted_features(path_to_file, screen=screen,
-                               survey=survey, sample=sample, function=method)
-
-        elif method == 'photometry':
+        if feature_extractor == "photometry":
             self.load_photometry_features(path_to_file, screen=screen,
                                           survey=survey, sample=sample)
+        elif feature_extractor in FEATURE_EXTRACTOR_MAPPING:
+            self.load_features(
+                path_to_file, screen=screen, survey=survey,
+                sample=sample, feature_extractor=feature_extractor)
         else:
-            raise ValueError('Only Bazin, bump or photometry features are implemented!'
+            raise ValueError('Only bazin, bump or photometry features are implemented!'
                              '\n Feel free to add other options.')
 
     def load_plasticc_mjd(self, path_to_data_dir):
diff --git a/resspect/feature_extractors/__init__.py b/resspect/feature_extractors/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/resspect/feature_extractors/bazin.py b/resspect/feature_extractors/bazin.py
new file mode 100644
index 00000000..b47709d2
--- /dev/null
+++ b/resspect/feature_extractors/bazin.py
@@ -0,0 +1,223 @@
+"""
+    Bazin feature extractor
+"""
+
+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
+
+__all__ = ['BazinFeatureExtractor']
+
+
+class BazinFeatureExtractor(LightCurve):
+    def __init__(self):
+        super().__init__()
+        self.features_names = ['a', 'b', 't0', 'tfall', 'trise']
+
+    def evaluate(self, time: np.array) -> dict:
+        """
+        Evaluate the Bazin function given parameter values.
+
+        Parameters
+        ----------
+        time: np.array or list
+          Time since first light curve observation.
+
+        Returns
+        -------
+        dict
+          Value of the Bazin flux in each required time per filter.
+        """
+        # store flux values and starting points
+        flux = {}
+
+        for k in range(len(self.filters)):
+            # store flux values per filter
+            flux[self.filters[k]] = []
+
+            # check if Bazin features exist
+            if 'None' not in self.features[k * 5: (k + 1) * 5]:
+                for item in time:
+                    flux[self.filters[k]].append(
+                        bazin(item, self.features[0 + k * 5],
+                              self.features[1 + k * 5],
+                              self.features[2 + k * 5],
+                              self.features[3 + k * 5],
+                              self.features[4 + k * 5]))
+            else:
+                flux[self.filters[k]].append(None)
+        return flux
+
+    def fit(self, band: str) -> np.ndarray:
+        """
+        Extracts Bazin features for one filter.
+
+        Parameters
+        ----------
+        band: str
+            Choice of broad band filter
+
+        Returns
+        -------
+        bazin_param: np.ndarray
+            Best fit parameters for the Bazin function:
+            [a, b, t0, tfall, trise].
+        """
+
+        # build filter flag
+        band_indices = self.photometry['band'] == band
+        if not sum(band_indices) > (len(self.features_names) - 1):
+            return np.array([])
+
+        # get info for this filter
+        time = self.photometry['mjd'].values[band_indices]
+        flux = self.photometry['flux'].values[band_indices]
+        flux_error = self.photometry['fluxerr'].values[band_indices]
+
+        # fit Bazin function
+        bazin_param = fit_scipy(time - time[0], flux, flux_error)
+        return bazin_param
+
+    def fit_all(self):
+        """
+        Performs bazin fit for all filters independently and concatenate results.
+        Populates the attributes: bazinfeatures.
+        """
+        default_bazinfeatures = ['None'] * len(self.features_names)
+
+        if self.photometry.shape[0] < 1:
+            self.features = ['None'] * len(
+                self.features_names) * len(self.filters)
+
+        elif 'None' not in self.features:
+            self.features = []
+            for each_band in self.filters:
+                best_fit = self.fit(band=each_band)
+                if (best_fit.size > 0) and (not np.isnan(np.sum(best_fit))):
+                    self.features.extend(best_fit.tolist())
+                else:
+                    self.features.extend(default_bazinfeatures)
+        else:
+            self.features.extend(default_bazinfeatures)
+
+    def plot_fit(self, save=True, show=False, output_file=' ',
+                 figscale=1, extrapolate=False,
+                 time_flux_pred=None, unit='flux'):
+        """
+        Plots data with bazin fitted function.
+
+        Parameters
+        ----------
+        figscale: float (optional)
+            Allow to control the size of the figure.
+        extrapolate: bool (optional)
+            If True, also plot the estimated flux values.
+            Default is False.
+        output_file: str (optional)
+            Name of file to store the plot.
+        save: bool (optional)
+             Save figure to file. Default is True.
+        show: bool (optinal)
+             Display plot in windown. Default is False.
+        time_flux_pred: list (optional)
+            Time since first observation where flux is to be
+            estimated. It is only used if "extrapolate == True".
+            Default is None.
+        unit: str (optional)
+            Unit for plot. Options are 'flux' or 'mag'.
+            Use zero point from SNANA for flux-to-mag conversion
+            ==> mag = 2.5 * (11 - np.log10(flux)).
+            Default is 'flux'.
+        """
+
+        # number of columns in the plot
+        ncols = len(self.filters) / 2 + len(self.filters) % 2
+        fsize = (figscale * 5 * ncols, figscale * 10)
+
+        plt.figure(figsize=fsize)
+
+        for i in range(len(self.filters)):
+            plt.subplot(2, int(ncols), i + 1)
+            plt.title('Filter: ' + self.filters[i])
+
+            # filter flag
+            filter_flag = self.photometry['band'] == self.filters[i]
+            x = self.photometry['mjd'][filter_flag].values
+            y = self.photometry['flux'][filter_flag].values
+            yerr = self.photometry['fluxerr'][filter_flag].values
+
+            # check Bazin fit convergence
+            if 'None' in self.features[i * 5: (i + 1) * 5]:
+                plot_fit = False
+            else:
+                plot_fit = True
+
+            # shift to avoid large numbers in x-axis
+            time = x - min(x)
+
+            if plot_fit:
+                xaxis = np.linspace(0, max(time), 500)[:, np.newaxis]
+                fitted_flux = self.evaluate(xaxis)
+                if unit == 'flux':
+                    plt.plot(xaxis, fitted_flux[self.filters[i]], color='red',
+                             lw=1.5, label='Bazin fit')
+                elif unit == 'mag':
+                    mag = self.conv_flux_mag(fitted_flux[self.filters[i]])
+                    mag_flag = mag < 50
+                    plt.plot(xaxis[mag_flag], mag[mag_flag], color='red',
+                             lw=1.5)
+                else:
+                    raise ValueError('Unit can only be "flux" or "mag".')
+
+                if extrapolate:
+                    xaxis_extrap = list(xaxis.flatten()) + list(time_flux_pred)
+                    xaxis_extrap = np.sort(np.array(xaxis_extrap))
+                    ext_flux = self.evaluate(xaxis_extrap)
+                    if unit == 'flux':
+                        plt.plot(xaxis_extrap, ext_flux[self.filters[i]],
+                                 color='red', lw=1.5, ls='--',
+                                 label='Bazin extrap')
+                    elif unit == 'mag':
+                        ext_mag = self.conv_flux_mag(ext_flux[self.filters[i]])
+                        ext_mag_flag = ext_mag < 50
+                        plt.plot(xaxis_extrap[ext_mag_flag],
+                                 ext_mag[ext_mag_flag],
+                                 color='red', lw=1.5, ls='--')
+
+            if unit == 'flux':
+                plt.errorbar(time, y, yerr=yerr, color='blue', fmt='o',
+                             label='obs')
+                plt.ylabel('FLUXCAL')
+            elif unit == 'mag':
+                mag_obs = self.conv_flux_mag(y)
+                mag_obs_flag = mag_obs < 50
+                time_mag = time[mag_obs_flag]
+
+                plt.scatter(time_mag, mag_obs[mag_obs_flag], color='blue',
+                            label='calc mag', marker='s')
+
+                # if MAG is provided in the table, also plot it
+                # this allows checking the flux mag conversion
+                if 'MAG' in self.photometry.keys():
+                    mag_flag = self.photometry['MAG'].values < 50
+                    mag_ff = np.logical_and(filter_flag, mag_flag)
+                    mag_table = self.photometry['MAG'][mag_ff].values
+                    mjd_table = self.photometry['mjd'][mag_ff].values - min(x)
+
+                    plt.scatter(mjd_table, mag_table, color='black',
+                                label='table mag', marker='x')
+
+                ax = plt.gca()
+                ax.set_ylim(ax.get_ylim()[::-1])
+                plt.ylabel('mag')
+
+            plt.xlabel('days since first observation')
+            plt.tight_layout()
+
+        if save:
+            plt.savefig(output_file)
+        if show:
+            plt.show()
diff --git a/resspect/feature_extractors/bump.py b/resspect/feature_extractors/bump.py
new file mode 100644
index 00000000..ebb4dfde
--- /dev/null
+++ b/resspect/feature_extractors/bump.py
@@ -0,0 +1,218 @@
+"""
+    Bump feature extractor
+"""
+
+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
+
+
+class BumpFeatureExtractor(LightCurve):
+    def __init__(self):
+        super().__init__()
+        self.features_names = ['p1', 'p2', 'p3', 'time_shift', 'max_flux']
+
+    def evaluate(self, time: np.array) -> dict:
+        """
+        Evaluate the Bump function given parameter values.
+
+        Parameters
+        ----------
+        time: np.array or list
+            Time since first light curve observation.
+
+        Returns
+        -------
+        dict
+            Value of the Bump flux in each required time per filter.
+        """
+        # store flux values and starting points
+        flux = {}
+        for k in range(len(self.filters)):
+            # store flux values per filter
+            flux[self.filters[k]] = []
+
+            # check if Bump features exist
+            if 'None' not in self.features[k * 5 : (k + 1) * 5]:
+                for item in time:
+                    flux[self.filters[k]].append(
+                        bump(item, self.features[0 + k * 5],
+                             self.features[1 + k * 5],
+                             self.features[2 + k * 5]))
+            else:
+                flux[self.filters[k]].append(None)
+        return flux
+
+    def fit(self, band: str) -> np.ndarray:
+        """
+        Extract Bump features for one filter.
+
+        Parameters
+        ----------
+        band: str
+            Choice of broad band filter
+
+        Returns
+        -------
+        bump_param: np.ndarray
+            Best fit parameters for the Bazin function:
+            [p1, p2, p3, time_shift, max_flux].
+        """
+
+        # build filter flag
+        band_indices = self.photometry['band'] == band
+        if not sum(band_indices) > (len(self.features_names) - 2):
+            return np.array([])
+
+        # get info for this filter
+        time = self.photometry['mjd'].values[band_indices]
+        flux = self.photometry['flux'].values[band_indices]
+        flux_error = self.photometry['fluxerr'].values[band_indices]
+
+        # fit Bump function
+        bump_param = fit_bump(time, flux, flux_error)
+        return bump_param
+
+    def fit_all(self):
+        """
+        Perform Bump fit for all filters independently and concatenate results.
+        Populates the attributes: bump_features.
+        """
+        default_bump_features = ['None'] * len(self.features_names)
+
+        if self.photometry.shape[0] < 1:
+            self.features = ['None'] * len(self.features_names) * len(self.filters)
+
+        elif 'None' not in self.features:
+            self.features = []
+            for each_band in self.filters:
+                best_fit = self.fit(band=each_band)
+                if (len(best_fit) > 0) and (not np.isnan(np.sum(best_fit))):
+                    self.features.extend(best_fit)
+                else:
+                    self.features.extend(default_bump_features)
+        else:
+            self.features.extend(default_bump_features)
+
+    def plot_fit(self, save=True, show=False, output_file=' ',
+                 figscale=1, extrapolate=False,
+                 time_flux_pred=None, unit='flux'):
+        """
+        Plot data and Bump fitted function.
+
+        Parameters
+        ----------
+        figscale: float (optional)
+            Allow to control the size of the figure.
+        extrapolate: bool (optional)
+            If True, also plot the estimated flux values.
+            Default is False.
+        output_file: str (optional)
+            Name of file to store the plot.
+        save: bool (optional)
+             Save figure to file. Default is True.
+        show: bool (optinal)
+             Display plot in windown. Default is False.
+        time_flux_pred: list (optional)
+            Time after last observation where flux is to be
+            estimated. It is only used if "extrapolate == True".
+            Default is None.
+        unit: str (optional)
+            Unit for plot. Options are 'flux' or 'mag'.
+            Use zero point from SNANA for flux-to-mag conversion
+            ==> mag = 2.5 * (11 - np.log10(flux)).
+            Default is 'flux'.
+        """
+
+        # number of columns in the plot
+        ncols = len(self.filters) / 2 + len(self.filters) % 2
+        fsize = (figscale * 5 * ncols , figscale * 10)
+
+        plt.figure(figsize=fsize)
+
+        for i in range(len(self.filters)):
+            plt.subplot(2, int(ncols), i + 1)
+            plt.title('Filter: ' + self.filters[i])
+
+            # filter flag
+            filter_flag = self.photometry['band'] == self.filters[i]
+            x = self.photometry['mjd'][filter_flag].values
+            y = self.photometry['flux'][filter_flag].values
+            yerr = self.photometry['fluxerr'][filter_flag].values
+
+            # check Bump fit convergence
+            if 'None' in self.features[i * 5 : (i + 1) * 5]:
+                plot_fit = False
+            else:
+                plot_fit = True
+
+            # shift to avoid large numbers in x-axis
+            time = x + self.features[i * 5 + 3]
+
+            if plot_fit:
+                xaxis = np.linspace(min(time), max(time), 500)[:, np.newaxis]
+                fitted_flux = np.array(self.evaluate(xaxis)[self.filters[i]]) * self.features[i * 5 + 4]
+                if unit == 'flux':
+                    plt.plot(xaxis, fitted_flux, color='red',
+                             lw=1.5, label='Bump fit')
+                elif unit == 'mag':
+                    mag = self.conv_flux_mag(fitted_flux[self.filters[i]])
+                    mag_flag = mag < 50
+                    plt.plot(xaxis[mag_flag], mag[mag_flag], color='red',
+                             lw=1.5)
+                else:
+                    raise ValueError('Unit can only be "flux" or "mag".')
+
+                if extrapolate:
+                    xaxis_extrap = list(xaxis.flatten()) + [time_flux_pred + max(time)]
+                    xaxis_extrap = np.sort(np.array(xaxis_extrap))
+                    ext_flux = np.array(self.evaluate(xaxis_extrap)[self.filters[i]]) \
+                               * self.features[i * 5 + 4]
+                    if unit == 'flux':
+                        plt.plot(xaxis_extrap, ext_flux,
+                                 color='red', lw=1.5, ls='--',
+                                 label='Bump extrap')
+                    elif unit == 'mag':
+                        ext_mag = self.conv_flux_mag(ext_flux[self.filters[i]])
+                        ext_mag_flag = ext_mag < 50
+                        plt.plot(xaxis_extrap[ext_mag_flag],
+                                 ext_mag[ext_mag_flag],
+                                 color='red', lw=1.5, ls='--')
+
+            if unit == 'flux':
+                plt.errorbar(time, y, yerr=yerr, color='blue', fmt='o',
+                             label='obs')
+                plt.ylabel('FLUXCAL')
+            elif unit == 'mag':
+                mag_obs  = self.conv_flux_mag(y)
+                mag_obs_flag = mag_obs < 50
+                time_mag = time[mag_obs_flag]
+
+                plt.scatter(time_mag, mag_obs[mag_obs_flag], color='blue',
+                            label='calc mag', marker='s')
+
+                # if MAG is provided in the table, also plot it
+                # this allows checking the flux mag conversion
+                if 'MAG' in self.photometry.keys():
+                    mag_flag = self.photometry['MAG'].values < 50
+                    mag_ff = np.logical_and(filter_flag, mag_flag)
+                    mag_table = self.photometry['MAG'][mag_ff].values
+                    mjd_table = self.photometry['mjd'][mag_ff].values - min(x)
+
+                    plt.scatter(mjd_table, mag_table, color='black',
+                                label='table mag', marker='x')
+
+                ax = plt.gca()
+                ax.set_ylim(ax.get_ylim()[::-1])
+                plt.ylabel('mag')
+
+            plt.xlabel('days since first observation')
+            plt.tight_layout()
+
+        if save:
+            plt.savefig(output_file)
+        if show:
+            plt.show()
diff --git a/resspect/feature_extractors/light_curve.py b/resspect/feature_extractors/light_curve.py
new file mode 100644
index 00000000..0579b7a6
--- /dev/null
+++ b/resspect/feature_extractors/light_curve.py
@@ -0,0 +1,412 @@
+# Copyright 2020 resspect software
+# Author: Emille E. O. Ishida
+#
+# created on 14 April 2020
+#
+# Licensed GNU General Public License v3.0;
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.gnu.org/licenses/gpl-3.0.en.html
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from typing import Tuple
+import warnings
+
+import numpy as np
+import pandas as pd
+
+from resspect.exposure_time_calculator import ExpTimeCalc
+from resspect.lightcurves_utils import read_file
+from resspect.lightcurves_utils import load_snpcc_photometry_df
+from resspect.lightcurves_utils import get_photometry_with_id_name_and_snid
+from resspect.lightcurves_utils import read_plasticc_full_photometry_data
+from resspect.lightcurves_utils import load_plasticc_photometry_df
+from resspect.lightcurves_utils import read_resspect_full_photometry_data
+from resspect.lightcurves_utils import insert_band_column_to_resspect_df
+from resspect.lightcurves_utils import load_resspect_photometry_df
+from resspect.lightcurves_utils import get_snpcc_sntype
+
+
+warnings.filterwarnings("ignore", category=RuntimeWarning)
+logging.basicConfig(level=logging.INFO)
+
+__all__ = ['LightCurve']
+
+
+class LightCurve:
+    """ Light Curve object, holding meta and photometric data.
+
+    Attributes
+    ----------
+    features_names: list
+        List of names of the feature extraction parameters.
+    features: list
+        List with the 5 best-fit feature extraction parameters in all filters.
+        Concatenated from blue to red.
+    dataset_name: str
+        Name of the survey or data set being analyzed.
+    exp_time: dict
+        Exposure time required to take a spectra.
+        Keywords indicate telescope e.g.['4m', '8m'].
+    filters: list
+        List of broad band filters.
+    full_photometry: pd.DataFrame
+        Photometry for a set of light curves read from file.
+    id: int
+        SN identification number.
+    id_name:
+        Column name of object identifier.
+    last_mag: float
+        r-band magnitude of last observed epoch.
+    photometry: pd.DataFrame
+        Photometry information.
+        Minimum keys --> [mjd, band, flux, fluxerr].
+    redshift: float
+        Redshift
+    sample: str
+        Original sample to which this light curve is assigned.
+    sim_peakmag: np.array
+        Simulated peak magnitude in each filter.
+    sim_pkmjd: float
+        Simulated day of maximum, observer frame.
+    sncode: int
+        Number identifying the SN model used in the simulation.
+    sntype: str
+        General classification, possibilities are: Ia, II or Ibc.
+    unique_ids: str or array
+        List of unique ids available in the photometry file.
+        Only used for PLAsTiCC data.
+
+    Methods
+    -------
+    calc_exp_time(telescope_diam: float, SNR: float, telescope_name: str)
+        Calculates time required to take a spectra in the last obs epoch.
+    check_queryable(mjd: float, r_lim: float)
+        Check if this light can be queried in a given day.
+    conv_flux_mag(flux: np.array)
+        Convert positive flux into magnitude.
+    evaluate_bazin(param: list, time: np.array) -> np.array
+        Evaluate the Bazin function given parameter values.
+    load_snpcc_lc(path_to_data: str)
+        Reads header and photometric information for 1 light curve.
+    load_plasticc_lc(photo_file: str, snid: int)
+    	Load photometric information for 1 PLAsTiCC light curve.
+    load_resspect_lc(photo_file: str, snid: int)
+    	Load photometric information for 1 RESSPECT light curve.
+    plot_bump_fit(save: bool, show: bool, output_file: srt)
+        Plot photometric points and Bump fitted curve.
+
+    """
+
+    def __init__(self):
+        self.queryable = None
+        self.features = []
+        self.features_names = ['p1', 'p2', 'p3', 'time_shift', 'max_flux']
+        self.dataset_name = ' '
+        self.exp_time = {}
+        self.filters = []
+        self.full_photometry = pd.DataFrame()
+        self.id = 0
+        self.id_name = None
+        self.last_mag = None
+        self.photometry = pd.DataFrame()
+        self.redshift = 0
+        self.sample = ' '
+        self.sim_peakmag = []
+        self.sim_pkmjd = None
+        self.sncode = 0
+        self.sntype = ' '
+
+    def _get_snpcc_photometry_raw_and_header(
+            self, lc_data: np.ndarray,
+            sntype_test_value: str = "-9") -> Tuple[np.ndarray, list]:
+        """
+        Reads SNPCC photometry raw and header data
+
+        Parameters
+        ----------
+        lc_data
+            SNPCC light curve data
+        sntype_test_value
+            test sample SNTYPE value
+        """
+        photometry_raw = []
+        header = []
+        for each_row in lc_data:
+            name = each_row[0]
+            value = each_row[1]
+            if name == 'SNID:':
+                self.id = int(value)
+                self.id_name = 'SNID'
+            elif name == 'SNTYPE:':
+                self.sample = 'test' if value == sntype_test_value else 'train'
+            elif name == 'SIM_REDSHIFT:':
+                self.redshift = float(value)
+            elif name == 'SIM_NON1a:':
+                self.sncode = value
+                self.sntype = get_snpcc_sntype(int(value))
+            elif name == 'VARLIST:':
+                header = each_row[1:]
+            elif name == 'OBS:':
+                photometry_raw.append(np.array(each_row[1:]))
+            elif name == 'SIM_PEAKMAG:':
+                self.sim_peakmag = np.array(each_row[1:5]).astype(float)
+            elif name == 'SIM_PEAKMJD:':
+                self.sim_pkmjd = float(value)
+        return np.array(photometry_raw), header
+
+    def load_snpcc_lc(self, path_to_data: str):
+        """Reads one LC from SNPCC data.
+
+        Populates the attributes: dataset_name, id, sample, redshift, sncode,
+        sntype, photometry, sim_peakmag and sim_pkmjd.
+
+        Parameters
+        ---------
+        path_to_data: str
+            Path to text file with data from a single SN.
+        """
+        self.dataset_name = 'SNPCC'
+        self.filters = ['g', 'r', 'i', 'z']
+
+        lc_data = np.array(read_file(path_to_data), dtype=object)
+        photometry_raw, header = self._get_snpcc_photometry_raw_and_header(
+            lc_data)
+
+        if photometry_raw.size > 0:
+            self.photometry = load_snpcc_photometry_df(photometry_raw, header)
+
+    def load_resspect_lc(self, photo_file: str, snid: int):
+        """
+        Return 1 light curve from RESSPECT simulations.
+
+        Parameters
+        ----------
+        photo_file: str
+            Complete path to light curves file.
+        snid: int
+            Identification number for the desired light curve.
+        """
+
+        self.dataset_name = 'RESSPECT'
+        self.filters = ['u', 'g', 'r', 'i', 'z', 'Y']
+        self.id = snid
+
+        if self.full_photometry.empty:
+            _, self.full_photometry = read_resspect_full_photometry_data(
+                photo_file)
+        id_names_list = ['SNID', 'snid', 'objid', 'id']
+        filtered_photometry, self.id_name = (
+            get_photometry_with_id_name_and_snid(
+                self.full_photometry, id_names_list, snid))
+
+        if not filtered_photometry.empty:
+            filtered_photometry = insert_band_column_to_resspect_df(
+                filtered_photometry, self.filters)
+            self.photometry = load_resspect_photometry_df(filtered_photometry)
+
+    def load_plasticc_lc(self, photo_file: str, snid: int):
+        """
+        Return 1 light curve from PLAsTiCC simulations.
+
+        Parameters
+        ----------
+        photo_file: str
+            Complete path to light curve file.
+        snid: int
+            Identification number for the desired light curve.
+        """
+        self.dataset_name = 'PLAsTiCC'
+        self.filters = ['u', 'g', 'r', 'i', 'z', 'Y']
+        self.id = snid
+
+        if self.full_photometry.empty:
+            self.full_photometry = read_plasticc_full_photometry_data(
+                photo_file)
+
+        id_names_list = ['object_id', 'SNID', 'snid']
+
+        filtered_photometry, self.id_name = (
+            get_photometry_with_id_name_and_snid(
+                self.full_photometry, id_names_list, snid))
+
+        filter_mapping_dict = {
+            0: 'u', 1: 'g', 2: 'r', 3: 'i', 4: 'z', 5: 'Y'
+        }
+
+        if not filtered_photometry.empty:
+            self.photometry = load_plasticc_photometry_df(
+                filtered_photometry, filter_mapping_dict)
+
+    @staticmethod
+    def conv_flux_mag(flux, zpt: float = 27.5):
+        """Convert FLUXCAL to magnitudes.
+
+        Parameters
+        ----------
+        flux: list or np.array
+            Values of flux to be converted into mag.
+        zpt: float (optional)
+            Zero point. Default is 27.5 (from SNANA).
+
+        Returns
+        -------
+        mag: list or np.array
+            Magnitude values. If flux < 1e-5 returns 9999.
+        """
+
+        if not isinstance(flux[0], float):
+            flux = np.array([item[0] for item in flux])
+        mag = [zpt - 2.5 * np.log10(f) if f > 1e-5 else 9999 for f in flux]
+        return np.array(mag)
+
+    def check_queryable(self, mjd: float, filter_lim: float, criteria: int =1,
+                        days_since_last_obs=2, feature_method='Bazin',
+                        filter_cut='r'):
+        """Check if this object can be queried in a given day.
+
+        Parameters
+        ----------
+        mjd: float
+            MJD where the query will take place.
+        filter_lim: float
+            Magnitude limit below which query is possible.
+        criteria: int [1, 2 or 3] (optional)
+            Criteria to determine if an obj is queryable.
+            1 -> Cut on last measured photometric point.
+            2 -> if last obs was further than days_since_last_obs,
+                 use Bazin estimate for today. Otherwise, use
+                 the last observed point.
+            Default is 1.
+        days_since_last_obs: int (optional)
+            If there is an observation within these days, use the
+            measured value, otherwise estimate current mag.
+            Only used if "criteria == 2". Default is 2.
+        feature_method: str (optional)
+            Feature extraction method. Only 'Bazin' is implemented.
+            Default is 'Bazin'.
+        filter_cut: str (optional)
+            Band in which cut is applied. Default is 'r'.
+
+        Returns
+        -------
+        bool
+            True if current magnitude lower than "filter_lim".
+        """
+
+        # create photo flag
+        photo_flag = self.photometry['mjd'].values <= mjd
+        rband_flag = self.photometry['band'].values == filter_cut
+        surv_flag = np.logical_and(photo_flag, rband_flag)
+
+        if criteria == 1 and sum(surv_flag):
+            if 'MAG' in self.photometry.keys():
+                # check surviving photometry
+                self.last_mag = self.photometry['MAG'].values[surv_flag][-1]
+
+            else:
+                surv_flux = self.photometry['flux'].values[surv_flag]
+                self.last_mag = self.conv_flux_mag([surv_flux[-1]])[0]
+
+        elif criteria == 2 and sum(surv_flag):
+            # check if there is an observation recently
+            surv_mjd = self.photometry['mjd'].values[surv_flag]
+            gap = mjd - surv_mjd[-1]
+
+            if gap <= days_since_last_obs:
+                if 'MAG' in self.photometry.keys():
+                    # check surviving photometry
+                    self.last_mag = \
+                        self.photometry['MAG'].values[surv_flag][-1]
+
+                else:
+                    surv_flux = self.photometry['flux'].values[surv_flag]
+                    self.last_mag = self.conv_flux_mag([surv_flux[-1]])[0]
+
+            elif feature_method == 'Bazin':
+                # get first day of observation in this filter
+                mjd_min = min(self.photometry['mjd'].values[surv_flag])
+
+                # estimate flux based on Bazin function
+                fitted_flux = self.evaluate([mjd - mjd_min])[filter_cut][0]
+                self.last_mag = self.conv_flux_mag([fitted_flux])[0]
+
+            else:
+                raise ValueError('Only "Bazin" features are implemented!')
+
+        elif sum(surv_flag):
+            raise ValueError('Criteria needs to be "1" or "2". \n ' + \
+                             'See docstring for further info.')
+        elif not sum(surv_flag):
+            return False
+
+        if self.last_mag <= filter_lim:
+            return True
+        else:
+            return False
+
+    def calc_exp_time(self, telescope_diam: float, SNR: float,
+                      telescope_name: str, **kwargs):
+        """Calculates time required to take a spectra in the last obs epoch.
+
+        Populates attribute exp_time.
+
+        Parameters
+        ----------
+        SNR: float
+            Required SNR.
+        telescope_diam: float
+            Diameter of primary mirror for spectroscopic telescope in meters.
+        telescope_name: str
+            Identification for telescope.
+        kwargs: extra parameters
+            Any input required by ExpTimeCalc.findexptime function.
+
+        Returns
+        -------
+        exp_time: float
+            Exposure time require for taking a spectra in seconds.
+        """
+
+        # check if last magnitude was calculated
+        if self.last_mag is None:
+            raise ValueError('Magnitude at last epoch not calculated.\n' + \
+                             'Run the check_queryable function.')
+
+        etc = ExpTimeCalc()
+        etc.diameter = telescope_diam
+        exp_time = etc.findexptime(SNRin=SNR, mag=self.last_mag, **kwargs)
+
+        if 60 < exp_time < 7200 and self.last_mag < 30:
+            self.exp_time[telescope_name] = exp_time
+            return exp_time
+        else:
+            self.exp_time[telescope_name] = 9999
+            return 9999
+
+    def evaluate(self, time: np.ndarray) -> dict:
+        raise NotImplementedError
+
+    def fit(self, band: str) -> np.ndarray:
+        raise NotImplementedError
+
+    def fit_all(self):
+        raise NotImplementedError
+
+    def plot_fit(self):
+        raise NotImplementedError
+
+    def clear_data(self):
+        """ Reset to default values """
+        self.photometry = []
+        self.redshift = None
+        self.sntype = None
+        self.sncode = None
+        self.sample = None
diff --git a/resspect/fit_lightcurves.py b/resspect/fit_lightcurves.py
index 8f61585d..a5758cfa 100644
--- a/resspect/fit_lightcurves.py
+++ b/resspect/fit_lightcurves.py
@@ -1,40 +1,19 @@
-# Copyright 2020 resspect software
-# Author: Emille E. O. Ishida
-#
-# created on 14 April 2020
-#
-# Licensed GNU General Public License v3.0;
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.gnu.org/licenses/gpl-3.0.en.html
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
+"""
+    Fit light curves for different datasets
+"""
 import multiprocessing
 import logging
 import os
 from copy import copy
 from itertools import repeat
 from typing import IO
-from typing import Tuple
 from typing import Union
-import warnings
 
-import matplotlib.pylab as plt
 import numpy as np
 import pandas as pd
 
-from resspect.bazin import bazin
-from resspect.bazin import fit_scipy
-from resspect.bump import bump
-from resspect.bump import fit_bump
-from resspect.exposure_time_calculator import ExpTimeCalc
-from resspect.lightcurves_utils import read_file
+from resspect.feature_extractors.bazin import BazinFeatureExtractor
+from resspect.feature_extractors.bump import BumpFeatureExtractor
 from resspect.lightcurves_utils import get_resspect_header_data
 from resspect.lightcurves_utils import load_snpcc_photometry_df
 from resspect.lightcurves_utils import get_photometry_with_id_name_and_snid
@@ -51,878 +30,15 @@
 from resspect.lightcurves_utils import PLASTICC_TARGET_TYPES
 from resspect.lightcurves_utils import PLASTICC_RESSPECT_FEATURES_HEADER
 
+__all__ = ["fit_snpcc", "fit_resspect", "fit_plasticc"]
 
-warnings.filterwarnings("ignore", category=RuntimeWarning)
-logging.basicConfig(level=logging.INFO)
-
-__all__ = ['LightCurve', 'fit_snpcc', 'fit_resspect_bazin',
-           'fit_plasticc_bazin']
-
-
-class LightCurve:
-    """ Light Curve object, holding meta and photometric data.
-
-    Attributes
-    ----------
-    bazin_features_names: list
-        List of names of the Bazin function parameters.
-    bazin_features: list
-        List with the 5 best-fit Bazin parameters in all filters.
-        Concatenated from blue to red.
-    dataset_name: str
-        Name of the survey or data set being analyzed.
-    exp_time: dict
-        Exposure time required to take a spectra.
-        Keywords indicate telescope e.g.['4m', '8m'].
-    filters: list
-        List of broad band filters.
-    full_photometry: pd.DataFrame
-        Photometry for a set of light curves read from file.
-    id: int
-        SN identification number.
-    id_name:
-        Column name of object identifier.
-    last_mag: float
-        r-band magnitude of last observed epoch.
-    photometry: pd.DataFrame
-        Photometry information.
-        Minimum keys --> [mjd, band, flux, fluxerr].
-    redshift: float
-        Redshift
-    sample: str
-        Original sample to which this light curve is assigned.
-    sim_peakmag: np.array
-        Simulated peak magnitude in each filter.
-    sim_pkmjd: float
-        Simulated day of maximum, observer frame.
-    sncode: int
-        Number identifying the SN model used in the simulation.
-    sntype: str
-        General classification, possibilities are: Ia, II or Ibc.
-    unique_ids: str or array
-        List of unique ids available in the photometry file.
-        Only used for PLAsTiCC data.
-
-    Methods
-    -------
-    calc_exp_time(telescope_diam: float, SNR: float, telescope_name: str)
-        Calculates time required to take a spectra in the last obs epoch.
-    check_queryable(mjd: float, r_lim: float)
-        Check if this light can be queried in a given day.
-    conv_flux_mag(flux: np.array)
-        Convert positive flux into magnitude.
-    evaluate_bazin(param: list, time: np.array) -> np.array
-        Evaluate the Bazin function given parameter values.
-    load_snpcc_lc(path_to_data: str)
-        Reads header and photometric information for 1 light curve.
-    load_plasticc_lc(photo_file: str, snid: int)
-    	Load photometric information for 1 PLAsTiCC light curve.
-    load_resspect_lc(photo_file: str, snid: int)
-    	Load photometric information for 1 RESSPECT light curve.
-    fit_bazin(band: str) -> list
-        Calculates best-fit parameters from the Bazin function in 1 filter.
-    fit_bazin_all()
-        Calculates  best-fit parameters from the Bazin func for all filters.
-    plot_bazin_fit(save: bool, show: bool, output_file: srt)
-        Plot photometric points and Bazin fitted curve.
-    fit_bump(band: str) -> list
-        Calculates best-fit parameters from the Bump function in 1 filter.
-    fit_bump_all()
-        Calculates  best-fit parameters from the Bump func for all filters.
-    plot_bump_fit(save: bool, show: bool, output_file: srt)
-        Plot photometric points and Bump fitted curve.
-
-    Examples
-    --------
-
-    ##### for RESSPECT and PLAsTiCC light curves it is necessary to
-    ##### input the object identification for dealing with 1 light curve
+FEATURE_EXTRACTOR_MAPPING = {
+    "bazin": BazinFeatureExtractor,
+    "bump": BumpFeatureExtractor
+}
 
-    >>> import io
-    >>> import pandas as pd
-    >>> import tarfile
 
-    >>> from resspect import LightCurve
-
-    # path to header file
-    >>> path_to_header = '~/RESSPECT_PERFECT_V2_TRAIN_HEADER.tar.gz'
-
-    # openning '.tar.gz' files requires some juggling ...
-    >>> tar = tarfile.open(path_to_header, 'r:gz')
-    >>> fname = tar.getmembers()[0]
-    >>> content = tar.extractfile(fname).read()
-    >>> header = pd.read_csv(io.BytesIO(content))
-    >>> tar.close()
-
-    # choose one object
-    >>> snid = header['objid'].values[4]
-
-    # once you have the identification you can use this class
-    >>> path_to_lightcurves = '~/RESSPECT_PERFECT_V2_TRAIN_LIGHTCURVES.tar.gz'
-
-    >>> lc = LightCurve()                        # create light curve instance
-    >>> lc.load_snpcc_lc(path_to_lightcurves)    # read data
-    >>> lc.photometry
-             mjd band       flux   fluxerr         SNR
-    0    53214.0    u   0.165249  0.142422    1.160276
-    1    53214.0    g  -0.041531  0.141841   -0.292803
-    ..       ...  ...        ...       ...         ...
-    472  53370.0    z  68.645930  0.297934  230.406460
-    473  53370.0    Y  63.254270  0.288744  219.067050
-
-    >>> lc.fit_bazin_all()               # perform Bazin fit in all filters
-    >>> lc.bazin_features                # display Bazin parameters
-    [198.63302952843623, -9.38297128588733, 43.99971014717201,
-    ... ...
-    -1.546372806815066]
-
-    for fitting the entire sample ...
-
-    >>> output_file = 'RESSPECT_PERFECT_TRAIN.DAT'
-    >>> fit_resspect_bazin(path_to_lightcurves, path_to_header,
-                           output_file, sample='train')
-    """
-
-    def __init__(self):
-        self.queryable = None
-        self.bazin_features = []
-        self.bazin_features_names = ['a', 'b', 't0', 'tfall', 'trise']
-        self.bump_features = []
-        self.bump_features_names = ['p1', 'p2', 'p3', 'time_shift', 'max_flux']
-        self.dataset_name = ' '
-        self.exp_time = {}
-        self.filters = []
-        self.full_photometry = pd.DataFrame()
-        self.id = 0
-        self.id_name = None
-        self.last_mag = None
-        self.photometry = pd.DataFrame()
-        self.redshift = 0
-        self.sample = ' '
-        self.sim_peakmag = []
-        self.sim_pkmjd = None
-        self.sncode = 0
-        self.sntype = ' '
-
-
-    def _get_snpcc_photometry_raw_and_header(
-            self, lc_data: np.ndarray,
-            sntype_test_value: str = "-9") -> Tuple[np.ndarray, list]:
-        """
-        Reads SNPCC photometry raw and header data
-
-        Parameters
-        ----------
-        lc_data
-            SNPCC light curve data
-        sntype_test_value
-            test sample SNTYPE value
-        """
-        photometry_raw = []
-        header = []
-        for each_row in lc_data:
-            name = each_row[0]
-            value = each_row[1]
-            if name == 'SNID:':
-                self.id = int(value)
-                self.id_name = 'SNID'
-            elif name == 'SNTYPE:':
-                self.sample = 'test' if value == sntype_test_value else 'train'
-            elif name == 'SIM_REDSHIFT:':
-                self.redshift = float(value)
-            elif name == 'SIM_NON1a:':
-                self.sncode = value
-                self.sntype = get_snpcc_sntype(int(value))
-            elif name == 'VARLIST:':
-                header = each_row[1:]
-            elif name == 'OBS:':
-                photometry_raw.append(np.array(each_row[1:]))
-            elif name == 'SIM_PEAKMAG:':
-                self.sim_peakmag = np.array(each_row[1:5]).astype(float)
-            elif name == 'SIM_PEAKMJD:':
-                self.sim_pkmjd = float(value)
-        return np.array(photometry_raw), header
-
-    def load_snpcc_lc(self, path_to_data: str):
-        """Reads one LC from SNPCC data.
-
-        Populates the attributes: dataset_name, id, sample, redshift, sncode,
-        sntype, photometry, sim_peakmag and sim_pkmjd.
-
-        Parameters
-        ---------
-        path_to_data: str
-            Path to text file with data from a single SN.
-        """
-        self.dataset_name = 'SNPCC'
-        self.filters = ['g', 'r', 'i', 'z']
-
-        lc_data = np.array(read_file(path_to_data), dtype=object)
-        photometry_raw, header = self._get_snpcc_photometry_raw_and_header(
-            lc_data)
-
-        if photometry_raw.size > 0:
-            self.photometry = load_snpcc_photometry_df(photometry_raw, header)
-
-    def load_resspect_lc(self, photo_file: str, snid: int):
-        """
-        Return 1 light curve from RESSPECT simulations.
-
-        Parameters
-        ----------
-        photo_file: str
-            Complete path to light curves file.
-        snid: int
-            Identification number for the desired light curve.
-        """
-
-        self.dataset_name = 'RESSPECT'
-        self.filters = ['u', 'g', 'r', 'i', 'z', 'Y']
-        self.id = snid
-
-        if self.full_photometry.empty:
-            _, self.full_photometry = read_resspect_full_photometry_data(
-                photo_file)
-        id_names_list = ['SNID', 'snid', 'objid', 'id']
-        filtered_photometry, self.id_name = (
-            get_photometry_with_id_name_and_snid(
-                self.full_photometry, id_names_list, snid))
-
-        if not filtered_photometry.empty:
-            filtered_photometry = insert_band_column_to_resspect_df(
-                filtered_photometry, self.filters)
-            self.photometry = load_resspect_photometry_df(filtered_photometry)
-
-    def load_elasticc_lc(self, photo_file: str, snid: int):
-        """
-        Return 1 light curve from ELAsTiCC simulations.
-
-        Parameters
-        ----------
-        photo_file: str
-            Complete path to light curve file.
-        snid: int
-            Identification number for the desired light curve.
-        """
-        self.dataset_name = 'ELAsTiCC'
-        self.filters = ['u', 'g', 'r', 'i', 'z', 'Y']
-        self.id = snid
-
-        if self.full_photometry.empty:
-            self.full_photometry = read_elasticc_full_photometry_data(
-                photo_file)[1]
-
-        id_names_list = ['object_id', 'SNID', 'snid']
-
-        filtered_photometry, self.id_name = (
-            get_photometry_with_id_name_and_snid(
-                self.full_photometry, id_names_list, snid))
-
-        filter_mapping_dict = {
-            b'u ': 'u', b'g ': 'g', b'r ': 'r', b'i ': 'i', b'z ': 'z', b'Y ': 'Y'
-        }
-
-        if not filtered_photometry.empty:
-            self.photometry = load_elasticc_photometry_df(
-                filtered_photometry, filter_mapping_dict)
-
-    def load_plasticc_lc(self, photo_file: str, snid: int):
-        """
-        Return 1 light curve from PLAsTiCC simulations.
-
-        Parameters
-        ----------
-        photo_file: str
-            Complete path to light curve file.
-        snid: int
-            Identification number for the desired light curve.
-        """
-        self.dataset_name = 'PLAsTiCC'
-        self.filters = ['u', 'g', 'r', 'i', 'z', 'Y']
-        self.id = snid
-
-        if self.full_photometry.empty:
-            self.full_photometry = read_plasticc_full_photometry_data(
-                photo_file)
-
-        id_names_list = ['object_id', 'SNID', 'snid']
-
-        filtered_photometry, self.id_name = (
-            get_photometry_with_id_name_and_snid(
-                self.full_photometry, id_names_list, snid))
-
-        filter_mapping_dict = {
-            0: 'u', 1: 'g', 2: 'r', 3: 'i', 4: 'z', 5: 'Y'
-        }
-
-        if not filtered_photometry.empty:
-            self.photometry = load_plasticc_photometry_df(
-                filtered_photometry, filter_mapping_dict)
-
-    def conv_flux_mag(self, flux, zpt=27.5):
-        """Convert FLUXCAL to magnitudes.
-
-        Parameters
-        ----------
-        flux: list or np.array
-            Values of flux to be converted into mag.
-        zpt: float (optional)
-            Zero point. Default is 27.5 (from SNANA).
-
-        Returns
-        -------
-        mag: list or np.array
-            Magnitude values. If flux < 1e-5 returns 9999.
-        """
-
-        if not isinstance(flux[0], float):
-            flux = np.array([item[0] for item in flux])
-
-        mag = [zpt - 2.5 * np.log10(f) if f > 1e-5 else 9999 for f in flux]
-
-        return np.array(mag)
-
-    def check_queryable(self, mjd: float, filter_lim: float, criteria=1,
-                        days_since_last_obs=2, feature_method='Bazin',
-                        filter_cut='r'):
-        """Check if this object can be queried in a given day.
-
-        Parameters
-        ----------
-        mjd: float
-            MJD where the query will take place.
-        filter_lim: float
-            Magnitude limit below which query is possible.
-        criteria: int [1, 2 or 3] (optional)
-            Criteria to determine if an obj is queryable.
-            1 -> Cut on last measured photometric point.
-            2 -> if last obs was further than days_since_last_obs,
-                 use Bazin estimate for today. Otherwise, use
-                 the last observed point.
-            Default is 1.
-        days_since_last_obs: int (optional)
-            If there is an observation within these days, use the
-            measured value, otherwise estimate current mag.
-            Only used if "criteria == 2". Default is 2.
-        feature_method: str (optional)
-            Feature extraction method. Only 'Bazin' is implemented.
-            Default is 'Bazin'.
-        filter_cut: str (optional)
-            Band in which cut is applied. Default is 'r'.
-
-        Returns
-        -------
-        bool
-            True if current magnitude lower than "filter_lim".
-        """
-
-        # create photo flag
-        photo_flag = self.photometry['mjd'].values <= mjd
-        rband_flag = self.photometry['band'].values == filter_cut
-        surv_flag = np.logical_and(photo_flag, rband_flag)
-
-        if criteria == 1 and sum(surv_flag):
-            if 'MAG' in self.photometry.keys():
-                # check surviving photometry
-                self.last_mag = self.photometry['MAG'].values[surv_flag][-1]
-
-            else:
-                surv_flux = self.photometry['flux'].values[surv_flag]
-                self.last_mag = self.conv_flux_mag([surv_flux[-1]])[0]
-
-        elif criteria == 2 and sum(surv_flag):
-            # check if there is an observation recently
-            surv_mjd = self.photometry['mjd'].values[surv_flag]
-            gap = mjd - surv_mjd[-1]
-
-            if gap <= days_since_last_obs:
-                if 'MAG' in self.photometry.keys():
-                    # check surviving photometry
-                    self.last_mag = \
-                            self.photometry['MAG'].values[surv_flag][-1]
-
-                else:
-                    surv_flux = self.photometry['flux'].values[surv_flag]
-                    self.last_mag = self.conv_flux_mag([surv_flux[-1]])[0]
-
-            elif feature_method == 'Bazin':
-                # get first day of observation in this filter
-                mjd_min = min(self.photometry['mjd'].values[surv_flag])
-
-                # estimate flux based on Bazin function
-                fitted_flux = self.evaluate_bazin([mjd - mjd_min])[filter_cut][0]
-                self.last_mag = self.conv_flux_mag([fitted_flux])[0]
-
-            else:
-                raise ValueError('Only "Bazin" features are implemented!')
-
-        elif sum(surv_flag):
-            raise ValueError('Criteria needs to be "1" or "2". \n ' + \
-                             'See docstring for further info.')
-        elif not sum(surv_flag):
-            return False
-
-        if self.last_mag <= filter_lim:
-            return True
-        else:
-            return False
-
-    def calc_exp_time(self, telescope_diam: float, SNR: float,
-                      telescope_name: str, **kwargs):
-        """Calculates time required to take a spectra in the last obs epoch.
-
-        Populates attribute exp_time.
-
-        Parameters
-        ----------
-        SNR: float
-            Required SNR.
-        telescope_diam: float
-            Diameter of primary mirror for spectroscopic telescope in meters.
-        telescope_name: str
-            Identification for telescope.
-        kwargs: extra parameters
-            Any input required by ExpTimeCalc.findexptime function.
-
-        Returns
-        -------
-        exp_time: float
-            Exposure time require for taking a spectra in seconds.
-        """
-
-        # check if last magnitude was calculated
-        if self.last_mag == None:
-            raise ValueError('Magnitude at last epoch not calculated.\n' + \
-                             'Run the check_queryable function.')
-
-        etc = ExpTimeCalc()
-        etc.diameter = telescope_diam
-        exp_time = etc.findexptime(SNRin=SNR, mag=self.last_mag, **kwargs)
-
-        if exp_time > 60 and exp_time < 7200 and self.last_mag < 30:
-            self.exp_time[telescope_name] = exp_time
-            return exp_time
-        else:
-            self.exp_time[telescope_name] = 9999
-            return 9999
-
-    def fit_bazin(self, band: str) -> np.ndarray:
-        """Extract Bazin features for one filter.
-
-        Parameters
-        ----------
-        band: str
-            Choice of broad band filter
-
-        Returns
-        -------
-        bazin_param: np.ndarray
-            Best fit parameters for the Bazin function:
-            [a, b, t0, tfall, trise].
-        """
-
-        # build filter flag
-        band_indices = self.photometry['band'] == band
-        if not sum(band_indices) > (len(self.bazin_features_names) - 1):
-            return np.array([])
-
-        # get info for this filter
-        time = self.photometry['mjd'].values[band_indices]
-        flux = self.photometry['flux'].values[band_indices]
-        fluxerr = self.photometry['fluxerr'].values[band_indices]
-
-        # fit Bazin function
-        bazin_param = fit_scipy(time - time[0], flux, fluxerr)
-
-        return bazin_param
-
-    def evaluate_bazin(self, time: np.array):
-        """Evaluate the Bazin function given parameter values.
-
-        Parameters
-        ----------
-        time: np.array or list
-            Time since first light curve observation.
-
-        Returns
-        -------
-        dict
-            Value of the Bazin flux in each required time per filter.
-        """
-        # store flux values and starting points
-        flux = {}
-
-        for k in range(len(self.filters)):
-            # store flux values per filter
-            flux[self.filters[k]] = []
-
-            # check if Bazin features exist
-            if 'None' not in self.bazin_features[k * 5 : (k + 1) * 5]:
-                for item in time:
-                    flux[self.filters[k]].append(\
-                           bazin(item, self.bazin_features[0 + k * 5],
-                                 self.bazin_features[1 + k * 5],
-                                 self.bazin_features[2 + k * 5],
-                                 self.bazin_features[3 + k * 5],
-                                 self.bazin_features[4 + k * 5]))
-            else:
-                flux[self.filters[k]].append(None)
-
-        return flux
-
-    def fit_bazin_all(self):
-        """
-        Perform Bazin fit for all filters independently and concatenate results.
-        Populates the attributes: bazin_features.
-        """
-        default_bazin_features = ['None'] * len(self.bazin_features_names)
-
-        if self.photometry.shape[0] < 1:
-            self.bazin_features = ['None'] * len(self.bazin_features_names) * len(self.filters)
-
-        elif 'None' not in self.bazin_features:
-            self.bazin_features = []
-            for each_band in self.filters:
-                best_fit = self.fit_bazin(band=each_band)
-                if (best_fit.size > 0) and (not np.isnan(np.sum(best_fit))):
-                    self.bazin_features.extend(best_fit.tolist())
-                else:
-                    self.bazin_features.extend(default_bazin_features)
-        else:
-            self.bazin_features.extend(default_bazin_features)
-            
-    def fit_bump(self, band: str) -> np.ndarray:
-        """Extract Bump features for one filter.
-
-        Parameters
-        ----------
-        band: str
-            Choice of broad band filter
-
-        Returns
-        -------
-        bump_param: np.ndarray
-            Best fit parameters for the Bazin function:
-            [p1, p2, p3, time_shift, max_flux].
-        """
-
-        # build filter flag
-        band_indices = self.photometry['band'] == band
-        if not sum(band_indices) > (len(self.bump_features_names) - 2):
-            return np.array([])
-
-        # get info for this filter
-        time = self.photometry['mjd'].values[band_indices]
-        flux = self.photometry['flux'].values[band_indices]
-        fluxerr = self.photometry['fluxerr'].values[band_indices]
-
-        # fit Bump function
-        bump_param = fit_bump(time, flux, fluxerr)
-
-        return bump_param
-
-    def evaluate_bump(self, time: np.array):
-        """Evaluate the Bump function given parameter values.
-
-        Parameters
-        ----------
-        time: np.array or list
-            Time since first light curve observation.
-
-        Returns
-        -------
-        dict
-            Value of the Bump flux in each required time per filter.
-        """
-        # store flux values and starting points
-        flux = {}
-
-        for k in range(len(self.filters)):
-            # store flux values per filter
-            flux[self.filters[k]] = []
-
-            # check if Bump features exist
-            if 'None' not in self.bump_features[k * 5 : (k + 1) * 5]:
-                for item in time:
-                    flux[self.filters[k]].append(\
-                           bump(item, self.bump_features[0 + k * 5],
-                                 self.bump_features[1 + k * 5],
-                                 self.bump_features[2 + k * 5]))
-            else:
-                flux[self.filters[k]].append(None)
-
-        return flux
-
-    def fit_bump_all(self):
-        """
-        Perform Bump fit for all filters independently and concatenate results.
-        Populates the attributes: bump_features.
-        """
-        default_bump_features = ['None'] * len(self.bump_features_names)
-
-        if self.photometry.shape[0] < 1:
-            self.bump_features = ['None'] * len(self.bump_features_names) * len(self.filters)
-
-        elif 'None' not in self.bump_features:
-            self.bump_features = []
-            for each_band in self.filters:
-                best_fit = self.fit_bump(band=each_band)
-                if (len(best_fit) > 0) and (not np.isnan(np.sum(best_fit))):
-                    self.bump_features.extend(best_fit)
-                else:
-                    self.bump_features.extend(default_bump_features)
-        else:
-            self.bump_features.extend(default_bump_features)
-
-    def clear_data(self):
-        """ Reset to default values """
-        self.photometry = []
-        self.redshift = None
-        self.sntype = None
-        self.sncode = None
-        self.sample = None
-
-    def plot_bazin_fit(self, save=True, show=False, output_file=' ',
-                       figscale=1, extrapolate=False,
-                       time_flux_pred=None, unit='flux'):
-        """
-        Plot data and Bazin fitted function.
-
-        Parameters
-        ----------
-        figscale: float (optional)
-            Allow to control the size of the figure.
-        extrapolate: bool (optional)
-            If True, also plot the estimated flux values.
-            Default is False.
-        output_file: str (optional)
-            Name of file to store the plot.
-        save: bool (optional)
-             Save figure to file. Default is True.
-        show: bool (optinal)
-             Display plot in windown. Default is False.
-        time_flux_pred: list (optional)
-            Time since first observation where flux is to be
-            estimated. It is only used if "extrapolate == True".
-            Default is None.
-        unit: str (optional)
-            Unit for plot. Options are 'flux' or 'mag'.
-            Use zero point from SNANA for flux-to-mag conversion
-            ==> mag = 2.5 * (11 - np.log10(flux)).
-            Default is 'flux'.
-        """
-
-        # number of columns in the plot
-        ncols = len(self.filters) / 2 + len(self.filters) % 2
-        fsize = (figscale * 5 * ncols , figscale * 10)
-
-        plt.figure(figsize=fsize)
-
-        for i in range(len(self.filters)):
-            plt.subplot(2, int(ncols), i + 1)
-            plt.title('Filter: ' + self.filters[i])
-
-            # filter flag
-            filter_flag = self.photometry['band'] == self.filters[i]
-            x = self.photometry['mjd'][filter_flag].values
-            y = self.photometry['flux'][filter_flag].values
-            yerr = self.photometry['fluxerr'][filter_flag].values
-
-            # check Bazin fit convergence
-            if 'None' in self.bazin_features[i * 5 : (i + 1) * 5]:
-                plot_fit = False
-            else:
-                plot_fit = True
-
-            # shift to avoid large numbers in x-axis
-            time = x - min(x)
-
-            if plot_fit:
-                xaxis = np.linspace(0, max(time), 500)[:, np.newaxis]
-                fitted_flux = self.evaluate_bazin(xaxis)
-                if unit == 'flux':
-                    plt.plot(xaxis, fitted_flux[self.filters[i]], color='red',
-                             lw=1.5, label='Bazin fit')
-                elif unit == 'mag':
-                    mag = self.conv_flux_mag(fitted_flux[self.filters[i]])
-                    mag_flag = mag < 50
-                    plt.plot(xaxis[mag_flag], mag[mag_flag], color='red',
-                             lw=1.5)
-                else:
-                    raise ValueError('Unit can only be "flux" or "mag".')
-
-                if extrapolate:
-                    xaxis_extrap = list(xaxis.flatten()) + list(time_flux_pred)
-                    xaxis_extrap = np.sort(np.array(xaxis_extrap))
-                    ext_flux = self.evaluate_bazin(xaxis_extrap)
-                    if unit == 'flux':
-                        plt.plot(xaxis_extrap, ext_flux[self.filters[i]],
-                                 color='red', lw=1.5, ls='--',
-                                 label='Bazin extrap')
-                    elif unit == 'mag':
-                        ext_mag = self.conv_flux_mag(ext_flux[self.filters[i]])
-                        ext_mag_flag = ext_mag < 50
-                        plt.plot(xaxis_extrap[ext_mag_flag],
-                                 ext_mag[ext_mag_flag],
-                                 color='red', lw=1.5, ls='--')
-
-            if unit == 'flux':
-                plt.errorbar(time, y, yerr=yerr, color='blue', fmt='o',
-                             label='obs')
-                plt.ylabel('FLUXCAL')
-            elif unit == 'mag':
-                mag_obs  = self.conv_flux_mag(y)
-                mag_obs_flag = mag_obs < 50
-                time_mag = time[mag_obs_flag]
-
-                plt.scatter(time_mag, mag_obs[mag_obs_flag], color='blue',
-                            label='calc mag', marker='s')
-
-                # if MAG is provided in the table, also plot it
-                # this allows checking the flux mag conversion
-                if 'MAG' in self.photometry.keys():
-                    mag_flag = self.photometry['MAG'].values < 50
-                    mag_ff = np.logical_and(filter_flag, mag_flag)
-                    mag_table = self.photometry['MAG'][mag_ff].values
-                    mjd_table = self.photometry['mjd'][mag_ff].values - min(x)
-
-                    plt.scatter(mjd_table, mag_table, color='black',
-                                label='table mag', marker='x')
-
-                ax = plt.gca()
-                ax.set_ylim(ax.get_ylim()[::-1])
-                plt.ylabel('mag')
-
-            plt.xlabel('days since first observation')
-            plt.tight_layout()
-
-        if save:
-            plt.savefig(output_file)
-        if show:
-            plt.show()
-                        
-    def plot_bump_fit(self, save=True, show=False, output_file=' ',
-                       figscale=1, extrapolate=False,
-                       time_flux_pred=None, unit='flux'):
-        """
-        Plot data and Bump fitted function.
-
-        Parameters
-        ----------
-        figscale: float (optional)
-            Allow to control the size of the figure.
-        extrapolate: bool (optional)
-            If True, also plot the estimated flux values.
-            Default is False.
-        output_file: str (optional)
-            Name of file to store the plot.
-        save: bool (optional)
-             Save figure to file. Default is True.
-        show: bool (optinal)
-             Display plot in windown. Default is False.
-        time_flux_pred: list (optional)
-            Time after last observation where flux is to be
-            estimated. It is only used if "extrapolate == True".
-            Default is None.
-        unit: str (optional)
-            Unit for plot. Options are 'flux' or 'mag'.
-            Use zero point from SNANA for flux-to-mag conversion
-            ==> mag = 2.5 * (11 - np.log10(flux)).
-            Default is 'flux'.
-        """
-
-        # number of columns in the plot
-        ncols = len(self.filters) / 2 + len(self.filters) % 2
-        fsize = (figscale * 5 * ncols , figscale * 10)
-
-        plt.figure(figsize=fsize)
-
-        for i in range(len(self.filters)):
-            plt.subplot(2, int(ncols), i + 1)
-            plt.title('Filter: ' + self.filters[i])
-
-            # filter flag
-            filter_flag = self.photometry['band'] == self.filters[i]
-            x = self.photometry['mjd'][filter_flag].values
-            y = self.photometry['flux'][filter_flag].values
-            yerr = self.photometry['fluxerr'][filter_flag].values
-
-            # check Bump fit convergence
-            if 'None' in self.bump_features[i * 5 : (i + 1) * 5]:
-                plot_fit = False
-            else:
-                plot_fit = True
-
-            # shift to avoid large numbers in x-axis
-            time = x + self.bump_features[i * 5 + 3]
-
-            if plot_fit:
-                xaxis = np.linspace(min(time), max(time), 500)[:, np.newaxis]
-                fitted_flux = np.array(self.evaluate_bump(xaxis)[self.filters[i]]) * self.bump_features[i * 5 + 4]
-                if unit == 'flux':
-                    plt.plot(xaxis, fitted_flux, color='red',
-                             lw=1.5, label='Bump fit')
-                elif unit == 'mag':
-                    mag = self.conv_flux_mag(fitted_flux[self.filters[i]])
-                    mag_flag = mag < 50
-                    plt.plot(xaxis[mag_flag], mag[mag_flag], color='red',
-                             lw=1.5)
-                else:
-                    raise ValueError('Unit can only be "flux" or "mag".')
-
-                if extrapolate:
-                    xaxis_extrap = list(xaxis.flatten()) + [time_flux_pred + max(time)]
-                    xaxis_extrap = np.sort(np.array(xaxis_extrap))
-                    ext_flux = np.array(self.evaluate_bump(xaxis_extrap)[self.filters[i]]) \
-                               * self.bump_features[i * 5 + 4]
-                    if unit == 'flux':
-                        plt.plot(xaxis_extrap, ext_flux,
-                                 color='red', lw=1.5, ls='--',
-                                 label='Bump extrap')
-                    elif unit == 'mag':
-                        ext_mag = self.conv_flux_mag(ext_flux[self.filters[i]])
-                        ext_mag_flag = ext_mag < 50
-                        plt.plot(xaxis_extrap[ext_mag_flag],
-                                 ext_mag[ext_mag_flag],
-                                 color='red', lw=1.5, ls='--')
-
-            if unit == 'flux':
-                plt.errorbar(time, y, yerr=yerr, color='blue', fmt='o',
-                             label='obs')
-                plt.ylabel('FLUXCAL')
-            elif unit == 'mag':
-                mag_obs  = self.conv_flux_mag(y)
-                mag_obs_flag = mag_obs < 50
-                time_mag = time[mag_obs_flag]
-
-                plt.scatter(time_mag, mag_obs[mag_obs_flag], color='blue',
-                            label='calc mag', marker='s')
-
-                # if MAG is provided in the table, also plot it
-                # this allows checking the flux mag conversion
-                if 'MAG' in self.photometry.keys():
-                    mag_flag = self.photometry['MAG'].values < 50
-                    mag_ff = np.logical_and(filter_flag, mag_flag)
-                    mag_table = self.photometry['MAG'][mag_ff].values
-                    mjd_table = self.photometry['mjd'][mag_ff].values - min(x)
-
-                    plt.scatter(mjd_table, mag_table, color='black',
-                                label='table mag', marker='x')
-
-                ax = plt.gca()
-                ax.set_ylim(ax.get_ylim()[::-1])
-                plt.ylabel('mag')
-
-            plt.xlabel('days since first observation')
-            plt.tight_layout()
-
-        if save:
-            plt.savefig(output_file)
-        if show:
-            plt.show()
-
-
-def _get_features_to_write(light_curve_data: LightCurve, function: str ='bazin') -> list:
+def _get_features_to_write(light_curve_data) -> list:
     """
     Returns features list to write
 
@@ -930,24 +46,16 @@ def _get_features_to_write(light_curve_data: LightCurve, function: str ='bazin')
     ----------
     light_curve_data
         fitted light curve data
-    function: str (optional)
-        Function used for feature extraction. 
-        Options are "bazin" or "bump". Default is "bazin".
     """
     features_list = [light_curve_data.id, light_curve_data.redshift,
                      light_curve_data.sntype, light_curve_data.sncode,
                      light_curve_data.sample]
-    
-    if function == 'bazin':
-        features_list.extend(light_curve_data.bazin_features)
-    elif function == 'bump':
-        features_list.extend(light_curve_data.bump_features)
-        
+    features_list.extend(light_curve_data.features)
     return features_list
 
 
 def write_features_to_output_file(
-        light_curve_data: LightCurve, features_file: IO):
+        light_curve_data, features_file: IO):
     """
     Writes fitted light curve data to output features file
 
@@ -966,9 +74,9 @@ def write_features_to_output_file(
 
 
 def _snpcc_sample_fit(
-        file_name: str, path_to_data_dir: str, function: str) -> LightCurve:
+        file_name: str, path_to_data_dir: str, feature_extractor: str):
     """
-    Reads SNPCC file and performs bazin fit
+    Reads SNPCC file and performs fit
     Parameters
     ----------
     file_name
@@ -976,30 +84,23 @@ def _snpcc_sample_fit(
     path_to_data_dir
          Path to directory containing the set of individual files,
          one for each light curve.
-    function
-        Function used for feature extraction. 
+    feature_extractor
+        Function used for feature extraction.
         Options are 'bazin' or 'bump'.
     """
-    light_curve_data = LightCurve()
+    light_curve_data = FEATURE_EXTRACTOR_MAPPING[feature_extractor]()
     light_curve_data.load_snpcc_lc(
         os.path.join(path_to_data_dir, file_name))
-    
-    if function == 'bazin':
-        light_curve_data.fit_bazin_all()
-    elif function == 'bump':
-        light_curve_data.fit_bump_all()
-    else:
-        raise ValueError('Possible options for "function" are "bazin" or "bump".')
-        
+    light_curve_data.fit_all()
     return light_curve_data
 
 
 def fit_snpcc(
         path_to_data_dir: str, features_file: str,
-        file_prefix: str = "DES_SN", number_of_processors: int = 1, 
-        function: str = 'bazin'):
+        file_prefix: str = "DES_SN", number_of_processors: int = 1,
+        feature_extractor: str = 'bazin'):
     """
-    Perform Bazin fit to all objects in the SNPCC data.
+    Perform fit to all objects in the SNPCC data.
 
      Parameters
      ----------
@@ -1012,34 +113,35 @@ def fit_snpcc(
         File names prefix
      number_of_processors: int, default 1
         Number of cpu processes to use.
-     function: str, default bazin
+     feature_extractor: str, default bazin
         Function used for feature extraction.
     """
     files_list = os.listdir(path_to_data_dir)
     files_list = [each_file for each_file in files_list
                   if each_file.startswith(file_prefix)]
     multi_process = multiprocessing.Pool(number_of_processors)
-    logging.info("Starting SNPCC " + function + " fit...")
+    logging.info("Starting SNPCC " + feature_extractor + " fit...")
     with open(features_file, 'w') as snpcc_features_file:
-        snpcc_features_file.write(','.join(SNPCC_FEATURES_HEADER) + '\n')
+        # TODO: Current implementation uses bazin features header for
+        #  all feature extraction
+        snpcc_features_file.write(' '.join(SNPCC_FEATURES_HEADER) + '\n')
+        
         for light_curve_data in multi_process.starmap(
                 _snpcc_sample_fit, zip(
-                    files_list, repeat(path_to_data_dir), repeat(function))):
-            features = {'bazin': light_curve_data.bazin_features,
-                        'bump': light_curve_data.bump_features}
-            if 'None' not in features[function]:
+                    files_list, repeat(path_to_data_dir), repeat(feature_extractor))):
+            if 'None' not in light_curve_data.features:
                 write_features_to_output_file(
                     light_curve_data, snpcc_features_file)
     logging.info("Features have been saved to: %s", features_file)
 
 
-def _resspect_sample_fit_bazin(
+def _resspect_sample_fit(
         index: int, snid: int, path_photo_file: str,
-        sample: str, light_curve_data: LightCurve, meta_header: pd.DataFrame,
+        sample: str, light_curve_data, meta_header: pd.DataFrame,
         redshift_name: Union[str, None], sncode_name: Union[str, None],
-        sntype_name: Union[str, None]) -> LightCurve:
+        sntype_name: Union[str, None]):
     """
-    Performs bazin fit for PLAsTiCC dataset with snid
+    Performs fit for PLAsTiCC dataset with snid
 
     Parameters
     ----------
@@ -1063,7 +165,7 @@ def _resspect_sample_fit_bazin(
         sntype meta header column name
     """
     light_curve_data.load_resspect_lc(path_photo_file, snid)
-    light_curve_data.fit_bazin_all()
+    light_curve_data.fit_all()
     light_curve_data.redshift = meta_header[redshift_name][index]
     light_curve_data.sncode = meta_header[sncode_name][index]
     light_curve_data.sntype = meta_header[sntype_name][index]
@@ -1073,10 +175,11 @@ def _resspect_sample_fit_bazin(
     return light_curve_data_copy
 
 
-def fit_resspect_bazin(path_photo_file: str, path_header_file: str,
-                       output_file: str, sample=None,
-                       number_of_processors: int = 1):
-    """Perform Bazin fit to all objects in a given RESSPECT data file.
+def fit_resspect(path_photo_file: str, path_header_file: str,
+                 output_file: str, sample=None,
+                 number_of_processors: int = 1,
+                 feature_extractor: str = "bazin"):
+    """Perform fit to all objects in a given RESSPECT data file.
 
     Parameters
     ----------
@@ -1090,6 +193,8 @@ def fit_resspect_bazin(path_photo_file: str, path_header_file: str,
         'train' or 'test'. Default is None.
     number_of_processors: int, default 1.
         Number of cpu processes to use
+    feature_extractor: str, default bazin
+        feature extraction method
     """
     meta_header = get_resspect_header_data(path_header_file, path_photo_file)
 
@@ -1103,12 +208,14 @@ def fit_resspect_bazin(path_photo_file: str, path_header_file: str,
     subtype_name = find_available_key_name_in_header(
         meta_header_keys, ['code', 'SIM_TYPE_INDEX', 'SNTYPE_SUBCLASS'])
 
-    light_curve_data = LightCurve()
+    light_curve_data = FEATURE_EXTRACTOR_MAPPING[feature_extractor]()
     snid_values = meta_header[id_name]
     snid_values = np.array(list(snid_values.items()))
     multi_process = multiprocessing.Pool(number_of_processors)
-    logging.info("Starting RESSPECT bazin fit...")
+    logging.info("Starting RESSPECT " + feature_extractor + " fit...")
     with open(output_file, 'w') as ressepect_features_file:
+        # TODO: Current implementation uses bazin features header
+        #  for all feature extraction
         ressepect_features_file.write(
             ','.join(PLASTICC_RESSPECT_FEATURES_HEADER) + '\n')
         iterator_list = zip(
@@ -1117,20 +224,20 @@ def fit_resspect_bazin(path_photo_file: str, path_header_file: str,
             repeat(meta_header), repeat(z_name), repeat(subtype_name),
             repeat(type_name))
         for light_curve_data in multi_process.starmap(
-                _resspect_sample_fit_bazin, iterator_list):
-            if 'None' not in light_curve_data.bazin_features:
+                _resspect_sample_fit, iterator_list):
+            if 'None' not in light_curve_data.features:
                 write_features_to_output_file(
                     light_curve_data, ressepect_features_file)
             light_curve_data.clear_data()
     logging.info("Features have been saved to: %s", output_file)
 
 
-def _plasticc_sample_fit_bazin(
+def _plasticc_sample_fit(
         index: int, snid: int, path_photo_file: str,
-        sample: str, light_curve_data: LightCurve,
-        meta_header: pd.DataFrame) -> LightCurve:
+        sample: str, light_curve_data,
+        meta_header: pd.DataFrame):
     """
-    Performs bazin fit for PLAsTiCC dataset with snid
+    Performs fit for PLAsTiCC dataset with snid
 
     Parameters
     ----------
@@ -1148,7 +255,7 @@ def _plasticc_sample_fit_bazin(
         photometry meta header data
     """
     light_curve_data.load_plasticc_lc(path_photo_file, snid)
-    light_curve_data.fit_bazin_all()
+    light_curve_data.fit_all()
     light_curve_data.redshift = meta_header['true_z'][index]
     light_curve_data.sncode = meta_header['true_target'][index]
     light_curve_data.sntype = PLASTICC_TARGET_TYPES[
@@ -1159,11 +266,12 @@ def _plasticc_sample_fit_bazin(
     return light_curve_data_copy
 
 
-def fit_plasticc_bazin(path_photo_file: str, path_header_file: str,
-                       output_file: str, sample='train',
-                       number_of_processors: int = 1):
+def fit_plasticc(path_photo_file: str, path_header_file: str,
+                 output_file: str, sample='train',
+                 feature_extractor: str = "bazin",
+                 number_of_processors: int = 1):
     """
-    Perform Bazin fit to all objects in a given PLAsTiCC data file.
+    Perform fit to all objects in a given PLAsTiCC data file.
     Parameters
     ----------
     path_photo_file: str
@@ -1176,6 +284,8 @@ def fit_plasticc_bazin(path_photo_file: str, path_header_file: str,
         'train' or 'test'. Default is 'train'.
     number_of_processors: int, default 1
         Number of cpu processes to use.
+    feature_extractor: str, default bazin
+        feature extraction method
     """
 
     name_list = ['SNID', 'snid', 'objid', 'object_id']
@@ -1183,7 +293,7 @@ def fit_plasticc_bazin(path_photo_file: str, path_header_file: str,
     meta_header_keys = meta_header.keys().tolist()
     id_name = find_available_key_name_in_header(
         meta_header_keys, name_list)
-    light_curve_data = LightCurve()
+    light_curve_data = FEATURE_EXTRACTOR_MAPPING[feature_extractor]()
 
     if sample == 'train':
         snid_values = meta_header[id_name]
@@ -1195,16 +305,18 @@ def fit_plasticc_bazin(path_photo_file: str, path_header_file: str,
             columns=[id_name])[id_name]
     snid_values = np.array(list(snid_values.items()))
     multi_process = multiprocessing.Pool(number_of_processors)
-    logging.info("Starting PLAsTiCC bazin fit...")
+    logging.info("Starting PLAsTiCC " + feature_extractor + " fit...")
     with open(output_file, 'w') as plasticc_features_file:
+        # TODO: Current implementation uses bazin features header for
+        #  all feature extraction
         plasticc_features_file.write(
             ','.join(PLASTICC_RESSPECT_FEATURES_HEADER) + '\n')
         iterator_list = zip(
             snid_values[:, 0].tolist(), snid_values[:, 1].tolist(), repeat(path_photo_file),
             repeat(sample), repeat(light_curve_data), repeat(meta_header))
         for light_curve_data in multi_process.starmap(
-                _plasticc_sample_fit_bazin, iterator_list):
-            if 'None' not in light_curve_data.bazin_features:
+                _plasticc_sample_fit, iterator_list):
+            if 'None' not in light_curve_data.features:
                 write_features_to_output_file(
                     light_curve_data, plasticc_features_file)
     logging.info("Features have been saved to: %s", output_file)
diff --git a/resspect/learn_loop.py b/resspect/learn_loop.py
index 9e9ec972..3c984faa 100644
--- a/resspect/learn_loop.py
+++ b/resspect/learn_loop.py
@@ -49,7 +49,7 @@ def load_features(database_class: DataBase,
        'DES' or 'LSST'. Default is 'DES'.
         Name of the survey which characterizes filter set.
     features_method
-        Feature extraction method. Currently only 'Bazin' and "bump" are implemented.
+        Feature extraction method. Currently only 'bazin' and "bump" are implemented.
     number_of_classes
         Number of classes to consider in the classification
         Currently only nclass == 2 is implemented.
@@ -71,14 +71,14 @@ def load_features(database_class: DataBase,
     """
     if isinstance(path_to_features, str):
         database_class.load_features(
-            path_to_file=path_to_features, method=features_method,
+            path_to_file=path_to_features, feature_extractor=features_method,
             survey=survey)
     else:
         features_set_names = ['train', 'test', 'validation', 'pool']
         for sample_name in features_set_names:
             if sample_name in path_to_features.keys():
                 database_class.load_features(
-                    path_to_features[sample_name], method=features_method,
+                    path_to_features[sample_name], feature_extractor=features_method,
                     survey=survey, sample=sample_name)
             else:
                 logging.warning(f'Path to {sample_name} not given.'
@@ -320,7 +320,7 @@ def update_alternative_label(database_class_alternative: DataBase,
 # TODO: too many arguments! refactor further and update docs
 def learn_loop(nloops: int, strategy: str, path_to_features: str,
                output_metrics_file: str, output_queried_file: str,
-               features_method: str ='Bazin', classifier: str = 'RandomForest',
+               features_method: str = 'bazin', classifier: str = 'RandomForest',
                training: str = 'original', batch: int =1, survey: str = 'DES',
                nclass: int = 2, photo_class_thr: float = 0.5,
                photo_ids_to_file: bool = False, photo_ids_froot: str =' ',
@@ -350,7 +350,7 @@ def learn_loop(nloops: int, strategy: str, path_to_features: str,
     output_queried_file: str
         Full path to output file to store the queried sample.
     features_method: str (optional)
-        Feature extraction method. Currently only 'Bazin' and 'Bump' are implemented.
+        Feature extraction method. Currently only 'bazin' and 'Bump' are implemented.
     classifier: str (optional)
         Machine Learning algorithm.
         Currently implemented options are 'RandomForest', 'GradientBoostedTrees',
diff --git a/resspect/lightcurves_utils.py b/resspect/lightcurves_utils.py
index f36ff1ef..ee5eb6a4 100644
--- a/resspect/lightcurves_utils.py
+++ b/resspect/lightcurves_utils.py
@@ -11,6 +11,7 @@
 
 import numpy as np
 import pandas as pd
+
 from resspect.snana_fits_to_pd import read_fits
 
 
@@ -549,3 +550,20 @@ def get_query_flags(light_curve_data, telescope_names: list,
         else:
             query_flags.append(False)
     return query_flags
+
+
+def get_files_list(path_to_data_dir: str,
+                    file_prefix: str) -> list:
+    """
+    loads file names available in the folder
+    Parameters
+    ----------
+    path_to_data_dir
+        folder path
+    file_prefix
+        files start name
+    """
+    files_list = os.listdir(path_to_data_dir)
+    files_list = [each_file for each_file in files_list
+                  if each_file.startswith(file_prefix)]
+    return files_list
diff --git a/resspect/scripts/fit_dataset.py b/resspect/scripts/fit_dataset.py
index 797dd923..8ff52239 100644
--- a/resspect/scripts/fit_dataset.py
+++ b/resspect/scripts/fit_dataset.py
@@ -18,8 +18,8 @@
 import argparse
 
 from resspect.fit_lightcurves import fit_snpcc
-from resspect.fit_lightcurves import fit_resspect_bazin
-from resspect.fit_lightcurves import fit_plasticc_bazin
+from resspect.fit_lightcurves import fit_resspect
+from resspect.fit_lightcurves import fit_plasticc
 
 __all__ = ['main']
 
@@ -76,17 +76,17 @@ def main(user_choices):
                   number_of_processors=ncores, function=function)
     
     elif user_choices.sim_name == 'RESSPECT':
-        fit_resspect_bazin(path_photo_file=user_choices.photo_file,
-                           path_header_file=user_choices.header_file,
-                           output_file=features_file, sample=user_choices.sample,
-                          number_of_processors=ncores)
+        fit_resspect(path_photo_file=user_choices.photo_file,
+                     path_header_file=user_choices.header_file,
+                     output_file=features_file, sample=user_choices.sample,
+                     number_of_processors=ncores)
 
     elif user_choices.sim_name == 'PLAsTiCC':
-        fit_plasticc_bazin(path_photo_file=user_choices.photo_file, 
-                           path_header_file=user_choices.header_file,
-                           output_file=features_file,
-                           sample=user_choices.sample,
-                           number_of_processors=ncores)
+        fit_plasticc(path_photo_file=user_choices.photo_file,
+                    path_header_file=user_choices.header_file,
+                    output_file=features_file,
+                    sample=user_choices.sample,
+                    number_of_processors=ncores)
 
     return None
 
diff --git a/resspect/time_domain_loop.py b/resspect/time_domain_loop.py
index 22a306f3..78097a2f 100644
--- a/resspect/time_domain_loop.py
+++ b/resspect/time_domain_loop.py
@@ -869,6 +869,19 @@ def run_time_domain_active_learning_loop(
 
     """
     learning_days = [int(each_day) for each_day in learning_days]
+    
+    # create dictionary with budgets
+    if bool(budgets):
+        if len(budgets) not in [2, len(np.arange(learning_days[0], learning_days[1]))]:
+            raise ValueError('There must be 1 budget per telescope or ' + \
+                            '1 budget per telescope per night!')
+        
+        c = 0
+        budgets_dict = {}
+        for epoch in range(learning_days[0], learning_days[-1] - 1):
+            budgets_dict[epoch] = list(budgets)[c]
+            c = c + 1
+    
     for epoch in progressbar.progressbar(
             range(learning_days[0], learning_days[-1] - 1)):
         if light_curve_data.pool_features.shape[0] > 0:
@@ -876,7 +889,7 @@ def run_time_domain_active_learning_loop(
                 light_curve_data, classifier, is_classifier_bootstrap, **kwargs)
             if light_curve_data.queryable_ids.shape[0] > 0:
                 object_indices = _get_indices_of_objects_to_be_queried(
-                    light_curve_data, strategy, budgets, is_queryable,
+                    light_curve_data, strategy, budgets_dict[epoch], is_queryable,
                     query_threshold, batch)
                 light_curve_data = _update_samples_with_object_indices(
                     light_curve_data, object_indices, is_queryable, epoch)
diff --git a/resspect/time_domain_plasticc.py b/resspect/time_domain_plasticc.py
index 65fde3f9..add93e72 100644
--- a/resspect/time_domain_plasticc.py
+++ b/resspect/time_domain_plasticc.py
@@ -27,7 +27,9 @@
 import numpy as np
 import pandas as pd
 import progressbar
-from resspect import LightCurve
+
+from resspect.feature_extractors.bazin import BazinFeatureExtractor
+from resspect.feature_extractors.bump import BumpFeatureExtractor
 from resspect.lightcurves_utils import BAZIN_HEADERS
 from resspect.lightcurves_utils import get_query_flags
 from resspect.lightcurves_utils import maybe_create_directory
@@ -35,6 +37,16 @@
 from resspect.lightcurves_utils import read_plasticc_full_photometry_data
 
 
+FEATURE_EXTRACTOR_HEADERS_MAPPING = {
+    "bazin": BAZIN_HEADERS
+}
+
+FEATURE_EXTRACTOR_MAPPING = {
+    "bazin": BazinFeatureExtractor,
+    "bump": BumpFeatureExtractor
+}
+
+
 class PLAsTiCCPhotometry:
     """
     Handles photometric information for the PLAsTiCC data.
@@ -75,7 +87,7 @@ def __init__(self,
         self._class_mappings = None  # type: dict
         self._last_file_index = None # type: int
         self._features_file_name = None  # type: str
-        self._bazin_header = None  # type: list
+        self._header = None  # type: list
         self._file_list_dict = {}  # type: dict
         self._today = None  # type: int
         self._number_of_telescopes = None  # type: int
@@ -110,30 +122,30 @@ def build(self, config='original', photo_file=None,
         else:
             self._file_list_dict[sample] = [photo_file]
 
-    def _set_bazin_header(self, get_cost: bool = False, header: str = 'Bazin'):
+    def _set_header(self, get_cost: bool = False, feature_extractor: str = 'bazin'):
         """
-        Initializes bazin header.
+        Initializes header
 
         Parameters
         ----------
         get_cost: bool (optional)
             If True, calculate cost of taking a spectra in the last
             observed photometric point. Default is False.
-        header: str (optional)
+        feature_extractor: str (optional)
             List of elements to be added to the header.
             Separate by 1 space.
             Default option uses header for Bazin features file.
         """
-        if header == 'Bazin':
-            self._bazin_header = BAZIN_HEADERS['plasticc_header']
-            if get_cost:
-                self._bazin_header = BAZIN_HEADERS[
-                    'plasticc_header_with_cost']
-        else:
-            raise ValueError('Only Bazin headers are supported')
+        if feature_extractor not in FEATURE_EXTRACTOR_HEADERS_MAPPING:
+            raise ValueError('Only bazin headers are supported')
+        self._header = FEATURE_EXTRACTOR_HEADERS_MAPPING[
+                feature_extractor]['plasticc_header']
+        if get_cost:
+            self._header = FEATURE_EXTRACTOR_HEADERS_MAPPING[
+                feature_extractor]['plasticc_header_with_cost']
 
     def create_daily_file(self, output_dir: str, day: int,
-                          header: str = 'Bazin', get_cost: bool = False):
+                          feature_extractor: str = 'bazin', get_cost: bool = False):
         """
         Create one file for a given day of the survey.
 
@@ -145,7 +157,7 @@ def create_daily_file(self, output_dir: str, day: int,
             Complete path to output directory.
         day: int
             Day passed since the beginning of the survey.
-        header: str (optional)
+        feature_extractor: str (optional)
             List of elements to be added to the header.
             Separate by 1 space.
             Default option uses header for Bazin features file.
@@ -157,8 +169,8 @@ def create_daily_file(self, output_dir: str, day: int,
         self._features_file_name = os.path.join(
             output_dir, 'day_' + str(day) + '.csv')
         with open(self._features_file_name, 'w') as features_file:
-            self._set_bazin_header(header)
-            features_file.write(','.join(self._bazin_header) + '\n')
+            self._set_header(get_cost, feature_extractor=feature_extractor)
+            features_file.write(' '.join(self._header) + '\n')
 
     def create_all_daily_files(self, output_dir:str,
                                get_cost=False):
@@ -229,13 +241,13 @@ def _verify_telescope_names(self, telescope_names: list, get_cost: bool):
         get_cost
            if cost of taking a spectra is computed
         """
-        if (('cost_' + telescope_names[0] not in self._bazin_header or
-            'cost_' + telescope_names[1] not in self._bazin_header)
+        if (('cost_' + telescope_names[0] not in self._header or
+            'cost_' + telescope_names[1] not in self._header)
                 and get_cost):
             raise ValueError('Unknown or not supported telescope names')
 
     def _load_plasticc_data(self, raw_data_dir: str, volume: int,
-                            snid: int, sample='test') -> LightCurve:
+                            snid: int, sample='test', feature_extractor: str = 'bazin'):
         """
         Loads PLAsTiCC dataset files to LightCurve class
 
@@ -251,8 +263,7 @@ def _load_plasticc_data(self, raw_data_dir: str, volume: int,
         sample: str (optional)
             Sample to load, 'train' or 'test'. Default is 'test'.
         """
-        light_curve_data = LightCurve()
-
+        light_curve_data = FEATURE_EXTRACTOR_MAPPING[feature_extractor]()
         if sample == 'test':
             file_name = os.path.join(
                     raw_data_dir, self._file_list_dict['test'][volume - 1])
@@ -265,7 +276,7 @@ def _load_plasticc_data(self, raw_data_dir: str, volume: int,
         return light_curve_data
 
     def _get_fit_days(self, day: Union[int, None],
-                      light_curve_data: LightCurve) -> list:
+                      light_curve_data) -> list:
         """
         Parameters
         ----------
@@ -284,7 +295,7 @@ def _get_fit_days(self, day: Union[int, None],
         return [day]
 
     @staticmethod
-    def _verify_features_method(feature_method: str):
+    def _verify_features_method(feature_extractor: str):
         """
         Verifies if valid dataset name and features method is passed
         Parameters
@@ -292,10 +303,11 @@ def _verify_features_method(feature_method: str):
         feature_method
             Feature extraction method, only possibility is 'Bazin'.
         """
-        if feature_method != 'Bazin':
-            raise ValueError('Only Bazin features are implemented!!')
+        if feature_extractor not in FEATURE_EXTRACTOR_MAPPING:
+            raise ValueError(
+                'Provided feature extractor method has not been implemented!!')
 
-    def _check_queryable(self, light_curve_data: LightCurve,
+    def _check_queryable(self, light_curve_data,
                          queryable_criteria: int,
                          days_since_last_observation: int) -> bool:
         """
@@ -322,9 +334,9 @@ def _check_queryable(self, light_curve_data: LightCurve,
             days_since_last_obs=days_since_last_observation)
 
     def _update_queryable_if_get_cost(
-            self, light_curve_data: LightCurve, telescope_names: list,
+            self, light_curve_data, telescope_names: list,
             telescope_sizes: list, spectroscopic_snr: int,
-            **kwargs) -> LightCurve:
+            **kwargs):
         """
         Updates time required to take a spectra calc_exp_time() method of
         LightCurve class
@@ -356,7 +368,7 @@ def _update_queryable_if_get_cost(
         return light_curve_data
 
     def _process_current_day(
-            self, light_curve_data_day: LightCurve, queryable_criteria: int,
+            self, light_curve_data_day, queryable_criteria: int,
             days_since_last_observation: int, telescope_names: list,
             telescope_sizes: list, spectroscopic_snr: int,
             min_available_points: int = 4, **kwargs):
@@ -395,10 +407,10 @@ def _process_current_day(
         if np.sum(photo_flag) > min_available_points:
             light_curve_data_day.photometry = light_curve_data_day.photometry[
                 photo_flag]
-            light_curve_data_day.fit_bazin_all()
+            light_curve_data_day.fit_all()
 
-            if (len(light_curve_data_day.bazin_features) > 0 and
-                    'None' not in light_curve_data_day.bazin_features):
+            if (len(light_curve_data_day.features) > 0 and
+                    'None' not in light_curve_data_day.features):
 
                 light_curve_data_day.queryable = self._check_queryable(
                     light_curve_data_day, queryable_criteria,
@@ -415,7 +427,7 @@ def _process_current_day(
         return None
 
     def _get_features_to_write(
-            self, light_curve_data_day: LightCurve,
+            self, light_curve_data_day,
             get_cost: bool, telescope_names: list) -> list:
         """
         Returns features list to write
@@ -444,12 +456,12 @@ def _get_features_to_write(
             for index in range(self._number_of_telescopes):
                 features_list.append(str(
                     light_curve_data_day.exp_time[telescope_names[index]]))
-        features_list.extend(light_curve_data_day.bazin_features)
+        features_list.extend(light_curve_data_day.features)
 
         return features_list
 
-    def _update_light_curve_meta_data(self, light_curve_data_day: LightCurve,
-                                      snid: int) -> Union[LightCurve, None]:
+    def _update_light_curve_meta_data(self, light_curve_data_day,
+                                      snid: int):
         """
         Loads light curve data of the given SNID
 
@@ -482,7 +494,7 @@ def _update_light_curve_meta_data(self, light_curve_data_day: LightCurve,
     def fit_one_lc(self, raw_data_dir: str, snid: int, output_dir: str,
                    vol=None, day=None, queryable_criteria=1,
                    days_since_last_obs=2, get_cost=False, tel_sizes=[4, 8],
-                   tel_names=['4m', '8m'], feature_method='Bazin', spec_SNR=10,
+                   tel_names=['4m', '8m'], feature_extractor='bazin', spec_SNR=10,
                    time_window=[0, 1095], sample='test', bar=False,
                    **kwargs):
         """
@@ -529,7 +541,7 @@ def fit_one_lc(self, raw_data_dir: str, snid: int, output_dir: str,
             Primary mirrors diameters of potential spectroscopic telescopes.
             Only used if "get_cost == True".
             Default is [4, 8].
-        feature_method: str (optional)
+        feature_extractor: str (optional)
             Feature extraction method.
             Only possibility is 'Bazin'.
         spec_SNR: float (optional)
@@ -545,10 +557,11 @@ def fit_one_lc(self, raw_data_dir: str, snid: int, output_dir: str,
             Any input required by ExpTimeCalc.findexptime function.
         """
         self._verify_telescope_names(tel_names, get_cost)
-        self._verify_features_method(feature_method)
+        self._verify_features_method(feature_extractor)
         self._number_of_telescopes = len(tel_names)
 
-        light_curve_data = self._load_plasticc_data(raw_data_dir, vol, snid, sample)
+        light_curve_data = self._load_plasticc_data(
+            raw_data_dir, vol, snid, sample, feature_extractor)
         fit_days = np.arange(time_window[0], time_window[1])
 
         if bar:
@@ -610,7 +623,7 @@ def _maybe_create_feature_to_file(self, features_file_name: str):
         """
         if not os.path.isfile(features_file_name):
             with open(features_file_name, 'w') as features_file:
-                features_file.write(','.join(self._bazin_header) + '\n')
+                features_file.write(' '.join(self._header) + '\n')
 
     def _maybe_create_daily_feature_files(
             self, time_window: list, output_dir: str, get_cost: bool, 
@@ -651,7 +664,7 @@ def fit_all_snids_lc(
             vol: int = None, queryable_criteria: int = 1,
             days_since_last_obs: int = 2, get_cost: bool = False,
             tel_sizes: list = [4, 8], tel_names: list = ['4m', '8m'],
-            feature_method: str = 'Bazin', spec_SNR: int = 10,
+            feature_extractor: str = 'bazin', spec_SNR: int = 10,
             time_window: list = [0, 1095], sample: str = 'test',
             number_of_processors: int = 1, create_daily_files: bool = False,
             ask_user: bool = True, **kwargs):
@@ -695,9 +708,8 @@ def fit_all_snids_lc(
             Names of the telescopes under consideration for spectroscopy.
             Only used if "get_cost == True".
             Default is ["4m", "8m"].
-        feature_method
+        feature_extractor
             Feature extraction method.
-            Only possibility is 'Bazin'.
         spec_SNR
             SNR required for spectroscopic follow-up. Default is 10.
         time_window
@@ -715,9 +727,9 @@ def fit_all_snids_lc(
         kwargs
             Any input required by ExpTimeCalc.findexptime function.
         """
-        self._set_bazin_header(get_cost=get_cost)
+        self._set_header(get_cost=get_cost)
         self._verify_telescope_names(tel_names, get_cost)
-        self._verify_features_method(feature_method)
+        self._verify_features_method(feature_extractor)
         self._number_of_telescopes = len(tel_names)
         self._kwargs = kwargs
         if create_daily_files:
@@ -755,12 +767,12 @@ def fit_all_snids_lc(
                 json.dump(previous_day_number_of_points, json_file)
 
     def _process_each_snid(
-            self, snid: int, raw_data_dir: str, vol: str, day_of_survey: int,
+            self, snid: int, raw_data_dir: str, vol: int, day_of_survey: int,
             sample: str, get_cost, telescope_names: list,
             queryable_criteria: int, days_since_last_obs: int,
             telescope_sizes: list, spectroscopic_snr: int,
-            previous_day_number_of_points: dict) -> Tuple[
-        Union[str, None], Union[int, None]]:
+            previous_day_number_of_points: dict) -> Tuple[Union[str, None],
+        Union[int, None]]:
         """
         Fits snid light curve for the day of survey and returns features
 
@@ -874,8 +886,10 @@ def _load_previous_day_features(
         return None, {}
     with open(previous_day_file_name, 'r') as file:
         previous_day_features = file.readlines()
-    previous_day_index_mapping = {line.split(",", 1)[0]: index for index, line
-                          in enumerate(previous_day_features)}
+    previous_day_index_mapping = {
+        line.split(" ", 1)[0]: index for index, line in enumerate(
+            previous_day_features)}
+
     return previous_day_features, previous_day_index_mapping
 
 
diff --git a/resspect/time_domain_snpcc.py b/resspect/time_domain_snpcc.py
index ed27294c..24c8fc58 100644
--- a/resspect/time_domain_snpcc.py
+++ b/resspect/time_domain_snpcc.py
@@ -19,11 +19,11 @@
 import multiprocessing
 import os
 from itertools import repeat
-from typing import Union
 
-import progressbar
-from resspect import LightCurve
+from resspect.feature_extractors.bazin import BazinFeatureExtractor
+from resspect.feature_extractors.bump import BumpFeatureExtractor
 from resspect.lightcurves_utils import BAZIN_HEADERS
+from resspect.lightcurves_utils import get_files_list
 from resspect.lightcurves_utils import get_query_flags
 from resspect.lightcurves_utils import maybe_create_directory
 
@@ -32,6 +32,17 @@
 __all__ = ['SNPCCPhotometry']
 
 
+FEATURE_EXTRACTOR_MAPPING = {
+    "bazin": BazinFeatureExtractor,
+    "bump": BumpFeatureExtractor
+}
+
+
+FEATURE_EXTRACTOR_HEADERS_MAPPING = {
+    "bazin": BAZIN_HEADERS
+}
+
+
 class SNPCCPhotometry:
     """
     Handles photometric information for entire SNPCC data.
@@ -70,46 +81,14 @@ def __init__(self,
         self.max_epoch = max_epoch
         self.min_epoch = min_epoch
         self.rmag_lim = rmag_lim
-        self._bazin_header = None  # type: list
+        self._header = None  # type: list
         self._features_file_name = None  # type: str
         self._today = None  # type: int
         self._number_of_telescopes = None  # type: int
         self._kwargs = None  # type: dict
 
-    def get_lim_mjds(self, raw_data_dir: str) -> list:
-        """
-        Get minimum and maximum MJD for complete sample.
-
-        This function is not necessary if you are working with
-        SNPCC data. The values are hard coded in the class.
-
-        Parameters
-        ----------
-        raw_data_dir: str
-            Complete path to raw data directory.
-
-        Returns
-        -------
-        limits: list
-            List of extreme MJDs for entire sample: [min_MJD, max_MJD].
-        """
-        files_list = _get_files_list(raw_data_dir, 'DES_SN')
-        # store MJDs
-        min_day = []
-        max_day = []
-
-        for each_file in progressbar.progressbar(files_list):
-            light_curve_data = LightCurve()
-            light_curve_data.load_snpcc_lc(os.path.join(raw_data_dir, each_file))
-            min_day.append(min(light_curve_data.photometry['mjd'].values))
-            max_day.append(max(light_curve_data.photometry['mjd'].values))
-            self.min_epoch = min(min_day)
-            self.max_epoch = max(max_day)
-
-        return [min(min_day), max(max_day)]
-
     def create_daily_file(self, output_dir: str,
-                          day: int, header: str = 'Bazin',
+                          day: int, feature_extractor: str = 'bazin',
                           get_cost: bool = False):
         """
         Create one file for a given day of the survey.
@@ -125,23 +104,23 @@ def create_daily_file(self, output_dir: str,
         get_cost: bool (optional)
             If True, calculate cost of taking a spectra in the last
             observed photometric point. Default is False.
-        header: str (optional)
-            List of elements to be added to the header.
-            Separate by 1 space.
-            Default option uses header for Bazin features file.
+        feature_extractor: str
+            Feature extraction method, only possibility is 'Bazin'.
         """
         maybe_create_directory(output_dir)
         self._features_file_name = os.path.join(
             output_dir, 'day_' + str(day) + '.dat')
         logging.info('Creating features file')
         with open(self._features_file_name, 'w') as features_file:
-            if header == 'Bazin':
-                self._bazin_header = BAZIN_HEADERS['snpcc_header']
-                if get_cost:
-                    self._bazin_header = BAZIN_HEADERS['snpcc_header_with_cost']
-            else:
+            if feature_extractor not in FEATURE_EXTRACTOR_HEADERS_MAPPING:
                 raise ValueError('Only Bazin headers are supported')
-            features_file.write(' '.join(self._bazin_header) + '\n')
+            self._header = FEATURE_EXTRACTOR_HEADERS_MAPPING[
+                feature_extractor]['snpcc_header']
+            if get_cost:
+                self._header = FEATURE_EXTRACTOR_HEADERS_MAPPING[
+                    feature_extractor]['snpcc_header_with_cost']
+
+            features_file.write(' '.join(self._header) + '\n')
 
     def _verify_telescope_names(self, telescope_names: list, get_cost: bool):
         """
@@ -155,14 +134,14 @@ def _verify_telescope_names(self, telescope_names: list, get_cost: bool):
            if cost of taking a spectra is computed
         """
 
-        if (('cost_' + telescope_names[0] not in self._bazin_header or
-            'cost_' + telescope_names[1] not in self._bazin_header)
+        if (('cost_' + telescope_names[0] not in self._header or
+            'cost_' + telescope_names[1] not in self._header)
                 and get_cost):
 
             raise ValueError('Unknown or not supported telescope names')
 
     def _maybe_create_features_file(self, output_dir: str, day_of_survey: int,
-                                    feature_method: str, get_cost: bool):
+                                    feature_extractor: str, get_cost: bool):
         """
         Creates features output file if not available
 
@@ -172,7 +151,7 @@ def _maybe_create_features_file(self, output_dir: str, day_of_survey: int,
             output directory path to save features file
         day_of_survey: int
             Day since the beginning of survey.
-        feature_method: str
+        feature_extractor: str
             Feature extraction method, only possibility is 'Bazin'.
         get_cost: bool
            if True, cost of taking a spectra is computed.
@@ -180,26 +159,27 @@ def _maybe_create_features_file(self, output_dir: str, day_of_survey: int,
         if not os.path.isfile(self._features_file_name):
             logging.info('Features file doesnt exist')
             self.create_daily_file(
-                output_dir, day_of_survey, feature_method, get_cost)
+                output_dir, day_of_survey, feature_extractor, get_cost)
 
     @staticmethod
     def _verify_dataset_and_features_method(dataset_name: str,
-                                            feature_method: str):
+                                            feature_extractor: str):
         """
         Verifies if valid dataset name and features method is passed
         Parameters
         ----------
         dataset_name: str
             name of the dataset used
-        feature_method: str
+        feature_extractor: str
             Feature extraction method, only possibility is 'Bazin'.
         """
         if dataset_name != 'SNPCC':
             raise ValueError('This class supports only SNPCC dataset!')
-        if feature_method != 'Bazin':
-            raise ValueError('Only Bazin features are implemented!!')
+        # TODO: Update when bump headers are available
+        if feature_extractor != 'bazin':
+            raise ValueError('Only bazin features are implemented!!')
 
-    def _check_queryable(self, light_curve_data: LightCurve,
+    def _check_queryable(self, light_curve_data,
                          queryable_criteria: int,
                          days_since_last_observation: int) -> bool:
         """
@@ -207,7 +187,7 @@ def _check_queryable(self, light_curve_data: LightCurve,
 
         Parameters
         ----------
-        light_curve_data: resspect.LightCurve
+        light_curve_data:
             An instance of LightCurve class
         queryable_criteria: [1 or 2]
             Criteria to determine if an obj is queryable.
@@ -226,9 +206,9 @@ def _check_queryable(self, light_curve_data: LightCurve,
             days_since_last_obs=days_since_last_observation)
 
     def _update_queryable_if_get_cost(
-            self, light_curve_data: LightCurve, telescope_names: list,
+            self, light_curve_data, telescope_names: list,
             telescope_sizes: list, spectroscopic_snr: int,
-            kwargs: dict) -> LightCurve:
+            kwargs: dict):
         """
         Updates time required tp take a spectra calc_exp_time() method of
         LightCurve class
@@ -259,16 +239,16 @@ def _update_queryable_if_get_cost(
         return light_curve_data
 
     def _process_each_light_curve(
-            self, light_curve_data: LightCurve, queryable_criteria: int,
+            self, light_curve_data, queryable_criteria: int,
             days_since_last_observation: int, telescope_names: list,
             telescope_sizes: list, spectroscopic_snr: int, kwargs: dict,
-            min_available_points: int = 5) -> Union[LightCurve, None]:
+            min_available_points: int = 5):
         """
         Processes each light curve files.
 
         Parameters
         ----------
-        light_curve_data: resspect.LightCurve
+        light_curve_data:
             An instance of LightCurve class
         queryable_criteria: int
             Criteria to determine if an obj is queryable.
@@ -296,10 +276,10 @@ def _process_each_light_curve(
         if sum(photo_flag) >= min_available_points:
             light_curve_data.photometry = light_curve_data.photometry[
                 photo_flag]
-            light_curve_data.fit_bazin_all()
+            light_curve_data.fit_all()
 
-            if (len(light_curve_data.bazin_features) > 0 and
-                    'None' not in light_curve_data.bazin_features):
+            if (len(light_curve_data.features) > 0 and
+                    'None' not in light_curve_data.features):
                 light_curve_data.queryable = self._check_queryable(
                     light_curve_data, queryable_criteria,
                     days_since_last_observation)
@@ -312,14 +292,14 @@ def _process_each_light_curve(
                 return light_curve_data
         return None
 
-    def _get_features_to_write(self, light_curve_data: LightCurve,
+    def _get_features_to_write(self, light_curve_data,
                                get_cost: bool, telescope_names: list) -> list:
         """
         Returns features list to write.
 
         Parameters
         ----------
-        light_curve_data: resspect.LightCurve
+        light_curve_data:
             fitted light curve data
         get_cost: bool
            if cost of taking a spectra is computed
@@ -336,13 +316,13 @@ def _get_features_to_write(self, light_curve_data: LightCurve,
             for index in range(self._number_of_telescopes):
                 features_list.append(str(
                     light_curve_data.exp_time[telescope_names[index]]))
-        features_list.extend(light_curve_data.bazin_features)
+        features_list.extend(light_curve_data.features)
         return features_list
 
     def _get_current_sample_features(
             self, file_name: str, raw_data_dir: str, queryable_criteria: int,
             days_since_observation: int, telescope_names: list,
-            telescope_sizes: list, spec_SNR: int) -> LightCurve:
+            telescope_sizes: list, spec_SNR: int, feature_extractor: str):
         """
         Reads a SNPCC file and updates time domain features
 
@@ -373,8 +353,10 @@ def _get_current_sample_features(
             Default is [4, 8].
         spec_SNR: float (optional)
             SNR required for spectroscopic follow-up. Default is 10.
+        feature_extractor: str
+            Feature extraction method
         """
-        light_curve_data = LightCurve()
+        light_curve_data = FEATURE_EXTRACTOR_MAPPING[feature_extractor]()
         light_curve_data.load_snpcc_lc(
             os.path.join(raw_data_dir, file_name))
         light_curve_data = self._process_each_light_curve(
@@ -384,14 +366,14 @@ def _get_current_sample_features(
 
     # TODO: Too many arguments. Refactor and update docs
     def build_one_epoch(self, raw_data_dir: str, day_of_survey: int,
-                        time_domain_dir: str, feature_method: str = 'Bazin',
+                        time_domain_dir: str, feature_extractor: str = 'bazin',
                         dataset: str = 'SNPCC', days_since_obs: int = 2,
                         queryable_criteria: int = 1, get_cost: bool = False,
                         tel_sizes: list = [4, 8],
                         tel_names: list = ['4m', '8m'], spec_SNR: int = 10,
                         number_of_processors: int = 1, **kwargs):
         """
-        Fit bazin for all objects with enough points in a given day.
+        Fit features for all objects with enough points in a given day.
 
         Generate 1 file containing best-fit Bazin parameters for a given
         day of the survey.
@@ -411,7 +393,7 @@ def build_one_epoch(self, raw_data_dir: str, day_of_survey: int,
             Day since last observation to consider for spectroscopic
             follow-up without the need to extrapolate light curve.
             Only used if "queryable_criteria == 2". Default is 2.
-        feature_method: str (optional)
+        feature_extractor: str (optional)
             Feature extraction method.
             Only possibility is 'Bazin'.
         get_cost: bool (optional)
@@ -437,14 +419,14 @@ def build_one_epoch(self, raw_data_dir: str, day_of_survey: int,
         number_of_processors: int, default 1
             Number of cpu processes to use
         kwargs: extra parameters
-            Any input required by ExpTimeCalc.findexptime function.
+            Any input required by ExpTimeCalc.find exptime function.
         """
         self._kwargs = kwargs
-        self._verify_dataset_and_features_method(dataset, feature_method)
+        self._verify_dataset_and_features_method(dataset, feature_extractor)
         self._verify_telescope_names(tel_names, get_cost)
         self._maybe_create_features_file(
-            time_domain_dir, day_of_survey, feature_method, get_cost)
-        files_list = _get_files_list(raw_data_dir, 'DES_SN')
+            time_domain_dir, day_of_survey, feature_extractor, get_cost)
+        files_list = get_files_list(raw_data_dir, 'DES_SN')
         self._today = day_of_survey + self.min_epoch
         self._number_of_telescopes = len(tel_names)
 
@@ -454,7 +436,7 @@ def build_one_epoch(self, raw_data_dir: str, day_of_survey: int,
             iterator_list = zip(
                 files_list, repeat(raw_data_dir), repeat(queryable_criteria),
                 repeat(days_since_obs), repeat(tel_names), repeat(tel_sizes),
-                repeat(spec_SNR))
+                repeat(spec_SNR), repeat(feature_extractor))
             for light_curve_data in multi_process.starmap(
                     self._get_current_sample_features, iterator_list):
                 if light_curve_data is not None:
@@ -466,23 +448,6 @@ def build_one_epoch(self, raw_data_dir: str, day_of_survey: int,
         logging.info("Features have been saved to: %s", self._features_file_name)
 
 
-def _get_files_list(path_to_data_dir: str,
-                    file_prefix: str) -> list:
-    """
-    loads file names available in the folder
-    Parameters
-    ----------
-    path_to_data_dir
-        folder path
-    file_prefix
-        files start name
-    """
-    files_list = os.listdir(path_to_data_dir)
-    files_list = [each_file for each_file in files_list
-                  if each_file.startswith(file_prefix)]
-    return files_list
-
-
 def main():
     return None
 


Test patch:
diff --git a/resspect/tests/test_fit_lightcurves.py b/resspect/tests/test_fit_lightcurves.py
index 1115af44..84778390 100644
--- a/resspect/tests/test_fit_lightcurves.py
+++ b/resspect/tests/test_fit_lightcurves.py
@@ -85,7 +85,7 @@ def test_conv_flux_mag(input_lc):
 def test_check_queryable(input_lc):
     """ Test consistency of queryable tests. """
     
-    input_lc.fit_bazin_all()
+    input_lc.fit_all()
     
     min_mjd = min(input_lc.photometry['mjd'].values) 
     epochs = [15, 80]
@@ -151,8 +151,8 @@ def test_fit_bazin(input_lc):
 def test_fit_bazin_all(input_lc):
     """ Test Bazin fit in all filters. """
     
-    input_lc.fit_bazin_all()
-    l1 = len(input_lc.bazin_features)
+    input_lc.fit_all()
+    l1 = len(input_lc.features)
     l2 = len(input_lc.filters) * 5
     
     assert l1 == l2
@@ -167,7 +167,7 @@ def test_evaluate_bazin(input_lc):
     # generate a grid of times
     t = np.arange(min_mjd, max_mjd, 0.5)
     
-    input_lc.fit_bazin_all()
+    input_lc.fit_all()
     flux = input_lc.evaluate_bazin(time=np.random.choice(t, size=5))
     
     res = []
diff --git a/resspect/tests/test_learn_loop.py b/resspect/tests/test_learn_loop.py
index a74258d4..049dc0f9 100644
--- a/resspect/tests/test_learn_loop.py
+++ b/resspect/tests/test_learn_loop.py
@@ -10,7 +10,7 @@ def test_can_run_learn_loop(extract_feature):
     from resspect.learn_loop import learn_loop
 
     learn_loop(nloops=1,
-               features_method="Bazin",
+               features_method="bazin",
                strategy="RandomSampling",
                path_to_features=extract_feature,
                output_metrics_file="just_a_name.csv",
@@ -19,12 +19,12 @@ def test_can_run_learn_loop(extract_feature):
 
 @pytest.fixture(scope="function")
 def extract_feature(path_to_test_data):
-    from resspect import fit_snpcc_bazin
+    from resspect import fit_snpcc
 
     path_to_data_dir = path_to_test_data
     output_file = 'output_file.dat'
 
-    fit_snpcc_bazin(path_to_data_dir=path_to_data_dir, features_file=output_file)
+    fit_snpcc(path_to_data_dir=path_to_data_dir, features_file=output_file)
 
     return output_file