Pre-inclusion of N2H+ (d6ec4ca4) · Commits · Andrea Giannetti / CoCoA_methanol_archive

etl/mdl/mdl_psk_decomposition.py

+105 −26

Changes for etl/mdl/mdl_psk_decomposition.py: 105 added lines, 26 removed lines.
Original line number	Diff line number	Diff line
		@@ -5,12 +5,14 @@ import matplotlib
		import astropy.units as u
		import pyspeckit
		import pickle
		from astropy.constants import c as speed_of_light
		from etl.assets.commons import validate_parameter
		import yaml
		import scipy
		from astropy.constants import c as speed_of_light
		from scipy.signal import peak_widths
		from etl.assets.commons import validate_parameter
		from typing import Union, Tuple

		matplotlib.use('Qt5Agg')
		matplotlib.use('TkAgg')


		def get_bic(residuals: np.array,
		@@ -53,8 +55,10 @@ def fill_tied_pars(tied_pars):
		def fill_guesses(line_peak, doppler_adjusted_frequencies, shift_hz, linewidth_guess_hz, rms):
		guesses = [line_peak, doppler_adjusted_frequencies[0].to(u.Hz).value + shift_hz, linewidth_guess_hz,
		line_peak, doppler_adjusted_frequencies[1].to(u.Hz).value + shift_hz, linewidth_guess_hz,
		np.max((0.3 * line_peak, rms)), doppler_adjusted_frequencies[2].to(u.Hz).value + shift_hz, linewidth_guess_hz,
		np.max((0.1 * line_peak, rms)), doppler_adjusted_frequencies[3].to(u.Hz).value + shift_hz, linewidth_guess_hz]
		np.max((0.3 * line_peak, rms)), doppler_adjusted_frequencies[2].to(u.Hz).value + shift_hz,
		linewidth_guess_hz,
		np.max((0.1 * line_peak, rms)), doppler_adjusted_frequencies[3].to(u.Hz).value + shift_hz,
		linewidth_guess_hz]
		return guesses


		@@ -76,15 +80,69 @@ def update_results(fitted_cube, final_pars, indices_to_keep, mask, rms, componen
		return _final_pars


		def get_guesses(peaks: Tuple[np.array, dict],
		ea_channel_separation: float,
		ref_freq: float,
		ref_chan: float,
		channelwidth_hz: float,
		rest_freq: float,
		guesses: list) -> Tuple[float, float, float]:
		_peak_shifts_channels = peaks[0][::-1]
		_widths_channels = peaks[1]['widths'][::-1]
		comps = len(guesses) // 12
		_guess_chans = []
		for i in range(comps):
		_guess_chans.append((guesses[1 + 12 * i] - ref_freq) / channelwidth_hz + ref_chan)
		mask = np.full(shape=_peak_shifts_channels.shape, fill_value=False)
		is_e = []
		for component_shift, component_width in zip(_peak_shifts_channels, _widths_channels):
		tolerance = np.min((ea_channel_separation / 2, component_width))
		shifts = _peak_shifts_channels - component_shift
		mask = mask \| np.where(np.abs(shifts + ea_channel_separation) <= tolerance, True, False)
		if np.sum(mask) == 1:
		is_e.append(True)
		else:
		is_e.append(False)
		print(component_shift, np.abs(shifts + ea_channel_separation), tolerance)
		_peak_shifts_channels = np.where(~mask, _peak_shifts_channels, -9999)

		mask_guesses = np.full(shape=_peak_shifts_channels.shape, fill_value=False)
		for i in range(comps):
		mask_guesses = mask_guesses \| np.where(
		np.min(np.abs(np.where(mask_guesses, -9999, _peak_shifts_channels) - _guess_chans[i])) == np.abs(
		_peak_shifts_channels - _guess_chans[i]), True, False)
		_peak_heights = np.where(~mask[::-1] & ~mask_guesses[::-1], peaks[1]['peak_heights'], 0)
		try:
		next_peak = np.argmax(_peak_heights)
		additional_shift = 0
		if not is_e[-next_peak - 1]:
		additional_shift = ea_channel_separation * np.abs(channelwidth_hz)
		return (_peak_heights[next_peak],
		(_peak_shifts_channels[::-1][
		next_peak] - ref_chan) * channelwidth_hz + ref_freq - additional_shift - rest_freq,
		_widths_channels[::-1][next_peak] * np.abs(channelwidth_hz) / 2.355)
		except IndexError:
		return (np.nanmean(peaks[1]['peak_heights']), 3 * np.abs(channelwidth_hz), 2 * np.abs(channelwidth_hz))


		datadir = os.path.join('assets', 'data')

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

		source = 'G12'
		species = 'CH3OH'
		source_names = ['G332.96', 'G28', 'G14']
		source_names = ['G332.96']
		linewidth_guess = 1.5 * u.km / u.s
		max_components = 3
		rest_frequencies = {
		'CH3OH': np.array([96.739358, 96.74142, 96.744545, 96.755501]) * u.GHz,
		'N2H+': np.array([93.1958,]) * u.GHz
		}

		for source in source_names:
		filename = source_properties[source][species]['data']
		vlsr = source_properties[source][species]['vlsr'] * u.km / u.s
		vlsr = source_properties[source]['vlsr'] * u.km / u.s
		rms = source_properties[source][species]['rms']
		frequency_limits = source_properties[source][species]['frequency_limits_ghz']

		@@ -97,39 +155,48 @@ if smooth_data is True:
		smoothed_data[i, :, :] = scipy.ndimage.gaussian_filter(cube.cube[i, :, :], sigma=3)
		cube.cube = smoothed_data

		rest_frequencies = np.array([96.739358, 96.74142, 96.744545, 96.755501]) * 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
		if frequency_limits is not None:
		cube = cube.slice(frequency_limits[0], frequency_limits[1], unit=u.GHz)
		delta_freq = rest_frequencies.diff()

		linewidth_guess = 1.5 * u.km / u.s
		linewidth_guess_hz = convert_dvelocity_to_hz(dvelocity=linewidth_guess).value / 2.355
		channelwidth_hz = convert_dvelocity_to_hz(dvelocity=0.4 * u.km / u.s).value

		max_components = 3

		channelwidth_hz = cube.header['CDELT3']
		if len(rest_frequencies) > 1:
		delta_freq = rest_frequencies[1:] - rest_frequencies[0]
		independent_indices = np.array((0, 1, 2, 3, 6, 9))
		indices_to_keep = np.array([], dtype=int)
		independent_pars_per_component = len(independent_indices)
		channel_separation = delta_freq.to(u.Hz).value / np.abs(channelwidth_hz)
		else:
		delta_freq = rest_frequencies[1:] - rest_frequencies[0]
		independent_indices = np.array((0, 1, 2, 3, 6, 9))
		indices_to_keep = np.array([], dtype=int)
		final_pars = np.zeros((6 * max_components, cube.cube.shape[1], cube.cube.shape[2])) * np.nan
		independent_pars_per_component = len(independent_indices)
		pars_per_component = 3 * len(rest_frequencies)
		final_pars = np.zeros((pars_per_component * max_components, cube.cube.shape[1], cube.cube.shape[2])) * np.nan



		spec_sum = np.nansum(cube.cube, axis=0)
		# spec_sum = -np.where(spec_sum < 0, 0, spec_sum)
		# spec_sum[260:270, 55:65] *= -1
		# spec_sum[230:240, 295:305] *= -1
		threshold_spec_sum = 4 * np.sqrt(cube.shape[0]) * rms
		starting_point = np.unravel_index(np.argmax(spec_sum), spec_sum.shape)
		mask = np.where(spec_sum > threshold_spec_sum, True, False)
		bic = -1e10 * np.ones_like(spec_sum)
		guesses = []
		tied = []
		peaks = scipy.signal.find_peaks(cube.cube[:, starting_point[0], starting_point[1]], height=3 * rms, width=3)
		peak_guess = np.max(peaks[1]['peak_heights']) * 0.8
		peak_shift = 0

		for components in range(max_components):
		final_pars_old = final_pars.copy()
		if np.sum(mask) > 0:
		idx_shift = components * 12
		idx_shift = components * pars_per_component

		if len(rest_frequencies) > 1:
		tied_pars = {
		'PAR_VELO': 1 + idx_shift,
		'DELTA_FREQ_1': delta_freq[0].to(u.Hz).value,
		@@ -138,14 +205,22 @@ for components in range(max_components):
		'PAR_SIGMA_LINE': 2 + idx_shift
		}
		tied += fill_tied_pars(tied_pars=tied_pars)

		guesses += fill_guesses(line_peak=0.029,
		else:
		tied = None
		# best_guesses = get_guesses(peaks=peaks,
		# ea_channel_separation=channel_separation[0],
		# ref_freq=cube.header['CRVAL3'],
		# ref_chan=cube.header['CRPIX3'],
		# channelwidth_hz=channelwidth_hz,
		# rest_freq=doppler_adjusted_frequencies[0].to(u.Hz).value,
		# guesses=guesses)
		guesses += fill_guesses(line_peak=peak_guess,
		doppler_adjusted_frequencies=doppler_adjusted_frequencies,
		shift_hz=2 * linewidth_guess_hz * idx_shift,
		shift_hz=linewidth_guess_hz / 2 * idx_shift,
		linewidth_guess_hz=linewidth_guess_hz,
		rms=rms)
		limits, limited = create_limits(par_number=len(tied),
		channelwidth_hz=channelwidth_hz,
		channelwidth_hz=np.abs(channelwidth_hz),
		rms=rms)
		indices_to_keep = np.concatenate((indices_to_keep, (independent_indices + idx_shift)))

		@@ -157,10 +232,15 @@ for components in range(max_components):
		limits=limits,
		limited=limited,
		maskmap=mask,
		use_nearest_as_guess=True)
		use_nearest_as_guess=True,
		use_neighbor_as_guess=True,
		multicore=10)
		# cube.plot_spectrum(x=starting_point[1], y=starting_point[0])
		# cube.plot_fit(x=starting_point[1], y=starting_point[0])
		# matplotlib.use('Qt5Agg')
		residuals = cube.cube - cube.get_modelcube(update=True)
		starting_point = np.unravel_index(np.argmax(np.nansum(residuals, axis=0)), spec_sum.shape)
		peaks = scipy.signal.find_peaks(residuals[:, starting_point[0], starting_point[1]], height=3 * rms, width=3)
		number_of_parameters = tied.count('')
		new_bic = get_bic(residuals=residuals,
		rms=rms,
		@@ -177,16 +257,15 @@ for components in range(max_components):
		bic_ratio=bic / new_bic)
		print(bic / new_bic)
		bic = new_bic
		threshold = (1.5 * cube.shape[0] + cube.shape[0] * np.log(2 * np.pi)
		threshold = (1.1 * cube.shape[0] + cube.shape[0] * np.log(2 * np.pi)
		+ cube.shape[0] * np.log(rms ** 2) + np.log(cube.shape[0]) * number_of_parameters)
		mask = np.where(bic > threshold, True, False)

		cube.write_fit(f'final_fit_{source}_{species}.fits', overwrite=True)
		cube.write_fit(f'final_fit_{source}_{species}_{components}.fits', overwrite=True)

		with open(f'final_results_{source}_{species}.pickle', 'wb') as handle:
		pickle.dump(final_pars, handle, protocol=pickle.HIGHEST_PROTOCOL)


		# spec = cube.get_spectrum(5,5)
		# spec.specfit(guesses = guesses, tied = tied, annotate = False)
		# spec.plotter()