Added missing file

import os

import warnings

import numpy as np

import matplotlib.pyplot as plt

import matplotlib

import pyspeckit

import pickle

from astropy.constants import c as speed_of_light

from etl.assets.commons import validate_parameter

import yaml

import scipy

matplotlib.use('Qt5Agg')

    return bic

def convert_dvelocity_to_hz(dvelocity):

    return dvelocity / speed_of_light.to(u.km / u.s) * rest_frequencies[0].to(u.Hz)

def create_limits(par_number, channelwidth_hz, rms):

    limits = []

    for idx in range(par_number):

        if idx % 3 == 0:

            limits.append((rms, 10))

        elif idx % 3 == 2:

            limits.append((channelwidth_hz, 20))

        else:

            limits.append((0, 0))

    limited = [(True, False) if (idx % 3 == 0) or (idx % 3 == 2) else (False, False) for idx in range(par_number)]

    return limits, limited

def fill_tied_pars(tied_pars):

    with open(os.path.join('assets', 'templates', 'tied_3mm.template'), 'r') as template:

        _tied_pars_list = template.read().replace('\n', '').format(**tied_pars).split(',')

    return _tied_pars_list

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]

    return guesses

def update_results(fitted_cube, final_pars, indices_to_keep, mask, rms, components, bic_ratio):

    _final_pars = final_pars.copy()

    fit_results = fitted_cube.parcube.copy()

    detections = mask.copy()

    for idx_to_change in range(components + 1):

        velocity_map = (fitted_cube.parcube[1 + (12 * idx_to_change), :, :] * u.Hz).to(u.km / u.s,

                                                                                       equivalencies=freq_to_vel).value

        linewidth_map = 2.355 * ((fitted_cube.parcube[2 + (12 * idx_to_change), :, :] * u.Hz) / rest_frequencies[0].to(

            u.Hz)).value * speed_of_light.to(u.km / u.s).value

        fit_results[1 + (12 * idx_to_change), :, :] = velocity_map

        fit_results[2 + (12 * idx_to_change), :, :] = linewidth_map

        detections = (detections) & (fit_results[0 + (12 * idx_to_change), :, :] >= 2 * rms)

    _final_pars[0:6 + 6 * components, :, :] = np.where((bic_ratio > 1) & mask & detections,

                                                             fit_results[indices_to_keep, :, :],

                                                             final_pars[0:6 + 6 * components, :, :])

    return _final_pars

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

# filename = 'G332.96_CH3OH_3mm_feather.fits'

# filename = 'G28_CH3OH_3mm_feather.fits'

# filename = 'G12_CH3OH_3mm_ACA_MA.fits'

filename = 'G14_CH3OH_3mm_feather.fits'

# filename = 'test_subcube.fits'

with open(os.path.join(datadir, 'source_properties.yml'), 'r') as infile:

    source_properties = yaml.load(infile, Loader=yaml.FullLoader)

source = 'G12'

species = 'CH3OH'

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

vlsr = source_properties[source][species]['vlsr'] * u.km / u.s

rms = source_properties[source][species]['rms']

frequency_limits = source_properties[source][species]['frequency_limits_ghz']

cube = pyspeckit.Cube(os.path.join(datadir, filename), format='fits')

smooth_data = False

if smooth_data is True:

    smoothed_data = np.ones_like(cube.cube)

    for i in range(cube.cube.shape[0]):

        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])

# vlsr = -67.6 * u.km / u.s

# vlsr = 79.6 * u.km / u.s

# vlsr = 34.4 * u.km / u.s

vlsr = 40.5 * u.km / u.s

frequency_offset = vlsr.to(u.Hz, equivalencies=freq_to_vel) - rest_frequencies[0]

doppler_adjusted_frequencies = rest_frequencies + frequency_offset

#cube = cube.slice(96.755, 96.775, unit=u.GHz)

# cube = cube.slice(96.700, 96.740, unit=u.GHz)

cube = cube.slice(96.710, 96.750, unit=u.GHz)

if frequency_limits is not None:

    cube = cube.slice(frequency_limits[0], frequency_limits[1], unit=u.GHz)

delta_freq = rest_frequencies.diff()

# rms = 0.0034

# rms = 0.003

# rms = 0.002

rms = 0.0032

# Test double component

# cube.cube[:, 0, 0] = np.roll(cube.cube[:, 0, 0], -15) + cube.cube[:, 0, 0]

spec_sum = np.nansum(cube.cube, axis=0)

# spec_sum = -np.where(spec_sum < 0, 0, spec_sum)

# spec_sum[255:275, 50:70] *= -1

threshold_spec_sum = 4 * np.sqrt(cube.shape[0]) * rms

starting_point = np.unravel_index(np.argmax(spec_sum), spec_sum.shape)

fit_mask = np.where(spec_sum > threshold_spec_sum, True, False)

linewidth_guess = 1.5 * u.km / u.s

linewidth_guess_hz = linewidth_guess / speed_of_light.to(u.km / u.s) * rest_frequencies[0].to(u.Hz).value / 2.355

tied_pars = {

    'PAR_VELO': 1,

    '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

}

with open(os.path.join('assets', 'templates', 'tied_3mm.template'), 'r') as template:

    tied = template.read().replace('\n', '').format(**tied_pars).split(',')

limits = [(rms, 10) for idx in range(len(tied))]

limited = [(True, False) for idx in range(len(tied))]

guesses = [0.029, doppler_adjusted_frequencies[0].to(u.Hz).value, linewidth_guess_hz,

           0.029, doppler_adjusted_frequencies[1].to(u.Hz).value, linewidth_guess_hz,

           0.01, doppler_adjusted_frequencies[2].to(u.Hz).value, linewidth_guess_hz,

           0.005, doppler_adjusted_frequencies[3].to(u.Hz).value, linewidth_guess_hz]

cube.fiteach(guesses=guesses,

             start_from_point=(starting_point[1], starting_point[0]),

             tied=tied,

             limits=limits,

             limited=limited,

             maskmap=fit_mask,

             use_nearest_as_guess=True)

# cube.plot_spectrum(6, 6)

# cube.plot_fit(6, 6)

residuals = cube.cube - cube.get_modelcube(update=True)

number_of_parameters = tied.count('')

bic = get_bic(residuals=residuals,

              rms=rms,

              number_of_parameters=number_of_parameters,

              number_of_samples=cube.shape[0],

              with_likehood=True)

threshold = (1.5 * 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)

fit_results = cube.parcube.copy()

for idx_frequency in range(number_of_parameters // 6):

    idx_shift = idx_frequency * 12

    velocity_map = (cube.parcube[1 + idx_shift, :, :] * u.Hz).to(u.km / u.s, equivalencies=freq_to_vel).value

    linewidth_map = 2.355 * (

            (cube.parcube[2 + idx_shift, :, :] * u.Hz) / rest_frequencies[0].to(u.Hz)).value * speed_of_light.to(

        u.km / u.s).value

    linewidth_map = np.where(linewidth_map == 0, np.nan, linewidth_map)

    velocity_map = np.where(velocity_map == 0, np.nan, velocity_map)

    ratio_map_87_86 = (cube.parcube[6 + idx_shift, :, :] / cube.parcube[idx_shift, :, :])

    fit_results[1 + idx_shift, :, :] = velocity_map

    fit_results[2 + idx_shift, :, :] = linewidth_map

# plt.clf()

# plt.imshow(velocity_map.value)

# plt.colorbar()

# plt.show()

# plt.clf()

# plt.imshow(linewidth_map.value)

# plt.colorbar()

# plt.show()

# plt.clf()

# plt.imshow(ratio_map_87_86)

# plt.colorbar()

# plt.show()

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

final_pars = np.zeros((number_of_parameters * max_components, cube.parcube.shape[1], cube.parcube.shape[2])) * np.nan

indices_to_keep = np.array((0, 1, 2, 3, 6, 9))

_indices_to_keep = indices_to_keep.copy()

final_pars[0:number_of_parameters, :, :] = np.where((fit_mask == True) & (fit_results[0, :, :] > 0),

                                                    fit_results[indices_to_keep, :, :], np.nan)

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

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

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=[]

for components in range(max_components - 1):

    refit_mask = np.where(bic > threshold, True, False)

for components in range(max_components):

    final_pars_old = final_pars.copy()

    if np.sum(refit_mask) > 0:

        idx_shift = (components + 1) * 12

    if np.sum(mask) > 0:

        idx_shift = components * 12

        tied_pars = {

            'PAR_VELO': 1 + idx_shift,

            'DELTA_FREQ_3': delta_freq[2].to(u.Hz).value,

            'PAR_SIGMA_LINE': 2 + idx_shift

        }

        with open(os.path.join('assets', 'templates', 'tied_3mm.template'), 'r') as template:

            tied += template.read().replace('\n', '').format(**tied_pars).split(',')

        tied += fill_tied_pars(tied_pars=tied_pars)

        guesses += [0.029, doppler_adjusted_frequencies[0].to(u.Hz).value + 1e5 * idx_shift, linewidth_guess_hz,

                    0.029, doppler_adjusted_frequencies[1].to(u.Hz).value + 1e5 * idx_shift, linewidth_guess_hz,

                    0.01, doppler_adjusted_frequencies[2].to(u.Hz).value + 1e5 * idx_shift, linewidth_guess_hz,

                    0.005, doppler_adjusted_frequencies[3].to(u.Hz).value + 1e5 * idx_shift, linewidth_guess_hz]

        limits = [(rms, 0) for idx in range(len(tied))]

        limited = [(True, False) for idx in range(len(tied))]

        guesses += fill_guesses(line_peak=0.029,

                                doppler_adjusted_frequencies=doppler_adjusted_frequencies,

                                shift_hz=2 * linewidth_guess_hz * idx_shift,

                                linewidth_guess_hz=linewidth_guess_hz,

                                rms=rms)

        limits, limited = create_limits(par_number=len(tied),

                                        channelwidth_hz=channelwidth_hz,

                                        rms=rms)

        indices_to_keep = np.concatenate((indices_to_keep, (independent_indices + idx_shift)))

        starting_point = np.unravel_index(np.argmax(np.where(refit_mask == True, spec_sum, -1)), spec_sum.shape)

        starting_point = np.unravel_index(np.argmax(np.where(mask == True, spec_sum, -1)), spec_sum.shape)

        cube.fiteach(guesses=guesses,

                     start_from_point=(starting_point[1], starting_point[0]),

                     tied=tied,

                     limits=limits,

                     limited=limited,

                     maskmap=refit_mask,

                     maskmap=mask,

                     use_nearest_as_guess=True)

        # cube.plot_spectrum(x=starting_point[1], y=starting_point[0])

        # cube.plot_fit(x=starting_point[1], y=starting_point[0])

        residuals = cube.cube - cube.get_modelcube(update=True)

        number_of_parameters = tied.count('')

        new_bic = get_bic(residuals=residuals,

                          number_of_samples=cube.shape[0],

                          with_likehood=True)

        fit_results = cube.parcube.copy()

        for idx_to_change in range(components + 2):

            velocity_map = (cube.parcube[1 + (12 * idx_to_change), :, :] * u.Hz).to(u.km / u.s,

                                                                                    equivalencies=freq_to_vel).value

            linewidth_map = 2.355 * ((cube.parcube[2 + (12 * idx_to_change), :, :] * u.Hz) / rest_frequencies[0].to(

                u.Hz)).value * speed_of_light.to(

                u.km / u.s).value

            ratio_map_87_86 = (cube.parcube[6 + (12 * idx_to_change), :, :] / cube.parcube[(12 * idx_to_change), :, :])

            fit_results[1 + (12 * idx_to_change), :, :] = velocity_map

            fit_results[2 + (12 * idx_to_change), :, :] = linewidth_map

        _indices_to_keep = np.concatenate((_indices_to_keep, (indices_to_keep + idx_shift)))

        final_pars[0:6 + 6 * (components + 1), :, :] = np.where((bic / new_bic > 1) & (refit_mask),

                                                                fit_results[

                                                                _indices_to_keep, :, :],

                                                                final_pars_old[0:6 + 6 * (components + 1), :, :])

        final_pars = update_results(fitted_cube=cube,

                                    final_pars=final_pars,

                                    indices_to_keep=indices_to_keep,

                                    mask=mask,

                                    rms=rms,

                                    components=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)

                     + cube.shape[0] * np.log(rms ** 2) + np.log(cube.shape[0]) * number_of_parameters)

        mask = np.where(bic > threshold, True, False)

with open('final_results.pickle', 'wb') as handle:

cube.write_fit(f'final_fit_{source}_{species}.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()