First, partial refactoring (d2eb632d) · Commits · Andrea Giannetti / CoCoA_methanol_archive

etl/mdl/mdl_extract_ch3oh_spectra.py

+296 −101

Changes for etl/mdl/mdl_extract_ch3oh_spectra.py: 296 added lines, 101 removed lines.
Original line number	Diff line number	Diff line
		@@ -85,24 +85,24 @@ def summarize_ratio_posterior(R_samples):
		def compute_reliability_score(component):
		score = 0
		thresholds = {
		'SNR_intensity1': 50,
		'SNR_intensity2': 15,
		'SNR_intensity3': 10,
		'SNR1': 25,
		'SNR2': 10,
		'SNR3': 5,
		'SNR_intensity86': 50,
		'SNR_intensity87': 15,
		'SNR_intensity88': 10,
		'SNR86': 25,
		'SNR87': 10,
		'SNR88': 5,
		'ratio1': [0.1, 0.8],
		'ratio2': [0.1, 0.75],
		'rel_prec_ratio1': 0.7,
		'rel_prec_ratio2': 0.5,
		}
		weights = {
		'SNR_intensity1': 1,
		'SNR_intensity2': 3,
		'SNR_intensity3': 5,
		'SNR1': 1,
		'SNR2': 5,
		'SNR3': 10,
		'SNR_intensity86': 1,
		'SNR_intensity87': 3,
		'SNR_intensity88': 5,
		'SNR86': 1,
		'SNR87': 5,
		'SNR88': 10,
		'ratio1': 1,
		'ratio2': 1,
		'rel_prec_ratio1': 3,
		@@ -197,22 +197,238 @@ def compute_density_pdf(density, density_axis=None):
		return _density_axis, pdf / norm


		def load_moment_map(source, species):
		fn = os.path.join("prs", "figures", "moments", f"moment_zero_{source}_{species}.fits")
		hdu = fits.open(fn)[0]
		return hdu.data, WCS(hdu.header)


		def load_continuum_map(filename, datadir):
		hdu = fits.open(os.path.join(datadir, filename))[0]
		data = np.where(hdu.data >= 2.5e-4, hdu.data, np.nan)
		return data, WCS(hdu.header)


		def prepare_line_model_params(source, species, config_mdl, source_props):
		sp = source_props[source][species]
		vlsr = source_props[source]["vlsr"] * u.km / u.s
		freqs = np.array(config_mdl["rest_frequencies"][species]) * u.GHz
		doppler = u.doppler_radio(freqs[0])
		offset = vlsr.to(u.Hz, equivalencies=doppler) - freqs[0]
		doppler_adjusted_frequencies = freqs + offset
		linewidth_guess_hz = convert_dvelocity_to_hz(dvelocity=sp["linewidth_guess_km_s"] * u.km / u.s,
		rest_frequencies=freqs).value / 2.355
		max_width_hz = (convert_dvelocity_to_hz(sp["max_linewidth_km_s"] * u.km / u.s, freqs).value / 2.355
		if "max_linewidth_km_s" in sp else None)
		return doppler_adjusted_frequencies, linewidth_guess_hz, max_width_hz, sp["rms"], sp["frequency_limits_ghz"]


		def should_skip_component(source: str, region: str, vel_component: int) -> bool:
		"""
		Determine if the current velocity component should be skipped based on source and region.

		Parameters:
		source (str): The source name.
		region (str): The region identifier.
		vel_component (int): The index of the velocity component.

		Returns:
		bool: True if the component should be skipped, False otherwise.
		"""
		if (
		(fwhm > 5 * u.km / u.second) or
		((source == 'G332.96') or
		((source == 'G14') and (region == 'envelope_01'))) and (vel_component > 0)):
		return True
		return False


		def extract_fit_parameters(parinfo, idx_shift: int) -> dict:
		"""
		Extract amplitude and width parameters from spectral fit results.

		Parameters:
		parinfo: The fitted spectrum parameters information.
		idx_shift (int): Index offset corresponding to the velocity component.

		Returns:
		dict: Contains WIDTH and AMPLITUDE values for transitions 86, 87, and 88.
		"""
		try:
		fwhm_hz = parinfo[f'WIDTH{idx_shift}'].value
		amp_86 = parinfo[f'AMPLITUDE{idx_shift}'].value
		amp_87 = parinfo[f'AMPLITUDE{idx_shift + 2}'].value
		amp_88 = parinfo[f'AMPLITUDE{idx_shift + 3}'].value
		except (KeyError, AttributeError) as e:
		raise RuntimeError(f"Failed to extract parameters: {e}")

		return {
		'fwhm_hz': fwhm_hz,
		'amplitudes': {
		'86': amp_86,
		'87': amp_87,
		'88': amp_88
		}
		}


		def compute_rms_integral(fwhm_hz: float, channelwidth_hz: float, spec_rms: float) -> float:
		"""
		Compute the RMS of the integrated line intensity assuming a Gaussian profile.

		Parameters:
		fwhm_hz (float): FWHM in Hz.
		channelwidth_hz (float): Spectral channel width in Hz.
		spec_rms (float): RMS noise per channel.

		Returns:
		float: The integrated RMS noise over the line.
		"""
		nchans = fwhm_hz / abs(channelwidth_hz)
		return np.sqrt(2 * nchans) * spec_rms * abs(channelwidth_hz)


		def compute_integrated_intensities(amplitudes: dict, fwhm_hz: float) -> dict:
		"""
		Compute integrated intensities for three transitions assuming Gaussian profiles.

		Parameters:
		amplitudes (dict): Dictionary with amplitudes for '86', '87', and '88'.
		fwhm_hz (float): FWHM in Hz.

		Returns:
		dict: Integrated intensities keyed by transition.
		"""
		return {
		line: 1.0645 * amp * fwhm_hz
		for line, amp in amplitudes.items()
		}


		def compute_snr_params(integrated_intensities: dict, spec_rms: float, rms_integral: float, amplitudes: dict):
		"""
		Compute SNR metrics for each transition and peak amplitude.

		Parameters:
		integrated_intensities (dict): Integrated intensities for transitions '86', '87', '88'.
		spec_rms (float): RMS noise in the spectrum (per channel).
		rms_integral (float): RMS noise of the integrated intensity.
		amplitudes (dict): Peak amplitudes for transitions '86', '87', '88'.

		Returns:
		dict: SNR metrics keyed by parameter name.
		"""
		snr_params = {}
		# SNR on integrated intensity
		for line, int_int in integrated_intensities.items():
		snr_params[f'SNR_intensity{line}'] = int_int / rms_integral

		# SNR on peak amplitude
		for line, amp in amplitudes.items():
		snr_params[f'SNR{line}'] = amp / spec_rms

		return snr_params


		def compute_ratio_posteriors(intensity1: float, intensity2: float, sigma: float,
		lower_limit_ratio: float):
		"""
		Build, sample, and summarize the posterior of the line-intensity ratio.

		Parameters:
		intensity1 (float): Integrated intensity of line 1.
		intensity2 (float): Integrated intensity of line 2.
		sigma (float): RMS noise for the integrated intensities.
		lower_limit_ratio (float): Lower limit for the ratio (for censored modeling).

		Returns:
		tuple: (posterior_samples, summary_dict)
		- posterior_samples: 1D array of sampled ratio values
		- summary_dict: summary statistics from summarize_ratio_posterior
		"""
		# Build the model
		model = build_ratio_model(hatI1=intensity1, hatI2=intensity2,
		sigma=sigma, lower_limit_ratio=lower_limit_ratio)
		# Sample the posterior
		sampled = sample_ratio_model(model=model)
		# Extract and flatten ratio samples
		ratio_samples = sampled.posterior["R"].values.flatten()
		# Summarize
		summary = summarize_ratio_posterior(R_samples=ratio_samples)

		return ratio_samples, summary


		def build_density_pdf_from_ratio(ratio_samples, ratio_name):
		"""
		Given a set of ratio samples, produce a density PDF.

		Parameters:
		ratio_samples (array-like): Samples of the line-intensity ratio.
		ratio_name (str): Identifier for the ratio (e.g., '87-86').

		Returns:
		dict: {
		'density_axis': 1D numpy array of density grid,
		'pdf': 1D numpy array of probabilities matching the axis
		}
		"""
		# Convert ratio samples to densities
		density = convert_ratio_to_density(ratio_samples, ratio_name=ratio_name)

		# Fallback: all NaN or ratio mean too high
		if np.all(np.isnan(density)) or (np.nanmean(ratio_samples) > 0.95):
		density_axis = np.linspace(2.5, 7.5, 500)
		pdf = np.ones_like(density_axis)
		else:
		density_axis, pdf = compute_density_pdf(density)

		return {'density_axis': density_axis, 'pdf': pdf}


		def combine_density_pdfs(density_pdfs, mix_weight):
		"""
		Combine two density PDFs with a weight for cross-ratio mixing.

		Parameters:
		density_pdfs (dict):
		e.g. {
		'87-86': {'density_axis': axis, 'pdf': pdf1},
		'88-86': {'density_axis': axis, 'pdf': pdf2}
		}
		mix_weight (float): Weight for the second PDF in mixing (0 <= mix_weight <= 1).

		Returns:
		tuple: (density_axis, combined_pdf) where combined_pdf is normalized.
		"""
		# Assume both PDFs share the same density_axis
		tags = list(density_pdfs.keys())
		axis = density_pdfs[tags[0]]['density_axis']
		pdf1 = density_pdfs[tags[0]]['pdf']
		pdf2 = density_pdfs[tags[1]]['pdf']

		# Combine: pdf1 * (mix_weight * pdf2 + (1 - mix_weight))
		combined = pdf1 * (mix_weight * pdf2 + (1 - mix_weight))

		# Normalize to unit integral
		area = np.trapezoid(combined, axis)
		if area > 0:
		combined /= area

		return axis, combined


		for source in sources:
		moment_map = fits.open(os.path.join('prs', 'figures', 'moments', f'moment_zero_{source}_{species}.fits'))[0]
		wcs = WCS(moment_map.header)
		vlsr = source_properties[source]['vlsr'] * u.km / u.s
		rms = source_properties[source][species]['rms']
		linewidth_guess = source_properties[source][species]['linewidth_guess_km_s'] * u.km / u.s
		frequency_limits = source_properties[source][species]['frequency_limits_ghz']
		rest_frequencies = np.array(config_mdl['rest_frequencies'][species]) * u.GHz
		freq_to_vel = u.doppler_radio(rest_frequencies[0])
		frequency_offset = vlsr.to(u.Hz, equivalencies=freq_to_vel) - rest_frequencies[0]
		doppler_adjusted_frequencies = rest_frequencies + frequency_offset
		freq_to_vel = u.doppler_radio(rest_frequencies[0])
		linewidth_guess_hz = convert_dvelocity_to_hz(dvelocity=linewidth_guess,
		rest_frequencies=rest_frequencies).value / 2.355
		delta_freq = rest_frequencies[1:] - rest_frequencies[0]
		pars_per_component = 3 * len(rest_frequencies)
		moment_map, wcs = load_moment_map(source, species)
		doppler_adjusted_frequencies, linewidth_guess_hz, max_linewidth_hz, rms, frequency_limits = prepare_line_model_params(
		source=source,
		species=species,
		config_mdl=config_mdl,
		source_props=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])

		tied = []
		line_guesses = []
		@@ -223,17 +439,8 @@ for source in sources:
		frequency_limits=frequency_limits,
		smooth_data=config_mdl['smooth_data'])
		channelwidth_hz = cube.header['CDELT3']
		max_linewidth = source_properties[source][species]['max_linewidth_km_s'] if \
		'max_linewidth_km_s' in source_properties[source][species].keys() else None
		if max_linewidth is not None:
		max_linewidth_hz = convert_dvelocity_to_hz(dvelocity=max_linewidth,
		rest_frequencies=rest_frequencies).value / 2.355
		else:
		max_linewidth_hz = None

		hdu_cont = fits.open(os.path.join(datadir, continuum_fits[source]))[0]
		continuum_data = np.where(hdu_cont.data >= 2.5e-4, hdu_cont.data, np.nan)
		wcs_cont = WCS(hdu_cont.header)
		continuum_data, wcs_cont = load_continuum_map(filename=continuum_fits[source], datadir=datadir)

		densities = {}
		fwhms = {}
		@@ -253,7 +460,7 @@ for source in sources:
		for component in range(components):
		idx_shift = component * pars_per_component

		if len(rest_frequencies) > 1:
		if len(config_mdl["rest_frequencies"][species]) > 1:
		tied_pars = {
		'PAR_VELO': 1 + idx_shift,
		'DELTA_FREQ_1': delta_freq[0].to(u.Hz).value,
		@@ -289,76 +496,64 @@ for source in sources:
		spec_rms = np.nanstd(apspec.specfit.residuals)

		for vel_component in range(components):
		skipped = False
		if ((source == 'G332.96') or ((source == 'G14') and (region == 'envelope_01'))) and (vel_component > 0):
		skipped = True
		else:
		idx_shift = vel_component * 4
		fwhm_hz = apspec.specfit.parinfo[f'WIDTH{idx_shift}'].value
		fwhm = convert_dfreq_to_velocity(dfrequency=fwhm_hz,
		rest_frequencies=rest_frequencies)
		nchans = fwhm_hz / np.abs(channelwidth_hz)
		rms_integral = np.sqrt(2 * nchans) * spec_rms * np.abs(channelwidth_hz)
		fit_params = extract_fit_parameters(parinfo=apspec.specfit.parinfo,
		idx_shift=idx_shift)
		fwhm = convert_dfreq_to_velocity(dfrequency=fit_params['fwhm_hz'],
		rest_frequencies=config_mdl["rest_frequencies"][species] * u.GHz)
		skipped = should_skip_component(source=source,
		region=region,
		vel_component=vel_component)
		if skipped:
		continue
		nchans = fit_params['fwhm_hz'] / np.abs(channelwidth_hz)
		rms_integral = compute_rms_integral(fwhm_hz=fit_params['fwhm_hz'],
		channelwidth_hz=channelwidth_hz,
		spec_rms=spec_rms)
		reliability_score[region][vel_component] = {'pars': {}, 'score': 0}
		if (fwhm < 5 * u.km / u.second):
		int_intensity_86 = 1.0645 * apspec.specfit.parinfo[f'AMPLITUDE{idx_shift}'].value * fwhm_hz
		int_intensity_87 = 1.0645 * apspec.specfit.parinfo[f'AMPLITUDE{idx_shift + 2}'].value * fwhm_hz
		int_intensity_88 = 1.0645 * apspec.specfit.parinfo[f'AMPLITUDE{idx_shift + 3}'].value * fwhm_hz
		reliability_score[region][vel_component]['pars']['SNR_intensity1'] = int_intensity_86 / rms_integral
		reliability_score[region][vel_component]['pars']['SNR_intensity2'] = int_intensity_87 / rms_integral
		reliability_score[region][vel_component]['pars']['SNR_intensity3'] = int_intensity_88 / rms_integral
		reliability_score[region][vel_component]['pars']['SNR1'] = apspec.specfit.parinfo[
		f'AMPLITUDE{idx_shift}'].value / spec_rms
		reliability_score[region][vel_component]['pars']['SNR2'] = apspec.specfit.parinfo[
		f'AMPLITUDE{idx_shift + 2}'].value / spec_rms
		reliability_score[region][vel_component]['pars']['SNR3'] = apspec.specfit.parinfo[
		f'AMPLITUDE{idx_shift + 3}'].value / spec_rms
		model_8786 = build_ratio_model(hatI1=int_intensity_86,
		hatI2=int_intensity_87,
		int_intensities = compute_integrated_intensities(amplitudes=fit_params['amplitudes'],
		fwhm_hz=fit_params['fwhm_hz'])
		snr_params = compute_snr_params(integrated_intensities=int_intensities,
		amplitudes=fit_params['amplitudes'],
		spec_rms=spec_rms,
		rms_integral=rms_integral)

		ratio_posterior_8786, ratio_posterior_summary_8786 = compute_ratio_posteriors(
		intensity1=int_intensities['86'],
		intensity2=int_intensities['87'],
		sigma=rms_integral,
		lower_limit_ratio=0.05)
		sampled_ratio_8786 = sample_ratio_model(model=model_8786)
		ratio_posterior_8786 = sampled_ratio_8786.posterior["R"].values.flatten()
		ratio_posterior_summary_8786 = summarize_ratio_posterior(R_samples=ratio_posterior_8786)

		model_8886 = build_ratio_model(hatI1=int_intensity_86,
		hatI2=int_intensity_88,
		lower_limit_ratio=0.05
		)
		ratio_posterior_8886, ratio_posterior_summary_8886 = compute_ratio_posteriors(
		intensity1=int_intensities['86'],
		intensity2=int_intensities['88'],
		sigma=rms_integral,
		lower_limit_ratio=0.001)
		sampled_ratio_8886 = sample_ratio_model(model=model_8886)
		ratio_posterior_8886 = sampled_ratio_8886.posterior["R"].values.flatten()
		ratio_posterior_summary_8886 = summarize_ratio_posterior(R_samples=ratio_posterior_8886)
		reliability_score[region][vel_component]['pars']['ratio1'] = np.mean(ratio_posterior_8786)
		reliability_score[region][vel_component]['pars']['ratio2'] = np.mean(ratio_posterior_8886)
		reliability_score[region][vel_component]['pars']['rel_prec_ratio1'] = 1 - np.std(
		ratio_posterior_8786) / np.mean(ratio_posterior_8786)
		reliability_score[region][vel_component]['pars']['rel_prec_ratio2'] = 1 - np.std(
		ratio_posterior_8886) / np.mean(ratio_posterior_8886)
		reliability_score[region][vel_component]['score'] = compute_reliability_score(
		component=reliability_score[region][vel_component]['pars'])
		lower_limit_ratio=0.001
		)
		all_pars = {
		**snr_params,
		'ratio1': np.mean(ratio_posterior_8786),
		'ratio2': np.mean(ratio_posterior_8886),
		'rel_prec_ratio1': 1 - np.std(ratio_posterior_8786) / np.mean(ratio_posterior_8786),
		'rel_prec_ratio2': 1 - np.std(ratio_posterior_8886) / np.mean(ratio_posterior_8886)
		}
		score = compute_reliability_score(component=all_pars)
		reliability_score[region][vel_component] = {'pars': all_pars, 'score': score}

		density_pdfs = {}
		for ratio_string, ratio in zip(['87-86', '88-86'], [ratio_posterior_8786, ratio_posterior_8886]):
		density = convert_ratio_to_density(ratio, ratio_name=ratio_string)
		if (np.isnan(density).sum() == len(density)) or (np.nanmean(ratio) > 0.95):
		pdf = np.ones(500)
		density_axis = np.linspace(2.5, 7.5, 500)
		else:
		density_axis, pdf = compute_density_pdf(density)
		density_pdfs[ratio_string] = {'density': density_axis, 'pdf': pdf}
		mix_weight = 0.3
		final_density_pdf = density_pdfs['87-86']['pdf'] * (
		mix_weight * density_pdfs['88-86']['pdf'] + (1 - mix_weight))
		norm = np.trapezoid(final_density_pdf, density_axis)
		final_density_pdf /= norm
		probability_mass = cumulative_trapezoid(final_density_pdf, density_axis)
		if (fwhm.value < 5) and (skipped is False):
		logger.debug(f'Reliability parameters: {reliability_score[region][vel_component]["pars"]}')
		logger.debug(f'Reliability score: {reliability_score[region][vel_component]["score"]}')
		for tag, samples in zip(['87-86', '88-86'], [ratio_posterior_8786, ratio_posterior_8886]):
		density_pdfs[tag] = build_density_pdf_from_ratio(
		ratio_samples=samples,
		ratio_name=tag
		)
		density_axis, final_density_pdf = combine_density_pdfs(density_pdfs, mix_weight=0.3)

		best = float(density_axis[final_density_pdf.argmax()])
		hpd = find_zeros(density_axis, final_density_pdf - prob_threshold)
		densities[region]['best_fit'].append(float(density_axis[np.argmax(final_density_pdf)]))
		densities[region]['best_fit'].append(best)
		densities[region]['uncertainty'].append(hpd)
		densities[region]['reliability_score'].append(float(reliability_score[region][vel_component]['score']))
		fwhms[region].append(fwhm.value)
		densities[region]['reliability_score'].append(score)
		fwhms[region].append(fwhm)

		# plt.clf()
		# plt.plot(density_axis, final_density_pdf)