Refactored density estimation script (d3e8292a) · Commits · Andrea Giannetti / CoCoA_methanol_archive

etl/mdl/mdl_extract_ch3oh_spectra.py

+119 −144

Changes for etl/mdl/mdl_extract_ch3oh_spectra.py: 119 added lines, 144 removed lines.
Original line number	Diff line number	Diff line
		@@ -28,7 +28,7 @@ logger = setup_logger(name='AP. SPECTRA', log_level='DEBUG')
		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
		# Latent ratio R with a lower bound to forbid R→0
		R = pm.TruncatedNormal("R", mu=hat_int2 / hat_int1, sigma=0.3, lower=lower_limit_ratio, upper=0.95)

		# data container for I1
		@@ -164,21 +164,6 @@ def find_zeros(x, y):
		return zeros


		sources = ['G14', 'G28', 'G332.96']
		sources = ['G28']
		species = 'CH3OH'
		samples = 1000
		prob_threshold = 0.67
		matplotlib.use('TkAgg')

		plotdir = os.path.join('prs', 'figures', 'aperture_spectra_fits')
		with open(os.path.join('config', 'ch3oh_regions.yml'), 'r') as infile:
		regions = yaml.load(infile, Loader=yaml.FullLoader)

		with open(os.path.join('config', 'guesses.yml'), 'r') as infile:
		guesses = yaml.load(infile, Loader=yaml.FullLoader)


		def compute_density_pdf(density, density_axis=None):
		_density_axis = validate_parameter(density_axis, np.linspace(2.5, 7.5, 500))
		kde = KernelDensity(bandwidth=0.1, kernel='epanechnikov')
		@@ -410,7 +395,10 @@ def combine_density_pdfs(density_pdfs, mix_weight):


		def prepare_fit_guesses_and_ties(source, region, component_count,
		guesses, delta_freq, rms, pars_per_component):
		guesses, delta_freq, rms, pars_per_component,
		doppler_adjusted_frequencies,
		channelwidth_hz,
		region_coordinates,max_linewidth_hz):
		"""
		Build up the `tied` list and `line_guesses` for all components.
		"""
		@@ -421,7 +409,7 @@ def prepare_fit_guesses_and_ties(source, region, component_count,

		tied_pars = {
		'PAR_VELO': 1 + idx,
		'PAR_SIGMA_LINE': 2 + idx,
		'PAR_SIGMA_LINE': 5 + idx,
		}
		for i, df in enumerate(delta_freq, start=1):
		tied_pars[f"DELTA_FREQ_{i}"] = df.to(u.Hz).value
		@@ -439,7 +427,7 @@ def prepare_fit_guesses_and_ties(source, region, component_count,
		limits, limited = create_limits(
		par_number=len(line_guesses),
		channelwidth_hz=abs(channelwidth_hz),
		rms=rms / np.sqrt(np.pi * _region_coordinates[2] ** 2),
		rms=rms / np.sqrt(np.pi * region_coordinates[2] ** 2),
		max_linewidth=max_linewidth_hz
		)
		return tied, line_guesses, limits, limited
		@@ -457,7 +445,12 @@ def run_and_plot_specfit(apspec, tied, line_guesses, limits, limited,
		tied=tied,
		limits=limits,
		limited=limited,
		clear=True
		clear=True,
		xtol=1e-15,
		gtol=1e-15,
		maxiter=5000,
		info=5
		# nocovar=True,
		)

		# plot fit components
		@@ -473,9 +466,10 @@ def run_and_plot_specfit(apspec, tied, line_guesses, limits, limited,


		def overlay_region_on_map(apspec,
		region,
		region_coordinates,
		moment_map,
		wcs,
		region_coordinates,
		continuum_data,
		wcs_cont,
		plotfile):
		@@ -513,58 +507,38 @@ def overlay_region_on_map(apspec,
		fig_map.savefig(plotfile)
		plt.close(fig_map)


		for source in sources:
		def process_source(source, species, guesses):
		# 1. Load maps and prepare line models
		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
		freqs_dop, lw_guess_hz, 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
		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 = []

		# 2. Open cube & continuum
		filename = source_properties[source][species]['data']
		cube = get_cube(filename=filename,
		datadir=datadir,
		frequency_limits=frequency_limits,
		smooth_data=config_mdl['smooth_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)

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

		densities = {}
		fwhms = {}
		reliability_score = {}
		for region in regions[source]:
		logger.info(f'Processing region: {region}')
		_region_coordinates = regions[source][region][0:3]
		components = regions[source][region][3]
		region_guesses = regions[source][region]
		apspec = cube.get_apspec(_region_coordinates)
		line_guesses = []
		tied = []
		# 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': []}
		fwhms[region] = []
		reliability_score[region] = {}
		fwhms[region], reliability_score[region] = [], {}

		# 4a. Build & run the spec‐fit
		apspec = cube.get_apspec(coords)
		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
		source, region, components, guesses, delta_freq, rms, pars_per_component, freqs_dop,
		channelwidth_hz, coords, max_lw_hz
		)

		spec_rms, specfig = run_and_plot_specfit(
		spec_rms, _ = run_and_plot_specfit(
		apspec=apspec,
		tied=tied,
		line_guesses=line_guesses,
		@@ -576,104 +550,87 @@ for source in sources:
		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,
		region=region,
		vel_component=vel_component,
		fwhm=fwhm
		)
		# 4b. Loop over each velocity component
		for vel_comp in range(components):
		idx = vel_comp * 4

		if skipped:
		# extract fit params & convert to physical units
		fitp = extract_fit_parameters(apspec.specfit.parinfo, idx)
		fwhm = convert_dfreq_to_velocity(fitp['fwhm_hz'], freqs_dop)
		if should_skip_component(source, region, vel_comp, fwhm):
		continue

		rms_integral = compute_rms_integral(fwhm_hz=fit_params['fwhm_hz'],
		channelwidth_hz=channelwidth_hz,
		spec_rms=spec_rms)
		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)

		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'],
		sigma=rms_integral,
		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
		)
		# compute SNR, ratios, reliability
		rms_int = compute_rms_integral(fitp['fwhm_hz'], channelwidth_hz, spec_rms)
		ints = compute_integrated_intensities(fitp['amplitudes'], fitp['fwhm_hz'])
		snr_p = compute_snr_params(integrated_intensities=ints, amplitudes=fitp['amplitudes'], spec_rms=spec_rms, rms_integral=rms_int)

		r87, s87 = compute_ratio_posteriors(ints['86'], ints['87'], rms_int, 0.05)
		r88, s88 = compute_ratio_posteriors(ints['86'], ints['88'], rms_int, 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)
		**snr_p,
		'ratio1': np.mean(r87),
		'ratio2': np.mean(r87),
		'rel_prec_ratio1': 1 - np.std(r87)/np.mean(r87),
		'rel_prec_ratio2': 1 - np.std(r88)/np.mean(r88)
		}
		score = compute_reliability_score(component=all_pars)
		reliability_score[region][vel_component] = {'pars': all_pars, 'score': score}

		density_pdfs = {}
		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(best)
		densities[region]['uncertainty'].append(hpd)
		densities[region]['reliability_score'].append(score)
		densities[region]['fwhm'].append(fwhm)
		fwhms[region].append(fwhm)
		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(density_axis, final_density_pdf)
		# plt.plot(density_axis, density_pdfs['87-86']['pdf'])
		# plt.plot(density_axis, density_pdfs['88-86']['pdf'])
		# 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)
		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))
		fwhms[region].append(fwhm.value)

		# 4c. Overlay region on map
		overlay_region_on_map(
		apspec=apspec,
		apspec=apspec, moment_map=moment_map, wcs=wcs, region=region, region_coordinates=coords,
		continuum_data=continuum_data, wcs_cont=wcs_cont,
		plotfile=os.path.join(plotdir, f"{source}_{region}_{species}.svg")
		)

		# 5. Final summary plots & YAML dump
		plot_density_vs_fwhm(densities=densities, fwhms=fwhms, source=source, species=species, regions=regions)
		overlay_regions_density(source=source,
		regions=regions,
		densities=densities,
		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")
		)
		wcs_cont=wcs_cont)
		with open(os.path.join('mdl', 'inferred_densities', f'densities_{source}.yml'), 'w') as outfile:
		yaml.dump(densities, outfile, default_flow_style=False)


		fig = plt.figure(figsize=(7, 5))
		def plot_density_vs_fwhm(densities, fwhms, source, species, regions):
		plt.clf()
		fig = plt.figure()
		for region in regions[source]:
		if not region.startswith('core_ch3oh'):
		color = 'blue' if region.startswith('core') else 'cyan'
		max_idx = np.argmax(densities[region]['best_fit'])
		plt.scatter(densities[region]['best_fit'], fwhms[region], label=region, edgecolors=color)
		plt.scatter(densities[region]['best_fit'], fwhms[region], label=region, edgecolors=color, figure=fig)
		plt.legend()
		plt.savefig(os.path.join(plotdir, f'aperture_spectra_extraction_density_FWHM_{source}_{species}.svg'))
		plt.close(fig)

		with open(os.path.join('mdl', 'inferred_densities', f'densities_{source}.yml'), 'w') as outfile:
		yaml.dump(densities, outfile, default_flow_style=False)

		def overlay_regions_density(source, regions, densities, moment_map, wcs, continuum_data, wcs_cont):
		region_types = ['envelope', 'core']

		for rtype in region_types:
		@@ -696,13 +653,31 @@ for source in sources:
		pc.set_facecolor('none')
		pc.set_clim(4.8, 5.8)
		ax.add_collection(pc)
		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='grey',
		origin='lower')
		plt.ylabel('Dec[J2000]')
		plt.xlabel('RA[J2000]')
		# 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='grey',
		# origin='lower')
		plt.ylabel('Dec')
		plt.xlabel('RA')
		plt.colorbar(pc, label=r'log($n($H$_2)$ [cm$^{-3}$])')
		plt.tight_layout()
		plt.savefig(os.path.join(plotdir, f'aperture_spectra_extraction_{source}_{rtype}_{species}.svg'))
		plt.close()

		sources = ['G14', 'G28', 'G332.96']
		# sources = ['G28']
		species = 'CH3OH'
		samples = 1000
		prob_threshold = 0.67
		matplotlib.use('TkAgg')

		plotdir = os.path.join('prs', 'figures', 'aperture_spectra_fits')
		with open(os.path.join('config', 'ch3oh_regions.yml'), 'r') as infile:
		regions = yaml.load(infile, Loader=yaml.FullLoader)

		with open(os.path.join('config', 'guesses.yml'), 'r') as infile:
		guesses = yaml.load(infile, Loader=yaml.FullLoader)

		for source in sources:
		process_source(source, species, guesses)