Commit 47018512 authored by Andrea Giannetti's avatar Andrea Giannetti
Browse files

Refactored process_source

parent b10e83f1

etl/mdl/mdl_extract_ch3oh_spectra.py

+135 −43
Changes for etl/mdl/mdl_extract_ch3oh_spectra.py: 135 added lines, 43 removed lines.
Original line number Diff line number Diff line
@@ -287,7 +287,7 @@ def compute_snr_params(integrated_intensities: dict, spec_rms: float, rms_integr 




def compute_ratio_posteriors(intensity1: float, intensity2: float, sigma: float,

                             lower_limit_ratio: float):

                             lower_limit_ratio: float, samples: int = 1000):

    """

    Build, sample, and summarize the posterior of the line-intensity ratio.
@@ -306,7 +306,7 @@ def compute_ratio_posteriors(intensity1: float, intensity2: float, sigma: float, 
    model = build_ratio_model(hat_int1=intensity1, hat_int2=intensity2,

                              sigma=sigma, lower_limit_ratio=lower_limit_ratio)

    # Sample the posterior

    sampled = sample_ratio_model(model=model)

    sampled = sample_ratio_model(model=model, draws=samples)

    # Extract and flatten ratio samples

    ratio_samples = sampled.posterior["R"].values.flatten()

    # Summarize
@@ -484,9 +484,8 @@ def overlay_region_on_map(apspec, 
    fig_map.savefig(plotfile, bbox_inches='tight')

    plt.close(fig_map)



def process_source(source, species, guesses):

    # 1. Load maps and prepare line models

    moment_map, wcs = load_moment_map(source, species)

def setup_spectral_analysis(source, species):

    """Setup frequency conversions and load spectral data."""

    freq_to_vel = u.doppler_radio(config_mdl["rest_frequencies"][species][0] * u.GHz)

    freqs_dop, _, max_lw_hz, rms, freq_limits = prepare_line_model_params(

        source, species, config_mdl, source_properties
@@ -495,22 +494,38 @@ def process_source(source, species, guesses): 
                  - config_mdl["rest_frequencies"][species][0]) * u.GHz

    pars_per_component = 3 * len(config_mdl["rest_frequencies"][species])



    # 2. Open cube & continuum

    filename = source_properties[source][species]['data']

    cube = get_cube(filename, datadir, freq_limits, config_mdl['smooth_data'])

    channelwidth_hz = cube.header['CDELT3']

    continuum_data, wcs_cont = load_continuum_map(continuum_fits[source], datadir)

    return freq_to_vel, freqs_dop, max_lw_hz, rms, delta_freq, pars_per_component, freq_limits



    # 3. Init outputs

    densities, fwhms, reliability_score = {}, {}, {}



    # 4. Loop over regions

    for region, props in regions[source].items():

        coords, components = props[:3], props[3]

        densities[region] = {'best_fit': [], 'uncertainty': [], 'reliability_score': [], 'fwhm': [], 'vlsr': []}

        fwhms[region], reliability_score[region] = [], {}

def run_spectral_fit(source, coords, components, region, guesses, cube,

                     delta_freq, rms, pars_per_component, freqs_dop, channelwidth_hz, max_lw_hz, species):

    """

    Perform spectral fitting on the provided data cube.



    This function sets up and executes a spectral fit for the specified source and region

    using the given parameters and data cube. It prepares the necessary fitting parameters

    and constraints, runs the fit, and plots the results.



    Parameters:

        source (str): The name of the source being analyzed.

        coords (tuple): The coordinates for the region of interest within the cube.

        components (int): The number of spectral components to fit.

        region (str): The region identifier within the source.

        guesses (dict): Initial guesses for the spectral fitting parameters.

        cube: The data cube containing spectral data.

        delta_freq (Quantity): Frequency offsets for the spectral lines.

        rms (float): The root mean square noise level in the spectrum.

        pars_per_component (int): Number of parameters per spectral component.

        freqs_dop (Quantity): Doppler shifted frequencies for the spectral lines.

        channelwidth_hz (float): Width of the spectral channels in Hz.

        max_lw_hz (float): Maximum allowable line width in Hz.

        species (str): The chemical species being analyzed.



        # 4a. Build & run the spec‐fit

    Returns:

        tuple: A tuple containing:

            - apspec: The fitted spectrum object.

            - spec_rms: The RMS of the residuals after fitting.

    """

    apspec = cube.get_apspec(coords)

    tied, line_guesses, limits, limited = prepare_fit_guesses_and_ties(

        source, region, components, guesses, delta_freq, rms, pars_per_component, freqs_dop,
@@ -528,55 +543,133 @@ def process_source(source, species, guesses): 
        plotdir=plotdir

    )



        # 4b. Loop over each velocity component

        for vel_comp in range(components):

    return apspec, spec_rms





def process_velocity_component(apspec, vel_comp, freq_to_vel, freqs_dop, channelwidth_hz,

                               spec_rms, source, region):

    """Process a single velocity component and return results."""

    idx = vel_comp * 4



            # extract fit params & convert to physical units

    # Extract fit params & convert to physical units

    fitpars = extract_fit_parameters(apspec.specfit.parinfo, idx)

    fwhm = convert_dfreq_to_velocity(fitpars['fwhm_hz'], freqs_dop)

    velocity = (fitpars['frequency'] * u.Hz).to(u.km / u.s, equivalencies=freq_to_vel).value



    if should_skip_component(source, region, vel_comp, fwhm):

                continue

        return None



            # compute SNR, ratios, reliability

    # Compute SNR, ratios, reliability

    rms_int = compute_rms_integral(fitpars['fwhm_hz'], channelwidth_hz, spec_rms)

    ints = compute_integrated_intensities(fitpars['amplitudes'], fitpars['fwhm_hz'])

            snr_p  = compute_snr_params(integrated_intensities=ints, amplitudes=fitpars['amplitudes'], spec_rms=spec_rms, rms_integral=rms_int)

    snr_p = compute_snr_params(integrated_intensities=ints, amplitudes=fitpars['amplitudes'],

                               spec_rms=spec_rms, rms_integral=rms_int)



    return {

        'fitpars': fitpars,

        'fwhm': fwhm,

        'velocity': velocity,

        'rms_int': rms_int,

        'ints': ints,

        'snr_p': snr_p

    }





def compute_density_from_ratios(ints, rms_int, prob_threshold, samples):

    """Compute density PDFs from line ratios."""

    r87, _ = compute_ratio_posteriors(ints['86'], ints['87'], rms_int, lower_limit_ratio=0.05, samples=samples)

    r88, _ = compute_ratio_posteriors(ints['86'], ints['88'], rms_int, lower_limit_ratio=0.001, samples=samples)



            r87, _ = compute_ratio_posteriors(ints['86'], ints['87'], rms_int, 0.05)

            r88, _ = compute_ratio_posteriors(ints['86'], ints['88'], rms_int, 0.001)

    all_pars = {

                **snr_p,

        'ratio1': np.mean(r87),

        'ratio2': np.mean(r88),

        'rel_prec_ratio1': 1 - np.std(r87) / np.mean(r87),

        'rel_prec_ratio2': 1 - np.std(r88) / np.mean(r88)

    }

            score = compute_reliability_score(all_pars)

            reliability_score[region][vel_comp] = {'pars': all_pars, 'score': score}



            # build & combine density PDFs

    # Build & combine density PDFs

    dpdfs = {

        '87-86': build_density_pdf_from_ratio(r87, '87-86'),

        '88-86': build_density_pdf_from_ratio(r88, '88-86')

    }

    axis, final_pdf = combine_density_pdfs(dpdfs, mix_weight=0.3)

            # plt.clf()

            # plt.plot(axis, final_pdf)

            # plt.plot(axis, dpdfs['87-86']['pdf'])

            # plt.plot(axis, dpdfs['88-86']['pdf'])

            # plt.show()



            # extract best & HPD

    # Extract best & HPD

    best = float(axis[final_pdf.argmax()])

    hpd = find_zeros(axis, final_pdf - prob_threshold)



    return all_pars, best, hpd





def save_results(densities, source):

    """Save density results to YAML file."""

    with open(os.path.join('mdl', 'inferred_densities', f'clump_densities_{source}.yml'), 'w') as outfile:

        yaml.dump(densities, outfile, default_flow_style=False)





def process_source(source, species, guesses, prob_threshold, samples):

    # 1. Load maps and prepare line models

    moment_map, wcs = load_moment_map(source, species)

    freq_to_vel, freqs_dop, max_lw_hz, rms, delta_freq, pars_per_component, freq_limits = setup_spectral_analysis(

        source, species)



    # 2. Open cube & continuum

    filename = source_properties[source][species]['data']

    cube = get_cube(filename, datadir, freq_limits, config_mdl['smooth_data'])

    channelwidth_hz = cube.header['CDELT3']

    continuum_data, wcs_cont = load_continuum_map(continuum_fits[source], datadir)



    # 3. Init outputs

    densities, fwhms, reliability_score = {}, {}, {}



    # 4. Loop over regions

    for region, props in regions[source].items():

        coords, components = props[:3], props[3]

        densities[region] = {'best_fit': [], 'uncertainty': [], 'reliability_score': [], 'fwhm': [], 'vlsr': []}

        fwhms[region], reliability_score[region] = [], {}

        # 4a. Run spectral fitting

        apspec, spec_rms = run_spectral_fit(

            source=source,

            region=region,

            components=components,

            guesses=guesses,

            cube=cube,

            delta_freq=delta_freq,

            rms=rms,

            pars_per_component=pars_per_component,

            freqs_dop=freqs_dop,

            channelwidth_hz=channelwidth_hz,

            coords=coords,

            max_lw_hz=max_lw_hz,

            species=species

        )

        # 4b. Loop over each velocity component

        for vel_comp in range(components):

            comp_result = process_velocity_component(

                apspec, vel_comp, freq_to_vel, freqs_dop, channelwidth_hz,

                spec_rms, source, region

            )



            if comp_result is None:

                continue



            # Compute density and reliability

            all_pars, best, hpd = compute_density_from_ratios(ints=comp_result['ints'],

                                                              rms_int=comp_result['rms_int'],

                                                              prob_threshold=prob_threshold,

                                                              samples=samples)

            all_pars.update(comp_result['snr_p'])



            score = compute_reliability_score(all_pars)

            reliability_score[region][vel_comp] = {'pars': all_pars, 'score': score}



            # Store results

            densities[region]['best_fit'].append(best)

            densities[region]['uncertainty'].append(hpd)

            densities[region]['reliability_score'].append(float(score))

            densities[region]['fwhm'].append(float(fwhm.value))

            densities[region]['vlsr'].append(float(velocity))

            fwhms[region].append(fwhm.value)

            densities[region]['fwhm'].append(float(comp_result['fwhm'].value))

            densities[region]['vlsr'].append(float(comp_result['velocity']))

            fwhms[region].append(comp_result['fwhm'].value)



        # 4c. Overlay region on map

        overlay_region_on_map(
@@ -592,8 +685,7 @@ def process_source(source, species, guesses): 
                            densities=densities,

                            moment_map=moment_map,

                            wcs=wcs)

    with open(os.path.join('mdl', 'inferred_densities', f'densities_{source}.yml'), 'w') as outfile:

        yaml.dump(densities, outfile, default_flow_style=False)

    save_results(densities, source)





def plot_density_vs_fwhm(densities, fwhms, source, species, regions):
@@ -648,7 +740,7 @@ def overlay_regions_density(source, regions, densities, moment_map, wcs): 
if __name__ == '__main__':

    sources = ['G14', 'G28', 'G332.96']

    species = 'CH3OH'

    samples = 1000

    samples = 2000

    prob_threshold = 0.67

    matplotlib.use('TkAgg')
@@ -660,4 +752,4 @@ if __name__ == '__main__': 
        guesses = yaml.load(infile, Loader=yaml.FullLoader)



    for source in sources:

        process_source(source, species, guesses)
 
        process_source(source, species, guesses, prob_threshold=prob_threshold, samples=samples)

etl/mdl/mdl_extract_clump_spectra_and_densities.py

+63 −77
Changes for etl/mdl/mdl_extract_clump_spectra_and_densities.py: 63 added lines, 77 removed lines.
Original line number Diff line number Diff line 
import yaml

import os

import matplotlib

import matplotlib.pyplot as plt
@@ -6,21 +5,14 @@ import numpy as np 
import pyspeckit

from astropy import units as u

from assets.commons import (setup_logger,

                            get_cube,

                            convert_dfreq_to_velocity,

                            find_zeros)

                            get_cube)

from mdl.mdl_extract_ch3oh_spectra import (compute_reliability_score,

                                           prepare_line_model_params,

                                           should_skip_component,

                                           extract_fit_parameters,

                                           compute_rms_integral,

                                           compute_integrated_intensities,

                                           compute_snr_params,

                                           compute_ratio_posteriors,

                                           build_density_pdf_from_ratio,

                                           combine_density_pdfs,

                                           prepare_fit_guesses_and_ties,

                                           run_and_plot_specfit,)

                                           run_and_plot_specfit,

                                           setup_spectral_analysis,

                                           process_velocity_component,

                                           compute_density_from_ratios,

                                           save_results)

from assets.constants import source_properties, config_mdl, datadir

from typing import Union
@@ -29,20 +21,13 @@ plt.rcParams['font.size'] = 11 
logger = setup_logger(name='AP. SPECTRA', log_level='DEBUG')





def process_source(source, species, guesses):

    freq_to_vel = u.doppler_radio(config_mdl["rest_frequencies"][species][0] * u.GHz)

    freqs_dop, _, max_lw_hz, rms, freq_limits = prepare_line_model_params(

        source, species, config_mdl, source_properties

    )

    delta_freq = (np.array(config_mdl["rest_frequencies"][species][1:])

                  - config_mdl["rest_frequencies"][species][0]) * u.GHz

    pars_per_component = 3 * len(config_mdl["rest_frequencies"][species])



def load_and_process_cube(source, species, freq_limits):

    """Load cube data and compute mean spectrum with antenna gain correction."""

    filename = source_properties[source][species]['data'].replace('.fits', '_pbcor.fits')

    cube = get_cube(filename, datadir, freq_limits, config_mdl['smooth_data'])



    fwhm_to_sigma = 1. / (8 * np.log(2)) ** 0.5

    antenna_gain = (cube.unit.to

        (

    antenna_gain = (cube.unit.to(

        'K', equivalencies=u.brightness_temperature(

            config_mdl["rest_frequencies"][species][0] * u.GHz,

            beam_area=2 * np.pi * fwhm_to_sigma ** 2 * cube.header['BMAJ'] * cube.header['BMIN'] * u.deg ** 2)
@@ -50,22 +35,20 @@ def process_source(source, species, guesses): 
    mean_spectrum = np.nanmean(cube.cube * antenna_gain, axis=(1, 2))

    channelwidth_hz = cube.header['CDELT3']



    densities, fwhms, reliability_score = {}, {}, {}

    region = 'clump'

    coords, components = regions[source][region][:3], regions[source][region][3]

    return cube, mean_spectrum, channelwidth_hz



    # Initialize storage for this single region

    densities = {region: {'best_fit': [], 'uncertainty': [], 'reliability_score': [], 'fwhm': [], 'vlsr': []}}

    fwhms = {region: []}

    reliability_score = {region: {}}



def run_spectral_fit(source, region, components, guesses, mean_spectrum, cube,

                     delta_freq, rms, pars_per_component, freqs_dop, channelwidth_hz,

                     coords, max_lw_hz, species):

    """Setup and run the spectral fitting."""

    apspec = pyspeckit.Spectrum(data=mean_spectrum, xarr=cube.xarr,

                                xarrkwargs={'unit': 'Hz'},

                                unit='K')



    tied, line_guesses, limits, limited = prepare_fit_guesses_and_ties(

        source, region, components, guesses, delta_freq, rms, pars_per_component, freqs_dop,

        channelwidth_hz, coords, max_lw_hz

        source, region, components, guesses, delta_freq, rms, pars_per_component,

        freqs_dop, channelwidth_hz, coords, max_lw_hz

    )



    spec_rms = run_and_plot_specfit(
@@ -81,64 +64,67 @@ def process_source(source, species, guesses): 
        name_prefix="clump"

    )



    return apspec, spec_rms





def process_source(source, species, guesses, prob_threshold=0.67, samples=1000):

    """Main function to process a single source."""

    # Setup

    freq_to_vel, freqs_dop, max_lw_hz, rms, delta_freq, pars_per_component, freq_limits = setup_spectral_analysis(

        source, species)

    cube, mean_spectrum, channelwidth_hz = load_and_process_cube(source, species, freq_limits)



    # Get region parameters

    region = 'clump'

    coords, components = regions[source][region][:3], regions[source][region][3]



    # Initialize storage

    densities = {region: {'best_fit': [], 'uncertainty': [], 'reliability_score': [], 'fwhm': [], 'vlsr': []}}

    fwhms = {region: []}

    reliability_score = {region: {}}



    # Run spectral fitting

    apspec, spec_rms = run_spectral_fit(

        source, region, components, guesses, mean_spectrum, cube,

        delta_freq, rms, pars_per_component, freqs_dop, channelwidth_hz,

        coords, max_lw_hz, species

    )



    # Process each velocity component

    for vel_comp in range(components):

        idx = vel_comp * 4

        comp_result = process_velocity_component(

            apspec, vel_comp, freq_to_vel, freqs_dop, channelwidth_hz,

            spec_rms, source, region

        )



        # extract fit params & convert to physical units

        fitpars = extract_fit_parameters(apspec.specfit.parinfo, idx)

        fwhm = convert_dfreq_to_velocity(fitpars['fwhm_hz'], freqs_dop)

        velocity = (fitpars['frequency'] * u.Hz).to(u.km / u.s, equivalencies=freq_to_vel).value

        if should_skip_component(source, region, vel_comp, fwhm):

        if comp_result is None:

            continue



        # compute SNR, ratios, reliability

        rms_int = compute_rms_integral(fitpars['fwhm_hz'], channelwidth_hz, spec_rms)

        ints = compute_integrated_intensities(fitpars['amplitudes'], fitpars['fwhm_hz'])

        snr_p = compute_snr_params(integrated_intensities=ints, amplitudes=fitpars['amplitudes'], spec_rms=spec_rms,

                                   rms_integral=rms_int)



        r87, _ = compute_ratio_posteriors(ints['86'], ints['87'], rms_int, 0.05)

        r88, _ = compute_ratio_posteriors(ints['86'], ints['88'], rms_int, 0.001)

        all_pars = {

            **snr_p,

            'ratio1': np.mean(r87),

            'ratio2': np.mean(r88),  # Note: this looks like it should be r88?

            'rel_prec_ratio1': 1 - np.std(r87) / np.mean(r87),

            'rel_prec_ratio2': 1 - np.std(r88) / np.mean(r88)

        }

        # Compute density and reliability

        all_pars, best, hpd = compute_density_from_ratios(ints=comp_result['ints'],

                                                          rms_int=comp_result['rms_int'],

                                                          prob_threshold=prob_threshold,

                                                          samples=samples)

        all_pars.update(comp_result['snr_p'])



        score = compute_reliability_score(all_pars)

        reliability_score[region][vel_comp] = {'pars': all_pars, 'score': score}



        # build & combine density PDFs

        dpdfs = {

            '87-86': build_density_pdf_from_ratio(r87, '87-86'),

            '88-86': build_density_pdf_from_ratio(r88, '88-86')

        }

        axis, final_pdf = combine_density_pdfs(dpdfs, mix_weight=0.3)

        # plt.clf()

        # plt.plot(axis, final_pdf)

        # plt.plot(axis, dpdfs['87-86']['pdf'])

        # plt.plot(axis, dpdfs['88-86']['pdf'])

        # plt.show()



        # extract best & HPD

        best = float(axis[final_pdf.argmax()])

        hpd = find_zeros(axis, final_pdf - prob_threshold)

        # Store results

        densities[region]['best_fit'].append(best)

        densities[region]['uncertainty'].append(hpd)

        densities[region]['reliability_score'].append(float(score))

        densities[region]['fwhm'].append(float(fwhm.value))

        densities[region]['vlsr'].append(float(velocity))

        fwhms[region].append(fwhm.value)

        densities[region]['fwhm'].append(float(comp_result['fwhm'].value))

        densities[region]['vlsr'].append(float(comp_result['velocity']))

        fwhms[region].append(comp_result['fwhm'].value)



    save_results(densities, source)



    with open(os.path.join('mdl', 'inferred_densities', f'clump_densities_{source}.yml'), 'w') as outfile:

        yaml.dump(densities, outfile, default_flow_style=False)



# sources = ['G14', 'G28', 'G332.96']

sources = ['G14', 'G28']

species = 'CH3OH'

samples = 1000

samples = 2000

prob_threshold = 0.67

matplotlib.use('TkAgg')

regions = {
@@ -167,4 +153,4 @@ guesses = { 
plotdir = os.path.join('prs', 'figures', 'aperture_spectra_fits')



for source in sources:

    process_source(source, species, guesses)
 
    process_source(source, species, guesses, prob_threshold=prob_threshold, samples=samples)