utils.plot_absorber modified, redshift estimate updated

docs/conf.py

@@ -17,7 +17,7 @@
     "github_repo": "qsoabsfind",  # Repo name
     "github_version": "main",  # Version
     "conf_py_path": "/docs/",  # Path in the checkout to the docs root
-    "last_updated": datetime.now(),
+    "last_updated": (datetime.now().date()),
     "commit": False,
 }
 

qsoabsfind/absfinder.py

@@ -5,7 +5,7 @@
 import numpy as np
 from functools import reduce
 from operator import add
-from .utils import convolution_fun, vel_dispersion
+from .utils import convolution_fun, vel_dispersion, elapsed
 from .absorberutils import (
     estimate_local_sigma_conv_array, group_and_weighted_mean_selection_function,
     median_selection_after_combining,
@@ -16,6 +16,7 @@
 from .config import load_constants
 from .spec import QSOSpecRead
 from numba import jit
+import time
 
 constants = load_constants()
 lines, oscillator_parameters, speed_of_light = constants.lines, constants.oscillator_parameters, constants.speed_of_light
@@ -99,14 +100,14 @@ def read_single_spectrum_and_find_absorber(fits_file, spec_index, absorber, **kw
         - This function assumes that the input spectra are already normalized (i.e., flux divided by continuum).
         - The wavelength search region is determined dynamically based on the observed wavelength range.
     """
-
+    start_time = time.time()
     # Read the specified QSO spectrum from the FITS file
     spectra = QSOSpecRead(fits_file, spec_index)
     z_qso = spectra.metadata['Z_QSO']
     lam_obs = spectra.wavelength
 
     # Define the wavelength range for searching the absorber
-    min_wave, max_wave = lam_obs.min() + 500, lam_obs.max() - 500
+    min_wave, max_wave = lam_obs.min() + 200, lam_obs.max() - 200  # avoiding edges
 
     # Retrieve flux and error data, ensuring consistent dtype for Numba compatibility
     residual, error = spectra.flux.astype('float64'), spectra.error.astype('float64')
@@ -123,12 +124,11 @@ def read_single_spectrum_and_find_absorber(fits_file, spec_index, absorber, **kw
     # Verify that the arrays are of equal size
     assert lam_search.size == unmsk_residual.size == unmsk_error.size, "Mismatch in array sizes of lam_search, unmsk_residual, and unmsk_error"
 
-    # Print progress for every 5000th spectrum processed
-    if spec_index % 5000 == 0:
-        print(f'Detection finished up to spec index = {spec_index}', flush=True)
-
     (index_spec, pure_z_abs, pure_gauss_fit, pure_gauss_fit_std, pure_ew_first_line_mean, pure_ew_second_line_mean, pure_ew_total_mean, pure_ew_first_line_error, pure_ew_second_line_error, pure_ew_total_error, redshift_err, sn1_all, sn2_all) = convolution_method_absorber_finder_in_QSO_spectra(spec_index, absorber, lam_obs, residual, error, lam_search, unmsk_residual, unmsk_error, **kwargs)
 
+    # Print progress for every spectrum processed
+    elapsed(start_time, f"\nTime taken to finish absorber detection for index = {spec_index} is: ")
+
     return (index_spec, pure_z_abs, pure_gauss_fit, pure_gauss_fit_std, pure_ew_first_line_mean, pure_ew_second_line_mean, pure_ew_total_mean, pure_ew_first_line_error, pure_ew_second_line_error, pure_ew_total_error, redshift_err, sn1_all, sn2_all)
 
 

qsoabsfind/ew.py

@@ -7,10 +7,33 @@
 from .utils import double_gaussian
 from .absorberutils import redshift_estimate
 from .config import load_constants
+from numba import jit
 
 constants = load_constants()
 lines = constants.lines
 
+def return_line_centers(use_kernel):
+    """
+    Return line centers for a given absorber
+
+    Args:
+        use_kerne (str): absorber (e.g. MgII, CIV)
+
+    Returns:
+        line centers (floats)
+    """
+
+    if use_kernel == 'MgII':
+        line_centre1 = lines['MgII_2796']
+        line_centre2 = lines['MgII_2803']
+    elif use_kernel == 'CIV':
+        line_centre1 = lines['CIV_1548']
+        line_centre2 = lines['CIV_1550']
+    else:
+        raise ValueError("Unsupported kernel type. Use 'MgII', or 'CIV'.")
+
+    return line_centre1, line_centre2
+
 # Example usage within double_curve_fit
 def double_curve_fit(index, fun_to_run, lam_fit_range, nmf_resi_fit, error_fit, bounds, init_cond, iter_n):
     """
@@ -134,6 +157,7 @@ def partial_derivatives(popt):
 
     return EW1_error, EW2_error, EW_total_error
 
+
 def measure_absorber_properties_double_gaussian(index, wavelength, flux, error, absorber_redshift, bound, use_kernel='Mg', d_pix=0.6, use_covariance=False):
     """
     Measures the properties of each potential absorber by fitting a double
@@ -180,17 +204,7 @@ def measure_absorber_properties_double_gaussian(index, wavelength, flux, error,
     EW_total_error = np.zeros(size_array, dtype='float32')
     z_abs_err = np.zeros(size_array, dtype='float32')
 
-    if use_kernel == 'MgII':
-        line_centre1 = lines['MgII_2796']
-        line_centre2 = lines['MgII_2803']
-    elif use_kernel == 'FeII':
-        line_centre1 = lines['FeII_2586']
-        line_centre2 = lines['FeII_2600']
-    elif use_kernel == 'CIV':
-        line_centre1 = lines['CIV_1548']
-        line_centre2 = lines['CIV_1550']
-    else:
-        raise ValueError("Unsupported kernel type. Use 'Mg', 'Fe', or 'CIV'.")
+    line_centre1, line_centre2 = return_line_centers(use_kernel)
 
     #defining wwavelength range for Gaussian fitting
     sigma = d_pix*15 # assuming maximum line width of d_pix * 15, can be larger/smaller, but this is a reasonable assumption.
@@ -229,6 +243,16 @@ def measure_absorber_properties_double_gaussian(index, wavelength, flux, error,
 
             z_abs_array[k], z_abs_err[k] = redshift_estimate(fitted_l1, fitted_l2, std_fitted_l1, std_fitted_l2, line_centre1, line_centre2)
 
+            #best-fit corresponding to this best redshift
+            absorber_rest_lam = wavelength / (1 + z_abs_array[k]) # rest-frame conversion of wavelength
+            lam_ind = np.where((absorber_rest_lam >= ix0) & (absorber_rest_lam <= ix1))[0]
+            lam_fit = absorber_rest_lam[lam_ind]
+            nmf_resi = flux[lam_ind]
+            error_flux = error[lam_ind]
+
+            fitting_param_for_spectrum[k], fitting_param_std_for_spectrum[k], EW_first_line[k], EW_second_line[k], EW_total[k] = double_curve_fit(
+                index, double_gaussian, lam_fit, nmf_resi, error_fit=error_flux, bounds=bound, init_cond=init_cond, iter_n=1000)
+
             #errors on EW
             if not use_covariance:
                 EW_first_line_error[k], EW_second_line_error[k], EW_total_error[k] = calculate_ew_errors(fitting_param_for_spectrum[k], fitting_param_std_for_spectrum[k])

qsoabsfind/io.py

@@ -90,3 +90,4 @@ def save_results_to_fits(results, input_file, output_file, headers, absorber):
     qso_hdu = fits.BinTableHDU(metadata, name='METADATA')
     hdul = fits.HDUList([fits.PrimaryHDU(header=hdr), hdu, qso_hdu])
     hdul.writeto(output_file, overwrite=True)
+    print(f'INFO: ouptut file {output_file} written.')

qsoabsfind/spec.py

@@ -10,20 +10,22 @@ class QSOSpecRead:
     """
     A class to read and handle QSO spectra from a FITS file containing FLUX, ERROR, WAVELENGTH, and METADATA extensions."""
 
-    def __init__(self, fits_file, index=None):
+    def __init__(self, fits_file, index=None, verbose=False):
         """
         Initializes the QSOSpecRead class.
 
         Args:
             fits_file (str): Path to the FITS file containing QSO spectra.
             index (int, list, or np.ndarray, optional): Index or indices of the rows to load. Default is None.
+            verbose (bool): if want to print time info
         """
         self.fits_file = fits_file
         self.flux = None
         self.error = None
         self.wavelength = None
         self.metadata = None
         self.index = index
+        self.verbose = verbose
         self.read_fits()
 
     def read_fits(self):
@@ -33,7 +35,8 @@ def read_fits(self):
         """
         start_time = time.time()
         self.flux, self.error, self.wavelength, self.metadata = read_fits_file(self.fits_file, self.index)
-        elapsed(start_time, "\nTime taken to read FITS file")
+        if self.verbose:
+            elapsed(start_time, "\nTime taken to read FITS file")
 
     def get_metadata(self):
         """

qsoabsfind/utils.py

@@ -9,6 +9,7 @@
 import matplotlib.pyplot as plt
 import os
 from astropy.io import fits
+from astropy.table import Table
 
 # Configure logging
 #logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -277,24 +278,44 @@ def vel_dispersion(c1, c2, sigma1, sigma2, resolution):
 
     return corr_del_v1_sq, corr_del_v2_sq
 
-def plot_absorber(lam, residual, absorber, zabs, xlabel='obs wave (ang)', ylabel='residual', title='QSO', plot_filename=None, zoom=False):
+def plot_absorber(lam, residual, absorber, zabs, xlabel='obs wave (ang)', ylabel='residual', title='QSO', plot_filename=None):
     """
     Saves a plot of spectra with absorber in the current working directory.
 
     Args:
         lam (array-like): Observed wavelength.
         residual (array-like): Residual flux.
         absorber (str): Type of absorber, e.g., 'MgII', 'CIV'.
-        zabs (list): List of absorber redshifts.
+        zabs (list, array, or Table): Absorber redshifts, or a Table with 'Z_ABS' and 'GAUSS_FIT' columns.
         xlabel (str): The label for the x-axis. Default is 'obs wave (ang)'.
         ylabel (str): The label for the y-axis. Default is 'residual'.
-        title (str): The title of the plot. Default is 'QSO'.
+        title (str): The super title of the plot. Default is 'QSO'.
         plot_filename (str): If provided, will save the plot to the given filename.
-        zoom (bool): If True, adds an inset plot to show a zoomed-in version of the absorber lines.
     """
-    # Create the main plot
-    plt.figure(figsize=(12, 4))
-    plt.plot(lam, residual, ls='-', lw=1.5)
+    # If zabs is a Table or structured array, extract redshifts and fit parameters
+    if isinstance(zabs, (Table, np.ndarray)) and ('Z_ABS' in zabs.colnames or 'Z_ABS' in zabs.dtype.names):
+        redshifts = zabs['Z_ABS']
+        fit_params = zabs['GAUSS_FIT']
+    else:
+        redshifts = zabs
+        fit_params = None
+
+    if isinstance(redshifts, float):
+        redshifts = [redshifts]
+
+    num_absorbers = len(redshifts)
+
+    # Create a grid with 2 rows: 1 for the main plot and 1 for zoomed plots
+    fig = plt.figure(figsize=(6 * num_absorbers, 8))
+    fig.subplots_adjust(hspace=0.15, wspace=0.15)  # Adjust space between plots
+
+    # Super title for the entire figure
+    fig.suptitle(title, fontsize=16)
+
+    # Create the main plot in the first row
+    ax_main = plt.subplot2grid((2, num_absorbers), (0, 0), colspan=num_absorbers)
+    ax_main.plot(lam, residual, ls='-', lw=1.5)
+    ax_main.set_xlim(3800, 9200)
 
     # Determine the absorber line labels
     if absorber == 'MgII':
@@ -304,35 +325,54 @@ def plot_absorber(lam, residual, absorber, zabs, xlabel='obs wave (ang)', ylabel
     else:
         raise ValueError(f"Unsupported absorber type: {absorber}")
 
-    # Plot vertical lines for the absorber lines
-    for z in zabs:
+    # Plot vertical lines for the absorber lines in the main plot
+    for z in redshifts:
         x1, x2 = lines[l1] * (1 + z), lines[l2] * (1 + z)
-        plt.axvline(x=x1, color='r', ls='--', label=f'{l1} z={z:.3f}')
-        plt.axvline(x=x2, color='r', ls='--', label=f'{l2} z={z:.3f}')
+        ax_main.axvline(x=x1, color='r', ls='--')
+        ax_main.axvline(x=x2, color='r', ls='--')
 
-    plt.xlabel(xlabel)
-    plt.ylabel(ylabel)
-    plt.title(title)
-    plt.grid(True)
-    plt.ylim(-1, 2)
-
-    # Add inset zoomed plot if requested
-    if zoom:
-        sep = 25
-        ax_inset = plt.gca().inset_axes([0.6, 0.15, 0.35, 0.4])  # position: [x0, y0, width, height]
-        for z in zabs:
-            x1, x2 = lines[l1] * (1 + z), lines[l2] * (1 + z)
-            mask = (lam > x1 - sep) & (lam < x2 + sep)  # zoom range around the lines
-            ax_inset.plot(lam[mask], residual[mask], ls='-', lw=1.5)
-            ax_inset.axvline(x=x1, color='r', ls='--')
-            ax_inset.axvline(x=x2, color='r', ls='--')
-        ax_inset.set_xlim([x1 - sep, x2 + sep])
-        #Determine appropriate y-limits for inset based on data
+    ax_main.set_xlabel(xlabel)
+    ax_main.set_ylabel(ylabel)
+    ax_main.grid(True)
+    ax_main.set_ylim(-1, 2)
+
+    # Add subplots for zoomed-in regions in the second row
+    sep = 25  # Set separation for zoomed plot ranges
+
+    for idx, z in enumerate(redshifts):
+        shift_z = 1 + z
+        ax_zoom = plt.subplot2grid((2, num_absorbers), (1, idx))
+        x1, x2 = lines[l1] * shift_z, lines[l2] * shift_z
+        mask = (lam > x1 - sep) & (lam < x2 + sep)  # Define zoom range around the lines
+
+        ax_zoom.plot(lam[mask], residual[mask], ls='-', lw=1.5, label='data')
+        ax_zoom.axvline(x=x1, color='r', ls='--')
+        ax_zoom.axvline(x=x2, color='r', ls='--')
+        ax_zoom.set_xlim([x1 - sep, x2 + sep])
+
+        # Determine appropriate y-limits for the subplot based on data
         y_min, y_max = residual[mask].min(), residual[mask].max()
         y_margin = 0.2 * (y_max - y_min)  # Add a margin for better visibility
-        ax_inset.set_ylim(y_min - y_margin, y_max + y_margin)
-        ax_inset.set_title(f'Zoomed-In {absorber} Lines')
-        ax_inset.grid(True)
+        ax_zoom.set_ylim(y_min - y_margin, y_max + y_margin)
+
+        ax_zoom.set_title(f'{absorber} at z={z:.3f}')
+        ax_zoom.grid(True)
+        ax_zoom.set_xlabel(xlabel)
+        ax_zoom.set_ylabel(ylabel)
+
+        # Add Gaussian fit
+        if fit_params is not None:
+            params = fit_params[idx]
+            # Adjust fit parameters for the redshift
+            # Plot the Gaussian fit
+            lam_fit = np.linspace(x1 - sep, x2 + sep, 1000)
+            fit_curve = double_gaussian(lam_fit, params[0], shift_z * params[1], shift_z * params[2],
+            params[3], shift_z * params[4], shift_z * params[5])
+            ax_zoom.plot(lam_fit, fit_curve, 'r-', label='Gaussian Fit')
+        ax_zoom.legend()
+
+    # Use tight_layout to ensure there are no overlaps
+    plt.tight_layout(rect=[0, 0, 1, 0.96])  # Reserve space for suptitle
 
     # Save or display the plot
     if plot_filename is not None:
@@ -350,6 +390,7 @@ def plot_absorber(lam, residual, absorber, zabs, xlabel='obs wave (ang)', ylabel
     else:
         plt.show()
 
+
 def read_nqso_from_header(file_path, hdu_name='METADATA'):
     """
     Read the NAXIS2 value from the header of a specified HDU in a FITS file.