redback 1.0.3__py3-none-any.whl → 1.1__py3-none-any.whl

redback/__init__.py

@@ -1,8 +1,9 @@
 from redback import analysis, constants, get_data, redback_errors, priors, result, sampler, transient, \
-    transient_models, utils, photosphere, sed, interaction_processes, constraints, plotting, model_library, simulate_transients
+    transient_models, utils, photosphere, sed, interaction_processes, constraints, plotting, model_library, \
+    simulate_transients
 from redback.transient import afterglow, kilonova, prompt, supernova, tde
 from redback.sampler import fit_model
 from redback.utils import setup_logger
 
-__version__ = "1.0.2"
+__version__ = "1.1.0"
 setup_logger(log_level='info')

redback/analysis.py

@@ -4,7 +4,7 @@ import numpy as np
 import pandas as pd
 
 import redback.model_library
-from redback.utils import logger, find_nearest
+from redback.utils import logger, find_nearest, bands_to_frequency
 from redback.result import RedbackResult
 from redback.constants import day_to_s
 import matplotlib
@@ -97,9 +97,9 @@ def plot_evolution_parameters(result, random_models=100):
         nn = output.nn
         mu = output.mu
         alpha = output.alpha
-        ax[0].plot(time, nn, "--", lw=1, color='red', alpha=2.5, zorder=-1)
-        ax[1].plot(time, np.rad2deg(alpha), "--", lw=1, color='red', alpha=2.5, zorder=-1)
-        ax[2].plot(time, mu, "--", lw=1, color='red', alpha=2.5, zorder=-1)
+        ax[0].plot(time, nn, "--", lw=1, color='red', zorder=-1)
+        ax[1].plot(time, np.rad2deg(alpha), "--", lw=1, color='red', zorder=-1)
+        ax[2].plot(time, mu, "--", lw=1, color='red', zorder=-1)
         ax[0].set_ylabel('braking index')
         ax[1].set_ylabel('inclination angle')
         ax[2].set_ylabel('magnetic moment')
@@ -151,3 +151,320 @@ def plot_spectrum(model, parameters, time_to_plot, axes=None, **kwargs):
     ax.set_ylabel(r'Flux ($10^{-17}$ erg s$^{-1}$ cm$^{-2}$ $\mathrm{\AA}$)')
     ax.legend(loc='upper left')
     return ax
+
+def plot_gp_lightcurves(transient, gp_output, axes=None, band_colors=None, band_scaling=None):
+    """
+    Plot the Gaussian Process lightcurves
+
+    :param transient: A transient object
+    :param gp_output: The output of the fit_gp function
+    :param axes: axes, ideally you should be passing the axes from the plot_data methods
+    :param band_colors: a dictionary of band colors; again ideally you should be passing the band_colors from the plot_data methods
+    :return: axes object with the GP lightcurves plotted
+    """
+    ax = axes or plt.gca()
+
+    if transient.use_phase_model:
+        ref_date = transient.x[0]
+    else:
+        ref_date = 0
+
+    t_new = np.linspace(transient.x.min() - 10, transient.x.max() + 20, 100)
+
+    if transient.data_mode in ['flux_density', 'flux', 'magnitude']:
+        if band_colors is None:
+            band_colors = dict(zip(transient.unique_bands, plt.cm.tab20(range(len(transient.unique_bands)))))
+        else:
+            band_colors = band_colors
+        if gp_output.use_frequency:
+            for band in transient.unique_bands:
+                if band_scaling:
+                    scaling = band_scaling[band]
+                else:
+                    scaling = 0
+                f_new = np.ones_like(t_new) * bands_to_frequency([band])
+                X_new = np.column_stack((f_new, t_new))
+                gp = gp_output.gp
+                y_pred, y_cov = gp.predict(gp_output.scaled_y, X_new, return_cov=True)
+                y_std = np.sqrt(np.diag(y_cov))
+                y_lower = y_pred - 0.5 * y_std
+                y_upper = y_pred + 0.5 * y_std
+                ax.plot(t_new - ref_date, (y_pred * gp_output.y_scaler) + scaling, color=band_colors[band])
+                ax.fill_between(t_new - ref_date, (y_lower * gp_output.y_scaler) + scaling,
+                                (y_upper * gp_output.y_scaler) + scaling, alpha=0.5,
+                                color=band_colors[band])
+        else:
+            for band in transient.unique_bands:
+                if band_scaling:
+                    scaling = band_scaling[band]
+                else:
+                    scaling = 0
+                gp = gp_output.gp[band]
+                y_pred, y_cov = gp.predict(gp_output.scaled_y[band], t_new, return_cov=True)
+                y_std = np.sqrt(np.diag(y_cov))
+                y_lower = y_pred - 0.5 * y_std
+                y_upper = y_pred + 0.5 * y_std
+                ax.plot(t_new - ref_date, (y_pred * gp_output.y_scaler) + scaling, color=band_colors[band])
+                ax.fill_between(t_new - ref_date, (y_lower * gp_output.y_scaler) + scaling,
+                                (y_upper * gp_output.y_scaler) + scaling, alpha=0.5,
+                                color=band_colors[band])
+    else:
+        y_pred, y_cov = gp_output.gp.predict(gp_output.scaled_y, t_new, return_cov=True)
+        y_std = np.sqrt(np.diag(y_cov))
+        y_lower = y_pred - 0.5 * y_std
+        y_upper = y_pred + 0.5 * y_std
+
+        ax.plot(t_new, y_pred * gp_output.y_scaler, color='red')
+        ax.fill_between(t_new, y_lower * gp_output.y_scaler, y_upper * gp_output.y_scaler, alpha=0.5, color='red')
+    return ax
+
+def fit_temperature_and_radius_gp(data, kernelT, kernelR, plot=False, **kwargs):
+    """
+    Fit a Gaussian Process to the temperature and radius data
+
+    :param data: DataFrame containing the temperature and radius data output of the transient.estimate_bb_params method.
+    :param kernelT: george kernel for the temperature
+    :param kernelR: george kernel for the radius
+    :param plot: Whether to make a two-panel plot of the temperature and radius GP evolution and the data
+    :param kwargs: Additional keyword arguments
+    :param inflate_errors: If True, inflate the errors by 20%, default is False
+    :return: Temperature and radius GP objects and plot fig and axes if requested
+    """
+    import george
+    from scipy.optimize import minimize
+
+    temperature = data['temperature']
+    radius = data['radius']
+    t_data = data['epoch_times']
+    T_err = data['temp_err']
+    R_err = data['radius_err']
+    inflate_errors = kwargs.get('inflate_errors', True)
+    if inflate_errors:
+        error = kwargs.get('error', 1.5)
+    else:
+        error = 1
+    gp_T_err_raw = T_err * error
+    gp_R_err = R_err * error
+
+    fit_in_log = kwargs.get("fit_in_log", False)
+    if fit_in_log:
+        # In log space, use: log10(T); propagate errors via: δ(log10T)=δT/(T*ln(10))
+        temperature_fit = np.log10(temperature)
+        gp_T_err = gp_T_err_raw / (temperature * np.log(10))
+    else:
+        temperature_fit = temperature
+        gp_T_err = gp_T_err_raw
+
+    gp_T = george.GP(kernelT)
+    gp_T.compute(t_data, gp_T_err + 1e-8)
+
+    def neg_ln_like_T(p):
+        gp_T.set_parameter_vector(p)
+        return -gp_T.log_likelihood(temperature_fit)
+
+    def grad_neg_ln_like_T(p):
+        gp_T.set_parameter_vector(p)
+        return -gp_T.grad_log_likelihood(temperature_fit)
+
+    p0_T = gp_T.get_parameter_vector()
+    result_T = minimize(neg_ln_like_T, p0_T, jac=grad_neg_ln_like_T)
+    gp_T.set_parameter_vector(result_T.x)
+
+    logger.info("Finished GP fit for temperature")
+    logger.info(f"GP final parameters: {gp_T.get_parameter_dict()}")
+
+    gp_R = george.GP(kernelR)
+    gp_R.compute(t_data, gp_R_err + 1e-8)
+
+    def neg_ln_like_R(p):
+        gp_R.set_parameter_vector(p)
+        return -gp_R.log_likelihood(radius)
+
+    def grad_neg_ln_like_R(p):
+        gp_R.set_parameter_vector(p)
+        return -gp_R.grad_log_likelihood(radius)
+
+    p0_R = gp_R.get_parameter_vector()
+    result_R = minimize(neg_ln_like_R, p0_R, jac=grad_neg_ln_like_R)
+    gp_R.set_parameter_vector(result_R.x)
+
+    logger.info("Finished GP fit for radius")
+    logger.info(f"GP final parameters: {gp_R.get_parameter_dict()}")
+
+    if plot:
+        sigma_to_plot = kwargs.get('sigma_to_plot', 1)
+        label = r"${}\sigma$ GP uncertainty".format(str(int(sigma_to_plot)))
+        t_pred = np.linspace(t_data.min(), t_data.max(), 100)
+        # Temperature prediction
+        T_pred, T_pred_var = gp_T.predict(temperature_fit, t_pred, return_var=True)
+        T_pred_std = np.sqrt(T_pred_var)
+
+        # If fitting in log space, convert the prediction back to linear units.
+        if fit_in_log:
+            T_pred_lin = 10**T_pred
+            # Propagate the uncertainty approximately: dT ≈ 10^x * ln(10) * sigma_x.
+            T_pred_std_lin = 10**T_pred * np.log(10) * T_pred_std
+        else:
+            T_pred_lin = T_pred
+            T_pred_std_lin = T_pred_std
+
+        # Radius prediction
+        R_pred, R_pred_var = gp_R.predict(radius, t_pred, return_var=True)
+        R_pred_std = np.sqrt(R_pred_var)
+
+        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(7, 3))
+        ax1.errorbar(t_data, temperature, yerr=T_err, fmt='o', label='Data', color='blue')
+        ax1.plot(t_pred, T_pred_lin, label='GP Prediction', color='red')
+        ax1.fill_between(t_pred, T_pred_lin - sigma_to_plot*T_pred_std_lin, T_pred_lin + sigma_to_plot*T_pred_std_lin,
+                         alpha=0.2, color='red', label=label)
+        ax2.errorbar(t_data, radius, yerr=R_err, fmt='o', label='Data', color='blue')
+        ax2.plot(t_pred, R_pred, label='GP Prediction', color='red')
+        ax2.fill_between(t_pred, R_pred - sigma_to_plot*R_pred_std, R_pred + sigma_to_plot*R_pred_std, alpha=0.2, color='red',
+                         label=label)
+
+        ax1.set_xlabel("Time", fontsize=15)
+        ax1.set_ylabel("Temperature [K]", fontsize=15)
+        ax1.set_title("Temperature Evolution", fontsize=15)
+        ax2.set_xlabel("Time", fontsize=15)
+        ax2.set_ylabel("Radius [cm]", fontsize=15)
+        ax2.set_title("Radius Evolution from GP", fontsize=15)
+
+        ax1.set_yscale('log')
+        ax2.set_yscale('log')
+
+        ax1.legend()
+        ax2.legend()
+        plt.subplots_adjust(wspace=0.3)
+        return gp_T, gp_R, fig, (ax1, ax2)
+    else:
+        return gp_T, gp_R
+
+def generate_new_transient_data_from_gp(gp_out, t_new, transient, **kwargs):
+    """
+    Generates new transient data based on Gaussian Process (GP) predictions for the given time array
+    and transient object. Depending on the data mode of the transient object
+    (e.g., 'flux_density', 'flux', 'magnitude', or 'luminosity'), this function updates the data
+    accordingly, adjusting errors and scaling by frequency if necessary.
+
+    :param gp_out: The GP output object containing the Gaussian Process model, scaled data,
+                   and other related attributes.
+    :type gp_out: object
+    :param t_new: Array of new time values for which GP predictions are to be generated.
+    :type t_new: array-like
+    :param transient: The transient object containing the original observation data and related
+                      properties such as data mode and unique frequencies or bands.
+    :type transient: object
+    :param kwargs: Additional parameters to modify behavior, such as:
+
+                   - **inflate_y_err** (bool): Flag to indicate whether to inflate GP errors.
+                   - **error** (float): Multiplier for adjusting GP error inflation.
+
+    :return: A new transient object with data updated using GP predictions.
+    :rtype: object
+    """
+    data_mode = transient.data_mode
+    logger.info(f"Data mode: {data_mode}")
+    logger.info("Creating new {} data".format(data_mode))
+
+    if data_mode not in ['flux_density', 'flux', 'magnitude', 'luminosity']:
+        raise ValueError("Data mode {} not understood".format(data_mode))
+
+    if kwargs.get('inflate_y_err', True):
+        error = kwargs.get('error', 10)
+    else:
+        logger.info("Using GP predicted errors, this is likely being too conservative")
+        error = 1.
+
+    if gp_out.use_frequency:
+        logger.info("GP is a 2D kernel with effective frequency")
+        freqs = transient.unique_frequencies
+        T, F = np.meshgrid(t_new, freqs)
+        try:
+            bands = redback.utils.frequency_to_bandname(F.flatten())
+        except Exception:
+            bands = F.flatten().astype(str)
+        X_new = np.column_stack((F.flatten(), T.flatten()))
+        y_pred, y_var = gp_out.gp.predict(gp_out.scaled_y, X_new, return_var=True)
+        y_std = np.sqrt(y_var)
+        y_err = y_std * error
+        y_pred = y_pred * gp_out.y_scaler
+        tts = T.flatten()
+        freqs = F.flatten()
+    else:
+        logger.info("GP is a 1D kernel")
+        if data_mode == 'flux_density':
+            logger.warning("Bandnames/frequency attributes for the transient object may be weird, "
+                           "Please check for yourself")
+            tts = []
+            ys = []
+            yerrs = []
+            bbs = []
+            for key in gp_out.gp.keys():
+                gp = gp_out.gp[key]
+                y_pred, y_cov = gp.predict(gp_out.scaled_y[key], t_new, return_cov=True)
+                y_std = np.sqrt(np.diag(y_cov))
+                y_err = y_std * error
+                y_pred = y_pred * gp_out.y_scaler
+                _bands = np.repeat(key, len(t_new))
+                bbs.append(key)
+                tts.append(t_new)
+                ys.append(y_pred)
+                yerrs.append(y_err)
+            temp_frame = pd.DataFrame({'time': tts, 'ys': ys, 'yerr': yerrs, 'band': bbs})
+            temp_frame.sort_values('time', inplace=True)
+            y_pred = temp_frame['ys']
+            y_err = temp_frame['yerr']
+            bands = temp_frame['band']
+            freqs = temp_frame['band']
+            tts = temp_frame['time']
+        elif data_mode in ['flux', 'magnitude']:
+            tts = []
+            ys = []
+            yerrs = []
+            bbs = []
+            for band in transient.unique_bands:
+                gp = gp_out.gp[band]
+                y_pred, y_cov = gp.predict(gp_out.scaled_y[band], t_new, return_cov=True)
+                y_std = np.sqrt(np.diag(y_cov))
+                y_err = y_std * error
+                y_pred = y_pred * gp_out.y_scaler
+                _bands = np.repeat(band, len(t_new))
+                bbs.append(_bands)
+                tts.append(t_new)
+                ys.append(y_pred)
+                yerrs.append(y_err)
+            temp_frame = pd.DataFrame({'time':tts, 'ys':ys, 'yerr':yerrs, 'band':bbs})
+            temp_frame.sort_values('time', inplace=True)
+            y_pred = temp_frame['ys']
+            y_err = temp_frame['yerr']
+            bands = temp_frame['band']
+            tts = temp_frame['time']
+        elif data_mode == 'luminosity':
+            y_pred, y_cov = gp_out.gp.predict(gp_out.scaled_y, t_new, return_cov=True)
+            y_std = np.sqrt(np.diag(y_cov))
+            y_err = y_std * error
+            y_pred = y_pred * gp_out.y_scaler
+            tts = t_new
+
+    logger.info(f"Data mode: {data_mode}")
+    logger.info("Creating new transient object with GP data")
+    if data_mode == 'flux_density':
+        new_transient = redback.transient.OpticalTransient(name=transient.name + '_gp',
+                                                           flux_density=y_pred, flux_density_err=y_err,
+                                                           time=tts, bands=bands, frequency=freqs,
+                                                           data_mode=data_mode, redshift=transient.redshift)
+    elif data_mode == 'flux':
+        new_transient = redback.transient.OpticalTransient(name=transient.name + '_gp',
+                                                           flux=y_pred, flux_err=y_err,
+                                                           time=tts, bands=bands,
+                                                           data_mode=data_mode, redshift=transient.redshift)
+    elif data_mode == 'magnitude':
+        new_transient = redback.transient.OpticalTransient(name=transient.name + '_gp',
+                                                           magnitude=y_pred, magnitude_err=y_err,
+                                                           time=tts, bands=bands,
+                                                           data_mode=data_mode, redshift=transient.redshift)
+    elif data_mode == 'luminosity':
+        new_transient = redback.transient.OpticalTransient(name=transient.name + '_gp',
+                                                           Lum50=y_pred, Lum50_err=y_err,
+                                                           time_rest_frame=tts, data_mode=data_mode)
+    return new_transient

redback/filters.py

@@ -5,7 +5,7 @@ import numpy as np
 import redback
 import sncosmo
 
-def add_to_database(LABEL, WAVELENGTH, ZEROFLUX, DATABASE, PLOT_LABEL):
+def add_to_database(LABEL, WAVELENGTH, ZEROFLUX, DATABASE, PLOT_LABEL, EFFECTIVE_WIDTH):
     
     """
     Add a filter to the Redback filter database.
@@ -15,8 +15,9 @@ def add_to_database(LABEL, WAVELENGTH, ZEROFLUX, DATABASE, PLOT_LABEL):
     """
 
     frequency = 3.0e8 / WAVELENGTH
-
-    DATABASE.add_row([LABEL, frequency, WAVELENGTH*1e10, 'black', ZEROFLUX, LABEL, PLOT_LABEL])
+    effective_width = 3.0e8 / EFFECTIVE_WIDTH
+    print(effective_width)
+    DATABASE.add_row([LABEL, frequency, WAVELENGTH*1e10, 'black', ZEROFLUX, LABEL, PLOT_LABEL, effective_width])
 
 def add_to_sncosmo(LABEL, TRANSMISSION):
 
@@ -30,7 +31,7 @@ def add_to_sncosmo(LABEL, TRANSMISSION):
     band  = sncosmo.Bandpass(TRANSMISSION['Wavelength'], TRANSMISSION['Transmission'], name=LABEL, wave_unit=u.angstrom)
     sncosmo.register(band, LABEL, force=True)
 
-def add_filter_svo(FILTER, LABEL, PLOT_LABEL=None):
+def add_filter_svo(FILTER, LABEL, PLOT_LABEL=None, OVERWRITE=False):
 
     """
     Wrapper to add a filter from SVO to SNCosmo and the Redback filter database
@@ -44,9 +45,13 @@ def add_filter_svo(FILTER, LABEL, PLOT_LABEL=None):
     
     mask = np.where((database_filters['bands'] == LABEL) & (database_filters['sncosmo_name'] == LABEL))[0]
     
-    # Only add filter to filter database if entry does not exist in the Redback database
+    # Only add filter to filter database if entry does not exist in the Redback database by default
     
-    if len(mask) == 0:
+    # If no entry exists or you choose to overwrite an entry
+    if (len(mask) == 0) or ( (len(mask) != 0) & OVERWRITE ):
+
+        if len(mask) > 0:
+            database_filters.remove_rows(mask)
 
         # Reference (=pivot) wavelength, unit: AA
         wavelength_pivot = FILTER['WavelengthRef']
@@ -71,7 +76,7 @@ def add_filter_svo(FILTER, LABEL, PLOT_LABEL=None):
 
         plot_label = PLOT_LABEL if PLOT_LABEL != None else LABEL
 
-        add_to_database(LABEL, wavelength_pivot * 1.0e-10, zeroflux, database_filters, plot_label)
+        add_to_database(LABEL, wavelength_pivot * 1.0e-10, zeroflux, database_filters, plot_label, effective_width)
 
     # Non-standard filters always needs to be re-added to SN Cosmo even if an entry exists in filter.csv
 
@@ -80,9 +85,10 @@ def add_filter_svo(FILTER, LABEL, PLOT_LABEL=None):
 
     # Prettify output
 
-    database_filters['wavelength [Hz]'].info.format = '.03e'
-    database_filters['wavelength [Angstrom]'].info.format = '.02f'
-    database_filters['reference_flux'].info.format = '.03e'
+    database_filters['wavelength [Hz]'].info.format = '.05e'
+    database_filters['wavelength [Angstrom]'].info.format = '.05f'
+    database_filters['reference_flux'].info.format = '.05e'
+    database_filters['effective_width [Hz]'].info.format = '.05e'
 
     database_filters.write(redback_db_fname, overwrite=True, format='csv')
 
@@ -154,13 +160,14 @@ def add_filter_user(FILE, LABEL, PLOT_LABEL=None, OVERWRITE=False):
 
         print(LABEL, wavelength_pivot * 1.0e-10, zeroflux, plot_label)
 
-        add_to_database(LABEL, wavelength_pivot * 1.0e-10, zeroflux, database_filters, plot_label)
+        add_to_database(LABEL, wavelength_pivot * 1.0e-10, zeroflux, database_filters, plot_label, effective_width)
 
         # Prettify output
 
-        database_filters['wavelength [Hz]'].info.format = '.03e'
-        database_filters['wavelength [Angstrom]'].info.format = '.02f'
-        database_filters['reference_flux'].info.format = '.03e'
+        database_filters['wavelength [Hz]'].info.format = '.05e'
+        database_filters['wavelength [Angstrom]'].info.format = '.05f'
+        database_filters['reference_flux'].info.format = '.05e'
+        database_filters['effective_width [Hz]'].info.format = '.05e'
 
         database_filters.write(redback_db_fname, overwrite=True, format='csv')
     
@@ -168,7 +175,7 @@ def add_filter_user(FILE, LABEL, PLOT_LABEL=None, OVERWRITE=False):
         
         print('Filter {} already exists. Set OVERWRITE to True if you want to overwrite the existing entry'.format(LABEL))
 
-def add_common_filters():
+def add_common_filters(overwrite=False):
 
     """
     Adds Euclid, NTT/EFOSC2, MPG/GROND, Spitzer and WISE filters from SVO
@@ -184,7 +191,7 @@ def add_common_filters():
     filter_label = ['grond::' + x.split('/')[1].split('.')[1] for x in filter_list['filterID']]
     plot_label   = ['GROND/' + x.split('/')[1].split('.')[1] for x in filter_list['filterID']]
 
-    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii]) for ii in range(len(filter_list))]
+    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii], OVERWRITE=overwrite) for ii in range(len(filter_list))]
 
     print('done.\n')
 
@@ -199,7 +206,7 @@ def add_common_filters():
     filter_label = ['efosc2::' + x for x in filter_list['Band']]
     plot_label   = ['EFOSC/' + x for x in filter_list['Band']]
 
-    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii]) for ii in range(len(filter_list))]
+    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii], OVERWRITE=overwrite) for ii in range(len(filter_list))]
 
     print('done.\n')
 
@@ -211,7 +218,7 @@ def add_common_filters():
     filter_label = ['euclid::' + x.split('/')[1].split('.')[1] for x in filter_list['filterID']]
     plot_label   = ['EUCLID/' + x.split('/')[1].split('.')[1].upper() for x in filter_list['filterID']]
 
-    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii]) for ii in range(len(filter_list))]
+    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii], OVERWRITE=overwrite) for ii in range(len(filter_list))]
 
 
     filter_list  = SvoFps.get_filter_list(facility='Euclid', instrument='NISP')
@@ -220,7 +227,7 @@ def add_common_filters():
     filter_label = ['euclid::' + x.split('/')[1].split('.')[1] for x in filter_list['filterID']]
     plot_label   = ['EUCLID/' + x.split('/')[1].split('.')[1] for x in filter_list['filterID']]
 
-    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii]) for ii in range(len(filter_list))]
+    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii], OVERWRITE=overwrite) for ii in range(len(filter_list))]
 
     print('done.\n')
 
@@ -232,7 +239,7 @@ def add_common_filters():
     filter_label = ['irac::' + x.split('/')[1].split('.')[1] for x in filter_list['filterID']]
     plot_label   = ['IRAC/' + x.split('/')[1].split('.')[1] for x in filter_list['filterID']]
 
-    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii]) for ii in range(len(filter_list))]
+    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii], OVERWRITE=overwrite) for ii in range(len(filter_list))]
 
     print('done.\n')
 
@@ -244,13 +251,40 @@ def add_common_filters():
     filter_label = ['wise::' + x.split('/')[1].split('.')[1] for x in filter_list['filterID']]
     plot_label   = ['WISE/' + x.split('/')[1].split('.')[1] for x in filter_list['filterID']]
 
-    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii]) for ii in range(len(filter_list))]
+    [add_filter_svo(filter_list[ii], filter_label[ii], plot_label[ii], OVERWRITE=overwrite) for ii in range(len(filter_list))]
 
     print('done.\n')
 
 def show_all_filters():
 
-    redback_db_fname = path = redback.__path__[0] + '/tables/filters.csv'
+    redback_db_fname = redback.__path__[0] + '/tables/filters.csv'
     database_filters = ascii.read(redback_db_fname)
     
-    return database_filters
+    return database_filters
+
+def add_effective_widths():
+    """
+    Adds effective widths to the Redback filter database
+    :return: None
+    """
+    import pandas as pd
+    db = pd.read_csv(redback.__path__[0] + '/tables/filters.csv')
+    import sncosmo
+    eff_width = np.zeros(len(db))
+    for ii, bb in enumerate(db['sncosmo_name']):
+        try:
+            band = sncosmo.get_bandpass(bb)
+            waves = band.wave  # wavelengths in Angstroms
+            trans = band.trans  # corresponding transmission values
+
+            # Calculate the effective width:
+            #   effective_width = ∫T(λ) dλ / max(T(λ))
+            effective_width = np.trapz(trans, waves) / np.max(trans)
+            effective_width = effective_width * u.Angstrom
+            eff_width[ii] = effective_width.to(u.Hz, equivalencies=u.spectral()).value
+        except Exception:
+            redback.utils.logger.warning("Failed for band={} at index={}".format(bb, ii))
+            eff_width[ii] = db['wavelength [Hz]'].iloc[ii]
+
+    db['effective_width [Hz]'] = eff_width
+    db.to_csv(redback.__path__[0] + '/tables/filters.csv', index=False)

redback/get_data/directory.py

@@ -12,6 +12,24 @@ SWIFT_PROMPT_BIN_SIZES = ['1s', '2ms', '8ms', '16ms', '64ms', '256ms']
 DirectoryStructure = namedtuple("DirectoryStructure", ['directory_path', 'raw_file_path', 'processed_file_path'])
 
 
+def spectrum_directory_structure(transient: str) -> DirectoryStructure:
+    """Provides directory structure for any spectrum data.
+
+    :param transient: Name of the GRB, e.g. GRB123456.
+    :type transient: str
+
+    :return: The directory structure, with 'directory_path', 'raw_file_path', and 'processed_file_path'
+    :rtype: namedtuple
+    """
+    directory_path = f'spectrum/'
+    check_directory_exists_and_if_not_mkdir(directory_path)
+
+    raw_file_path = f"{directory_path}{transient}_rawdata.csv"
+    processed_file_path = f"{directory_path}{transient}.csv"
+
+    return DirectoryStructure(
+        directory_path=directory_path, raw_file_path=raw_file_path, processed_file_path=processed_file_path)
+
 def afterglow_directory_structure(grb: str, data_mode: str, instrument: str = 'BAT+XRT') -> DirectoryStructure:
     """Provides directory structure for Swift afterglow data.