Commit 4c654b37 authored by Andrea Giannetti's avatar Andrea Giannetti
Browse files

Pre-refactoring

parent 5143027c

etl/mdl/mdl_psk_decomposition.py

+71 −28
Changes for etl/mdl/mdl_psk_decomposition.py: 71 added lines, 28 removed lines.
Original line number Diff line number Diff line
@@ -10,27 +10,57 @@ from astropy.constants import c as speed_of_light 


matplotlib.use('Qt5Agg')





def get_bic(residuals: np.array,

            rms: float,

            number_of_parameters: int,

            number_of_samples: int,

            with_likehood: bool = False):

    if with_likehood is True:

        bic = (np.nansum(residuals ** 2, axis=0) / (rms ** 2) + number_of_samples * np.log(2 * np.pi)

               + number_of_samples * np.log(rms ** 2) + np.log(number_of_samples) * number_of_parameters)

    else:

        bic = (number_of_samples * np.log(np.nansum(residuals ** 2, axis=0) / number_of_samples)

               + np.log(number_of_samples) * number_of_parameters)

    return bic





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

filename = 'G332.96_CH3OH_3mm_feather.fits'

# 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'

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

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

delta_freq = rest_frequencies.diff()

rms = 0.0034

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



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

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,
@@ -40,12 +70,12 @@ tied_pars = { 
}

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

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

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

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, 50000,

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

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

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

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]),
@@ -59,18 +89,24 @@ cube.fiteach(guesses=guesses, 


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

number_of_parameters = tied.count('')

bic = - np.nansum(-residuals ** 2, axis=0) / (rms ** 2) - cube.shape[0] * np.log(np.sqrt(2 * np.pi) * rms) + np.log(

    cube.shape[0]) * number_of_parameters

bic = get_bic(residuals=residuals,

              rms=rms,

              number_of_parameters=number_of_parameters,

              number_of_samples=cube.shape[0],

              with_likehood=True)



threshold = 2 * (cube.shape[0]) - cube.shape[0] * np.log(np.sqrt(2 * np.pi) * rms) + np.log(

    cube.shape[0]) * number_of_parameters

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(

    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
@@ -91,11 +127,13 @@ 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 = np.array((0, 1, 2, 3, 6, 9))

final_pars[0:number_of_parameters, :, :] = np.where(fit_mask, fit_results[indices_to_keep, :, :], np.nan)

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



for components in range(max_components - 1):

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

    final_pars_old = final_pars.copy()

    if np.sum(refit_mask) > 0:

        idx_shift = (components + 1) * 12
@@ -109,14 +147,14 @@ for components in range(max_components - 1): 
        with open(os.path.join('assets', 'templates', 'tied_3mm.template'), 'r') as template:

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



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

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

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

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

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



        starting_point = np.unravel_index(np.argmax(spec_sum[refit_mask]), spec_sum.shape)

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



        cube.fiteach(guesses=guesses,

                     start_from_point=(starting_point[1], starting_point[0]),
@@ -127,13 +165,18 @@ for components in range(max_components - 1): 
                     use_nearest_as_guess=True)

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

        number_of_parameters = tied.count('')

        new_bic = - np.nansum(-residuals ** 2, axis=0) / (rms ** 2) - cube.shape[0] * np.log(

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

        new_bic = get_bic(residuals=residuals,

                          rms=rms,

                          number_of_parameters=number_of_parameters,

                          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(

            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
@@ -142,7 +185,7 @@ for components in range(max_components - 1): 
        final_pars[0:6 + 6 * (components + 1), :, :] = np.where((bic / new_bic > 1) & (refit_mask),

                                                                fit_results[

                                                                _indices_to_keep, :, :],

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

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

        print(bic / new_bic)

        bic = new_bic