Second round of refactoring (0e564748) · Commits · Andrea Giannetti / CoCoA_methanol_archive

etl/mdl/mdl_extract_ch3oh_spectra.py

+162 −84

Changes for etl/mdl/mdl_extract_ch3oh_spectra.py: 162 added lines, 84 removed lines.
Original line number	Diff line number	Diff line
		@@ -25,31 +25,23 @@ from typing import Union
		logger = setup_logger(name='AP. SPECTRA', log_level='DEBUG')


		def build_ratio_model(hatI1, hatI2, sigma, lower_limit_ratio):
		"""
		Build a PyMC model for the ratio R = I2 / I1, where:
		- hatI1: observed intensity of line 1
		- sigma1: noise std of line 1
		- L: detection limit (upper limit) for line 2
		- sigma2: noise std of line 2
		Priors for I1 and I2 are LogNormal(mu, sigma_prior).
		"""
		def build_ratio_model(hat_int1, hat_int2, sigma, lower_limit_ratio):
		model = pm.Model()
		with model:
		# latent ratio R with a lower bound to forbid R→0
		R = pm.TruncatedNormal("R", mu=hatI2 / hatI1, sigma=0.3, lower=lower_limit_ratio, upper=0.95)
		R = pm.TruncatedNormal("R", mu=hat_int2 / hat_int1, sigma=0.3, lower=lower_limit_ratio, upper=0.95)

		# data container for I1
		I1 = pm.Data("I1", hatI1)
		I1 = pm.Data("I1", hat_int1)

		# predicted I2 = R * I1
		mu2 = R * I1

		# detection mask
		is_det = hatI2 >= 3 * sigma
		is_det = hat_int2 >= 3 * sigma

		# likelihood if detected
		logp_det = pm.Normal.logp(hatI2, mu=mu2, sigma=sigma)
		logp_det = pm.Normal.logp(hat_int2, mu=mu2, sigma=sigma)

		# censored likelihood if not detected
		z = (3 * sigma - mu2) / sigma
		@@ -72,13 +64,13 @@ def sample_ratio_model(model, draws=2000, tune=1000, target_accept=0.9, random_s
		return trace


		def summarize_ratio_posterior(R_samples):
		def summarize_ratio_posterior(ratio_samples):
		"""
		Summarize the ratio posterior with mean, median, and 95% HPD.
		"""
		mean = np.mean(R_samples)
		median = np.median(R_samples)
		hpd = az.hdi(R_samples, hdi_prob=0.95)
		mean = np.mean(ratio_samples)
		median = np.median(ratio_samples)
		hpd = az.hdi(ratio_samples, hdi_prob=0.95)
		return {"mean": mean, "median": median, "95%_HPD": hpd}


		@@ -223,7 +215,7 @@ def prepare_line_model_params(source, species, config_mdl, source_props):
		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:
		def should_skip_component(source: str, region: str, vel_component: int, fwhm: float) -> bool:
		"""
		Determine if the current velocity component should be skipped based on source and region.

		@@ -347,14 +339,14 @@ def compute_ratio_posteriors(intensity1: float, intensity2: float, sigma: float,
		- summary_dict: summary statistics from summarize_ratio_posterior
		"""
		# Build the model
		model = build_ratio_model(hatI1=intensity1, hatI2=intensity2,
		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)
		# Extract and flatten ratio samples
		ratio_samples = sampled.posterior["R"].values.flatten()
		# Summarize
		summary = summarize_ratio_posterior(R_samples=ratio_samples)
		summary = summarize_ratio_posterior(ratio_samples=ratio_samples)

		return ratio_samples, summary

		@@ -407,7 +399,6 @@ def combine_density_pdfs(density_pdfs, mix_weight):
		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
		@@ -418,6 +409,111 @@ def combine_density_pdfs(density_pdfs, mix_weight):
		return axis, combined


		def prepare_fit_guesses_and_ties(source, region, component_count,
		guesses, delta_freq, rms, pars_per_component):
		"""
		Build up the `tied` list and `line_guesses` for all components.
		"""
		tied, line_guesses = [], []

		for comp in range(component_count):
		idx = comp * pars_per_component

		tied_pars = {
		'PAR_VELO': 1 + idx,
		'PAR_SIGMA_LINE': 2 + idx,
		}
		for i, df in enumerate(delta_freq, start=1):
		tied_pars[f"DELTA_FREQ_{i}"] = df.to(u.Hz).value

		tied += fill_tied_pars(tied_pars=tied_pars)
		line_guesses += fill_guesses(
		line_peak=guesses[source][region][0][comp],
		doppler_adjusted_frequencies=doppler_adjusted_frequencies,
		shift_hz=guesses[source][region][1][comp],
		linewidth_guess_hz=guesses[source][region][2][comp],
		rms=rms,
		)

		# after building all guesses, compute the fitting limits once
		limits, limited = create_limits(
		par_number=len(line_guesses),
		channelwidth_hz=abs(channelwidth_hz),
		rms=rms / np.sqrt(np.pi * _region_coordinates[2] ** 2),
		max_linewidth=max_linewidth_hz
		)
		return tied, line_guesses, limits, limited


		def run_and_plot_specfit(apspec, tied, line_guesses, limits, limited,
		source, region, species, plotdir):
		"""
		Execute the Gaussian specfit, plot the result, and return the residual RMS.
		"""
		fig = plt.figure()
		apspec.specfit(
		fittype="gaussian",
		guesses=line_guesses,
		tied=tied,
		limits=limits,
		limited=limited,
		clear=True
		)

		# plot fit components
		apspec.plotter(title=f"{source} {region}")
		apspec.specfit.plot_fit(show_components=True, annotate=False)
		fig.savefig(os.path.join(
		plotdir, f"aperture_spectra_{source}_{species}_{region}.svg"
		))
		plt.close(fig)

		# return the RMS of the residuals
		return float(np.nanstd(apspec.specfit.residuals)), fig


		def overlay_region_on_map(apspec,
		moment_map,
		wcs,
		region_coordinates,
		continuum_data,
		wcs_cont,
		plotfile):
		"""
		Add the aperture circle and continuum contours to the left panel.
		"""
		fig_map = plt.figure(figsize=(15, 5))
		_ = fig_map.add_subplot(122)
		apspec.plotter(title=f"{source} {region}", figure=fig_map)
		apspec.specfit.plot_fit(show_components=True, annotate=False)

		ax_map = fig_map.add_subplot(121, projection=wcs)
		ax_map.imshow(np.log10(moment_map.data), origin="lower")
		ax_map.set_aspect("equal", "box")
		circ = Circle(
		(region_coordinates[0], region_coordinates[1]),
		region_coordinates[2],
		edgecolor="k", facecolor="none"
		)
		ax_map.add_patch(circ)

		levels = np.logspace(
		np.log10(2.5e-4),
		np.log10(0.9 * np.nanmax(continuum_data)),
		num=5
		)
		ax_map.contour(
		continuum_data, levels=levels,
		transform=ax_map.get_transform(wcs_cont),
		colors="red", origin="lower"
		)
		ax_map.set_xlabel("RA [J2000]")
		ax_map.set_ylabel("Dec [J2000]")
		fig_map.tight_layout()
		fig_map.savefig(plotfile)
		plt.close(fig_map)


		for source in sources:
		moment_map, wcs = load_moment_map(source, species)
		doppler_adjusted_frequencies, linewidth_guess_hz, max_linewidth_hz, rms, frequency_limits = prepare_line_model_params(
		@@ -446,7 +542,7 @@ for source in sources:
		fwhms = {}
		reliability_score = {}
		for region in regions[source]:
		print(region)
		logger.info(f'Processing region: {region}')
		_region_coordinates = regions[source][region][0:3]
		components = regions[source][region][3]
		region_guesses = regions[source][region]
		@@ -454,63 +550,55 @@ for source in sources:
		line_guesses = []
		tied = []

		densities[region] = {'best_fit': [], 'uncertainty': [], 'reliability_score': []}
		densities[region] = {'best_fit': [], 'uncertainty': [], 'reliability_score': [], 'fwhm': []}
		fwhms[region] = []
		reliability_score[region] = {}
		for component in range(components):
		idx_shift = component * pars_per_component

		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,
		'DELTA_FREQ_2': delta_freq[1].to(u.Hz).value,
		'DELTA_FREQ_3': delta_freq[2].to(u.Hz).value,
		'PAR_SIGMA_LINE': 2 + idx_shift
		}
		tied += fill_tied_pars(tied_pars=tied_pars)

		line_guesses += fill_guesses(line_peak=guesses[source][region][0][component],
		doppler_adjusted_frequencies=doppler_adjusted_frequencies,
		shift_hz=guesses[source][region][1][component],
		linewidth_guess_hz=guesses[source][region][2][component],
		rms=rms)

		limits, limited = create_limits(par_number=len(line_guesses),
		channelwidth_hz=np.abs(channelwidth_hz),
		rms=rms / np.sqrt(np.pi * _region_coordinates[2] ** 2),
		max_linewidth=max_linewidth_hz)
		tied, line_guesses, limits, limited = prepare_fit_guesses_and_ties(
		source=source,
		region=region,
		component_count=components,
		guesses=guesses,
		delta_freq=delta_freq,
		rms=rms,
		pars_per_component=pars_per_component
		)

		apspec.specfit(fittype='gaussian',
		guesses=line_guesses,
		spec_rms, specfig = run_and_plot_specfit(
		apspec=apspec,
		tied=tied,
		line_guesses=line_guesses,
		limits=limits,
		limited=limited,
		clear=True)
		fig = plt.figure(figsize=(15, 5))
		ax_s = fig.add_subplot(122)
		apspec.plotter(title=f'{source} {region}', figure=fig)
		apspec.specfit.plot_fit(show_components=True, annotate=False)
		apspec.plotter.savefig(os.path.join(plotdir, f'aperture_spectra_{source}_{species}_{region}.svg'))

		spec_rms = np.nanstd(apspec.specfit.residuals)
		source=source,
		region=region,
		species=species,
		plotdir=plotdir
		)

		for vel_component in range(components):
		idx_shift = vel_component * 4
		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,
		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)
		vel_component=vel_component,
		fwhm=fwhm
		)

		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}
		int_intensities = compute_integrated_intensities(amplitudes=fit_params['amplitudes'],
		fwhm_hz=fit_params['fwhm_hz'])
		snr_params = compute_snr_params(integrated_intensities=int_intensities,
		@@ -518,6 +606,7 @@ for source in sources:
		spec_rms=spec_rms,
		rms_integral=rms_integral)

		reliability_score[region][vel_component] = {'pars': {}, 'score': 0}
		ratio_posterior_8786, ratio_posterior_summary_8786 = compute_ratio_posteriors(
		intensity1=int_intensities['86'],
		intensity2=int_intensities['87'],
		@@ -553,6 +642,7 @@ for source in sources:
		densities[region]['best_fit'].append(best)
		densities[region]['uncertainty'].append(hpd)
		densities[region]['reliability_score'].append(score)
		densities[region]['fwhm'].append(fwhm)
		fwhms[region].append(fwhm)

		# plt.clf()
		@@ -561,27 +651,15 @@ for source in sources:
		# plt.plot(density_axis, density_pdfs['88-86']['pdf'])
		# plt.show()

		# plt.clf()
		# 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 = WCS(hdu_cont.header)
		# fig = plt.figure()

		ax = fig.add_subplot(121, projection=wcs)
		plt.imshow(np.log10(moment_map.data), origin='lower')
		plt.colorbar(label=r'log(K km s$^{-1}$)')
		patch = Circle(_region_coordinates[0:2], _region_coordinates[2], edgecolor='black', facecolor='none')
		ax.set_autoscale_on(False)
		ax.add_patch(patch)
		plt.contour(continuum_data,
		levels=np.logspace(np.log10(2.5e-4), np.log10(0.9 * np.nanmax(continuum_data)), num=5),
		transform=ax.get_transform(wcs_cont),
		colors='red',
		origin='lower')
		plt.ylabel('Dec[J2000]')
		plt.xlabel('RA[J2000]')
		plt.tight_layout()
		plt.savefig(os.path.join(plotdir, f'aperture_spectra_extraction_{source}_{region}_{species}.svg'))
		overlay_region_on_map(
		apspec=apspec,
		moment_map=moment_map,
		wcs=wcs,
		region_coordinates=_region_coordinates,
		continuum_data=continuum_data,
		wcs_cont=wcs_cont,
		plotfile=os.path.join(plotdir, f"aperture_spectra_extraction_{source}_{region}_{species}.svg")
		)

		fig = plt.figure(figsize=(7, 5))
		plt.clf()