hrt_pipe.py 36.3 KB
Newer Older
jonas's avatar
jonas committed
1
from re import L
jonas's avatar
jonas committed
2
3
4
5
6
import numpy as np 
import os.path
from astropy.io import fits
import subprocess
from scipy.ndimage import gaussian_filter
7
import time
jonas's avatar
jonas committed
8
from operator import itemgetter
jonas's avatar
jonas committed
9

Jonas Sinjan's avatar
Jonas Sinjan committed
10
from utils import *
jonas's avatar
jonas committed
11
12


13
def get_data(path, scaling = True):
jonas's avatar
jonas committed
14
15
16
    """load science data from path"""
    try:
        data, header = load_fits(path)
17
18

        if scaling:
jonas's avatar
jonas committed
19
      
20
            data /=  256. #conversion from 24.8bit to 32bit
jonas's avatar
jonas committed
21

22
            accu = header['ACCACCUM']*header['ACCROWIT']*header['ACCCOLIT'] #getting the number of accu from header
jonas's avatar
jonas committed
23

24
            data /= accu
jonas's avatar
jonas committed
25

26
            printc(f"Dividing by number of accumulations: {accu}",color=bcolors.OKGREEN)
Jonas Sinjan's avatar
Jonas Sinjan committed
27
28
        
        return data, header
jonas's avatar
jonas committed
29
30
31
32

    except Exception:
        printc("ERROR, Unable to open fits file: {}",path,color=bcolors.FAIL)

Jonas Sinjan's avatar
Jonas Sinjan committed
33
   
jonas's avatar
jonas committed
34

Jonas Sinjan's avatar
Jonas Sinjan committed
35
def demod_hrt(data,pmp_temp):
jonas's avatar
jonas committed
36
    '''
Jonas Sinjan's avatar
Jonas Sinjan committed
37
    Use constant demodulation matrices to demodulate data
jonas's avatar
jonas committed
38
    '''
Jonas Sinjan's avatar
Jonas Sinjan committed
39
40
41
42
43
44
 
    if pmp_temp == '50':
        demod_data = np.array([[ 0.28037298,  0.18741922,  0.25307596,  0.28119895],
                     [ 0.40408596,  0.10412157, -0.7225681,   0.20825675],
                     [-0.19126636, -0.5348939,   0.08181918,  0.64422774],
                     [-0.56897295,  0.58620095, -0.2579202,   0.2414017 ]])
Jonas Sinjan's avatar
Jonas Sinjan committed
45
        
Jonas Sinjan's avatar
Jonas Sinjan committed
46
47
48
49
50
    elif pmp_temp == '40':
        demod_data = np.array([[ 0.26450154,  0.2839626,   0.12642948,  0.3216773 ],
                     [ 0.59873885,  0.11278069, -0.74991184,  0.03091451],
                     [ 0.10833212, -0.5317737,  -0.1677862,   0.5923593 ],
                     [-0.46916953,  0.47738808, -0.43824592,  0.42579797]])
jonas's avatar
jonas committed
51
    
jonas's avatar
jonas committed
52
    else:
Jonas Sinjan's avatar
Jonas Sinjan committed
53
        printc("Demodulation Matrix for PMP TEMP of {pmp_temp} deg is not available", color = bcolors.FAIL)
jonas's avatar
jonas committed
54

Jonas Sinjan's avatar
Jonas Sinjan committed
55
    printc(f'Using a constant demodulation matrix for a PMP TEMP of {pmp_temp} deg',color = bcolors.OKGREEN)
jonas's avatar
jonas committed
56
    
Jonas Sinjan's avatar
Jonas Sinjan committed
57
58
59
    demod_data = demod_data.reshape((4,4))
    shape = data.shape
    demod = np.tile(demod_data, (shape[0],shape[1],1,1))
jonas's avatar
jonas committed
60
61
62
63

    if data.ndim == 5:
        #if data array has more than one scan
        data = np.moveaxis(data,-1,0) #moving number of scans to first dimension
jonas's avatar
jonas committed
64

jonas's avatar
jonas committed
65
66
        data = np.matmul(demod,data)
        data = np.moveaxis(data,0,-1) #move scans back to the end
jonas's avatar
jonas committed
67
68
69
    
    elif data.ndim == 4:
        #for if data has just one scan
jonas's avatar
jonas committed
70
        data = np.matmul(demod,data)
jonas's avatar
jonas committed
71
    
jonas's avatar
jonas committed
72
    return data, demod
jonas's avatar
jonas committed
73
74
75
    


76
def phihrt_pipe(data_f, dark_f = '', flat_f = '', scale_data = True, norm_f = True, clean_f = False, 
jonas's avatar
jonas committed
77
                sigma = 59, flat_states = 24, prefilter_f = None,flat_c = True, dark_c = True, fs_c = True,
78
79
                demod = True, norm_stokes = True, out_dir = './',  out_demod_file = False,  
                correct_ghost = False, ItoQUV = False, ctalk_params = None, rte = False):
jonas's avatar
jonas committed
80
81
82

    '''
    PHI-HRT data reduction pipeline
jonas's avatar
jonas committed
83
    1. read in science data (+ OPTION: scaling) open path option + open for several scans at once
jonas's avatar
jonas committed
84
85
86
87
    2. read in flat field (+scaling)- just accepts one flat field fits file
    3. read in dark field (+scaling)
    4. apply dark field
    5. option to clean flat field with unsharp masking
jonas's avatar
jonas committed
88
89
    6. normalise flat field
    7. apply flat field
jonas's avatar
jonas committed
90
    8. apply prefilter
jonas's avatar
jonas committed
91
92
93
    9. read in field stop
    10. apply field stop
    11. demodulate with const demod matrix
jonas's avatar
jonas committed
94
        a) option to output demod to fits file
jonas's avatar
jonas committed
95
96
    12. normalise to quiet sun
    13. calibration
jonas's avatar
jonas committed
97
        a) ghost correction - not implemented yet
jonas.sinjan's avatar
jonas.sinjan committed
98
        b) cross talk correction
jonas's avatar
jonas committed
99
    14. rte inversion with cmilos (CE, RTE or CE+RTE)
jonas's avatar
jonas committed
100
        a) output rte data products to fits file
jonas's avatar
jonas committed
101
102
103
104
105

    Parameters
    ----------
        Input:
    data_f : list or string
jonas's avatar
jonas committed
106
        list containing paths to fits files of the raw HRT data OR string of path to one file  - must have last 10 characters before .fits as the DID - for naming purposes of output files
107
    dark_f : string, DEFAULT ''
jonas's avatar
jonas committed
108
        Fits file of a dark file (ONLY ONE FILE)
109
    flat_f : string, DEFAULT ''
jonas's avatar
jonas committed
110
111
112
        Fits file of a HRT flatfield (ONLY ONE FILE)

    ** Options:
113
114
    read_scale_data: bool, DEFAULT True
        reads in science data and performs appropriate scaling
jonas's avatar
jonas committed
115
116
117
    norm_f: bool, DEFAULT: True
        to normalise the flat fields before applying
    clean_f: bool, DEFAULT: False
jonas's avatar
jonas committed
118
119
120
        clean the flat field with unsharp masking
    sigma: int, DEFAULT: 59
        sigma of the gaussian convolution used for unsharp masking if clean_f == True      
jonas's avatar
jonas committed
121
122
    flat_states: int, DEFAULT: 24
        Number of flat fields to be applied, options are 4 (one for each pol state), 6 (one for each wavelength), 24 (one for each image)
jonas's avatar
jonas committed
123
124
    prefilter_f: str, DEFAULT None
        file path location to prefilter fits file, apply prefilter correction
jonas's avatar
jonas committed
125
126
127
128
    flat_c: bool, DEFAULT: True
        apply flat field correction
    dark_c: bool, DEFAULT: True
        apply dark field correction
jonas's avatar
jonas committed
129
    fs_c: bool, DEFAULT True
130
        apply HRT field stop
jonas's avatar
jonas committed
131
132
133
134
135
136
137
138
139
140
141
142
    demod: bool, DEFAULT: True
        apply demodulate to the stokes
    norm_stokes: bool, DEFAULT: True
        normalise the stokes vector to the quiet sun (I_continuum)
    out_dir : string, DEFUALT: './'
        directory for the output files
    out_demod_file: bool, DEFAULT: False
        output file with the stokes vectors to fits file
    correct_ghost: bool, DEFAULT: False 
        correct the ghost in bottom left corner
    ItoQUV: bool, DEFAULT: False 
        apply I -> Q,U,V correction
jonas's avatar
jonas committed
143
144
    ctalk_params: numpy arr, DEFAULT: None 
        cross talk parameters for ItoQUV, (2,3) numpy array required: first axis: Slope, Offset (Normalised to I_c) - second axis:  Q,U,V
jonas's avatar
jonas committed
145
146
    rte: str, DEFAULT: False 
        invert using cmilos, options: 'RTE' for Milne Eddington Inversion, 'CE' for Classical Estimates, 'CE+RTE' for combined
jonas's avatar
jonas committed
147
148
149
    
    Returns
    -------
jonas's avatar
jonas committed
150
151
    data: nump array
        stokes vector 
jonas's avatar
jonas committed
152
153
154
155
156
157
158
159
160
161
162

    References
    ----------
    SPGYlib

    '''

    printc('--------------------------------------------------------------',bcolors.OKGREEN)
    printc('PHI HRT data reduction software  ',bcolors.OKGREEN)
    printc('--------------------------------------------------------------',bcolors.OKGREEN)

Jonas Sinjan's avatar
Jonas Sinjan committed
163
    overall_time = time.time()
jonas's avatar
jonas committed
164

jonas's avatar
jonas committed
165
166
167
    #-----------------
    # READ DATA
    #-----------------
168

Jonas Sinjan's avatar
Jonas Sinjan committed
169
    print(" ")
Jonas Sinjan's avatar
Jonas Sinjan committed
170
    printc('-->>>>>>> Reading Data              ',color=bcolors.OKGREEN) 
jonas's avatar
jonas committed
171

172
    start_time = time.time()
jonas's avatar
jonas committed
173
174

    if isinstance(data_f, list):
175
        #if the data_f contains several scans
Jonas Sinjan's avatar
Jonas Sinjan committed
176
        printc(f'Input contains {len(data_f)} scan(s)',color=bcolors.OKGREEN)
177
178
179
180
181
182
        
        number_of_scans = len(data_f)

        data_arr = [0]*number_of_scans
        hdr_arr = [0]*number_of_scans

jonas's avatar
jonas committed
183
184
        wve_axis_arr = [0]*number_of_scans
        cpos_arr = [0]*number_of_scans
jonas's avatar
jonas committed
185
186
        voltagesData_arr = [0]*number_of_scans
        tuning_constant_arr = [0]*number_of_scans
jonas's avatar
jonas committed
187

188
        for scan in range(number_of_scans):
189
            data_arr[scan], hdr_arr[scan] = get_data(data_f[scan], scaling = scale_data)
190

191
            wve_axis_arr[scan], voltagesData_arr[scan], tuning_constant_arr[scan], cpos_arr[scan] = fits_get_sampling(data_f[scan],verbose = True)
192

193
            if scale_data:
194

195
                if hdr_arr[scan]['BITPIX'] == 16:
196

197
198
199
                    print(f"This scan: {data_f[scan]} has a bits per pixel of: 16 \nPerforming the extra scaling")

                    data_arr[scan] *= 81920/127 #conversion factor if 16 bits
jonas's avatar
jonas committed
200

jonas's avatar
jonas committed
201
202
            if 'IMGDIRX' in hdr_arr[scan] and hdr_arr[scan]['IMGDIRX'] == 'YES':
                print("This scan has been flipped in the Y axis to conform to orientation standards.")
jonas's avatar
jonas committed
203
204

        #--------
jonas's avatar
jonas committed
205
        # test if the scans have different sizes
jonas's avatar
jonas committed
206
207
        #--------

208
209
210
        first_shape = data_arr[scan].shape
        result = all(element.shape == first_shape for element in data_arr)
        if (result):
Jonas Sinjan's avatar
Jonas Sinjan committed
211
            print("All the scans have the same dimension")
jonas's avatar
jonas committed
212

213
214
215
216
217
        else:
            print("The scans have different dimensions! \n Ending process")

            exit()

jonas's avatar
jonas committed
218
219

        #--------
jonas's avatar
jonas committed
220
        # test if the scans have different continuum wavelength_positions
jonas's avatar
jonas committed
221
222
223
224
225
        #--------

        first_cpos = cpos_arr[0]
        result = all(c_position == first_cpos for c_position in cpos_arr)
        if (result):
226
            print("All the scans have the same continuum wavelength position")
jonas's avatar
jonas committed
227
228

        else:
229
230
231
            print("The scans have different continuum_wavelength postitions! Please fix \n Ending Process")

            exit()
jonas's avatar
jonas committed
232

jonas's avatar
jonas committed
233
        #--------
jonas's avatar
jonas committed
234
        # test if the scans have different pmp temperatures
jonas's avatar
jonas committed
235
236
237
238
239
240
241
242
243
244
245
246
247
        #--------

        first_pmp_temp = hdr_arr[0]['HPMPTSP1']
        result = all(hdr['HPMPTSP1'] == first_pmp_temp for hdr in hdr_arr)
        if (result):
            print(f"All the scans have the same PMP Temperature Set Point: {first_pmp_temp}")
            pmp_temp = str(first_pmp_temp)

        else:
            print("The scans have different PMP Temperatures! Please fix \n Ending Process")

            exit()

248
249
        data = np.stack(data_arr, axis = -1)
        data = np.moveaxis(data, 0,-2) #so that it is [y,x,24,scans]
jonas's avatar
jonas committed
250

251
        print(f"Data shape is {data.shape}")
jonas's avatar
jonas committed
252
253

    elif isinstance(data_f, str):
254
        #case when data f is just one file
255
        data, header = get_data(data_f, scaling = scale_data)
256
257
258
259
260
261
262
263
264
        data = np.expand_dims(data, axis = -1) #so that it has the same dimensions as several scans
        data = np.moveaxis(data, 0, -2) #so that it is [y,x,24,1]

        if header['BITPIX'] == 16:
    
            print(f"This scan: {data_f} has a bits per pixel is: 16 \n Performing the extra scaling")

            data *= 81920/127 #conversion factor if 16 bits

265

266
267
        print(f"Data shape is {data.shape}")

jonas's avatar
jonas committed
268
        wave_axis, voltagesData, tuning_constant, cpos = fits_get_sampling(data_f,verbose = True)
jonas's avatar
jonas committed
269
270
271

        print("The data continuum wavelength position is at index: ", cpos)

jonas's avatar
jonas committed
272
        cpos_arr = [cpos]
jonas's avatar
jonas committed
273
274
275
276
277
278

        if cpos_arr[0] != 0 and cpos_arr[0] != 5:
            print("Data continuum position not at 0 or 5th index. Please reconcile. \n Ending Process")

            exit()

jonas's avatar
jonas committed
279
280
281
        hdr_arr = [header]
        voltagesData_arr = [voltagesData]
        tuning_constant_arr = [tuning_constant]
jonas's avatar
jonas committed
282

jonas's avatar
jonas committed
283
284
285
        pmp_temp = str(header['HPMPTSP1'])
        print(f"Data PMP Set Point Temperature is {pmp_temp}")

jonas's avatar
jonas committed
286
    else:
287
288
289
        printc("ERROR, data_f argument is neither a string nor list containing strings: {} \n Ending Process",data_f,color=bcolors.FAIL)
        exit()

290
291
292
293
294
295
296
297
298
299
    data_shape = data.shape

    data_size = data_shape[:2]
    
    #converting to [y,x,pol,wv,scans]

    data = data.reshape(data_size[0],data_size[1],6,4,data_shape[-1]) #separate 24 images, into 6 wavelengths, with each 4 pol states
    data = np.moveaxis(data, 2,-2) #need to swap back to work

    #enabling cropped datasets, so that the correct regions of the dark field and flat field are applied
jonas.sinjan's avatar
jonas.sinjan committed
300
    print("Data reshaped to: ", data.shape)
301

Jonas Sinjan's avatar
Jonas Sinjan committed
302
    diff = 2048-data_size[0] #handling 0/2 errors
303
304
305
306
307
308
309
310
    
    if np.abs(diff) > 0:
    
        start_row = int((2048-data_size[0])/2)
        start_col = int((2048-data_size[1])/2)
        
    else:
        start_row, start_col = 0, 0
Jonas Sinjan's avatar
Jonas Sinjan committed
311
312
313
314
        
    printc('--------------------------------------------------------------',bcolors.OKGREEN)
    printc(f"------------ Load science data time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
    printc('--------------------------------------------------------------',bcolors.OKGREEN)
jonas.sinjan's avatar
jonas.sinjan committed
315

jonas's avatar
jonas committed
316
317
318
319
320
    #-----------------
    # READ FLAT FIELDS
    #-----------------

    if flat_c:
Jonas Sinjan's avatar
Jonas Sinjan committed
321
        print(" ")
322
        printc('-->>>>>>> Reading Flats',color=bcolors.OKGREEN)
jonas's avatar
jonas committed
323

324
325
        start_time = time.time()
    
jonas's avatar
jonas committed
326
        try:
Jonas Sinjan's avatar
Jonas Sinjan committed
327
            flat, header_flat = get_data(flat_f)
328
329
330

            print(f"Flat field shape is {flat.shape}")
            
Jonas Sinjan's avatar
Jonas Sinjan committed
331
            if header_flat['BITPIX'] == 16:
332
333
334
335
336
337
    
                print("Number of bits per pixel is: 16")

                flat *= 614400/128

            flat = np.moveaxis(flat, 0,-1) #so that it is [y,x,24]
338
            flat = flat.reshape(2048,2048,6,4) #separate 24 images, into 6 wavelengths, with each 4 pol states
Jonas Sinjan's avatar
Jonas Sinjan committed
339
340
            flat = np.moveaxis(flat, 2,-1)
            
jonas.sinjan's avatar
jonas.sinjan committed
341
            print(flat.shape)
jonas's avatar
jonas committed
342
343
344
345

            _, _, _, cpos_f = fits_get_sampling(flat_f,verbose = True)

            print(f"The continuum position of the flat field is at {cpos_f} index position")
Jonas Sinjan's avatar
Jonas Sinjan committed
346
347
348
349
350
351
            
            if flat_f[-62:] == 'solo_L0_phi-hrt-flat_0667134081_V202103221851C_0162201100.fits':
                
                print("This flat needs to be rolled")
                
                flat = np.roll(flat, 1, axis = -1)
352

jonas's avatar
jonas committed
353
354
355
356
                cpos_f = cpos_arr[0]

            #if continuum wavelength of the flat is not the same as the data, attempt to roll
            if cpos_f != cpos_arr[0]:
jonas's avatar
jonas committed
357
                print("The flat field continuum position is not the same as the data, please check your input data. \n Ending Process")
jonas's avatar
jonas committed
358
359

                exit()
jonas's avatar
jonas committed
360
361
362

            flat_pmp_temp = str(header_flat['HPMPTSP1'])

jonas's avatar
jonas committed
363
            print(f"Flat PMP Temperature Set Point: {flat_pmp_temp}")
jonas's avatar
jonas committed
364
                
Jonas Sinjan's avatar
Jonas Sinjan committed
365
366
367
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
            printc(f"------------ Load flats time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
368

jonas's avatar
jonas committed
369
370
371
372
373
        except Exception:
            printc("ERROR, Unable to open flats file: {}",flat_f,color=bcolors.FAIL)


    else:
Jonas Sinjan's avatar
Jonas Sinjan committed
374
        print(" ")
375
376
        printc('-->>>>>>> No flats mode',color=bcolors.WARNING)

jonas's avatar
jonas committed
377
378
379
380

    #-----------------
    # READ AND CORRECT DARK FIELD
    #-----------------
381

jonas's avatar
jonas committed
382
    if dark_c:
Jonas Sinjan's avatar
Jonas Sinjan committed
383
        print(" ")
jonas's avatar
jonas committed
384
        printc('-->>>>>>> Reading Darks                   ',color=bcolors.OKGREEN)
385
386
387

        start_time = time.time()

jonas's avatar
jonas committed
388
        try:
jonas's avatar
jonas committed
389
390
391
392
393
394
395
396
397
398
399
400
401
402
            dark,h = get_data(dark_f)

            dark_shape = dark.shape

            if dark_shape != (2048,2048):
                
                printc("Dark Field Input File not in 2048,2048 format: {}",dark_f,color=bcolors.WARNING)
                printc("Attempting to correct ",color=bcolors.WARNING)

                
                try:
                    if dark_shape[0] > 2048:
                        dark = dark[dark_shape[0]-2048:,:]
                
403
                except Exception:
jonas's avatar
jonas committed
404
405
                    printc("ERROR, Unable to correct shape of dark field data: {}",dark_f,color=bcolors.FAIL)

Jonas Sinjan's avatar
Jonas Sinjan committed
406
407
408
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
            printc(f"------------ Load darks time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
409

jonas's avatar
jonas committed
410
411
412
        except Exception:
            printc("ERROR, Unable to open darks file: {}",dark_f,color=bcolors.FAIL)

jonas's avatar
jonas committed
413

414
415
416
        #-----------------
        # APPLY DARK CORRECTION 
        #-----------------    
Jonas Sinjan's avatar
Jonas Sinjan committed
417
        print(" ")
jonas's avatar
jonas committed
418
        print("-->>>>>>> Subtracting dark field")
Jonas Sinjan's avatar
Jonas Sinjan committed
419
        
420
421
        start_time = time.time()

Jonas Sinjan's avatar
Jonas Sinjan committed
422
        flat -= dark[..., np.newaxis, np.newaxis]
423
        
Jonas Sinjan's avatar
Jonas Sinjan committed
424
        data -= dark[start_row:start_row + data_size[0],start_col:start_col + data_size[1], np.newaxis, np.newaxis, np.newaxis] 
425

Jonas Sinjan's avatar
Jonas Sinjan committed
426
427
428
        printc('--------------------------------------------------------------',bcolors.OKGREEN)
        printc(f"------------- Dark Field correction time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
        printc('--------------------------------------------------------------',bcolors.OKGREEN)
jonas's avatar
jonas committed
429

430
431
432
    else:
        print(" ")
        printc('-->>>>>>> No dark mode',color=bcolors.WARNING)
jonas's avatar
jonas committed
433

jonas's avatar
jonas committed
434

jonas's avatar
jonas committed
435
436
437
438
439
440
441
442
443
444
445
446
447
    #-----------------
    # OPTIONAL Unsharp Masking clean the flat field stokes V images
    #-----------------

    if clean_f:
        print(" ")
        printc('-->>>>>>> Cleaning flats with Unsharp Masking (Stokes V only)',color=bcolors.OKGREEN)

        start_time = time.time()

        #cleaning the stripe in the flats for a particular flat

        if flat_f[-62:] == 'solo_L0_phi-hrt-flat_0667134081_V202103221851C_0162201100.fits': 
jonas's avatar
jonas committed
448
449
            flat[1345, 296:, 1, 2] = flat[1344, 296:, 1, 2]
            flat[1346, :291, 1, 2] = flat[1345, :291, 1, 2]
jonas's avatar
jonas committed
450
451
452

        #demod the flats

jonas's avatar
jonas committed
453
454
455
        

        flat_demod, demodM = demod_hrt(flat, flat_pmp_temp)
jonas's avatar
jonas committed
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490

        norm_factor = norm_factor = np.mean(flat_demod[512:1536,512:1536,0,0])

        flat_demod /= norm_factor

        new_demod_flats = np.copy(flat_demod)

        b_arr = np.zeros((2048,2048,3,5))

        if cpos_arr[0] == 0:
            wv_range = range(1,6)

        elif cpos_arr[0] == 5:
            wv_range = range(5)

        for pol in range(3,4):

            for wv in wv_range: #not the continuum

                a = np.copy(np.clip(flat_demod[:,:,pol,wv], -0.02, 0.02))
                b = a - gaussian_filter(a,sigma)
                b_arr[:,:,pol-1,wv-1] = b
                c = a - b

                new_demod_flats[:,:,pol,wv] = c

        invM = np.linalg.inv(demodM)

        flat = np.matmul(invM, new_demod_flats*norm_factor)


        printc('--------------------------------------------------------------',bcolors.OKGREEN)
        printc(f"------------- Cleaning flat time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
        printc('--------------------------------------------------------------',bcolors.OKGREEN)

491
492
493
    else:
        print(" ")
        printc('-->>>>>>> No clean flats mode',color=bcolors.WARNING)
jonas's avatar
jonas committed
494
495


496
497
498
499
500
    #-----------------
    # NORM FLAT FIELDS
    #-----------------

    if norm_f and flat_c:
Jonas Sinjan's avatar
Jonas Sinjan committed
501
502
        
        print(" ")
503
504
505
506
507
        printc('-->>>>>>> Normalising Flats',color=bcolors.OKGREEN)

        start_time = time.time()

        try:
Jonas Sinjan's avatar
Jonas Sinjan committed
508
509
            norm_fac = np.mean(flat[512:1536,512:1536, :, :], axis = (0,1))  #mean of the central 1k x 1k
            flat /= norm_fac[np.newaxis, np.newaxis, ...]
510

Jonas Sinjan's avatar
Jonas Sinjan committed
511
512
513
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
            printc(f"------------- Normalising flat time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
514
515
516
517

        except Exception:
            printc("ERROR, Unable to normalise the flat fields: {}",flat_f,color=bcolors.FAIL)

518
519
520
521
    else:
        print(" ")
        printc('-->>>>>>> No normalising flats mode',color=bcolors.WARNING)

jonas's avatar
jonas committed
522

jonas's avatar
jonas committed
523
524
525
526
527
    #-----------------
    # APPLY FLAT CORRECTION 
    #-----------------

    if flat_c:
Jonas Sinjan's avatar
Jonas Sinjan committed
528
        print(" ")
jonas's avatar
jonas committed
529
        printc('-->>>>>>> Correcting Flatfield',color=bcolors.OKGREEN)
530
531
532

        start_time = time.time()

jonas's avatar
jonas committed
533
        try:
534
535

            if flat_states == 6:
jonas's avatar
jonas committed
536
        
537
538
539
540
                printc("Dividing by 6 flats, one for each wavelength",color=bcolors.OKGREEN)
                    
                tmp = np.mean(flat,axis=-2) #avg over pol states for the wavelength

Jonas Sinjan's avatar
Jonas Sinjan committed
541
                data /= tmp[start_row:start_row + data_size[0],start_col:start_col + data_size[1], np.newaxis, np.newaxis]
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556


            elif flat_states == 24:

                printc("Dividing by 24 flats, one for each image",color=bcolors.OKGREEN)

                data /= flat[start_row:start_row + data_size[0],start_col:start_col + data_size[1], :, :, np.newaxis] #only one new axis for the scans
                    
            elif flat_states == 4:

                printc("Dividing by 4 flats, one for each pol state",color=bcolors.OKGREEN)

                tmp = np.mean(flat,axis=-1) #avg over wavelength

                data /= tmp[start_row:start_row + data_size[0],start_col:start_col + data_size[1], np.newaxis, np.newaxis]
jonas's avatar
jonas committed
557
        
Jonas Sinjan's avatar
Jonas Sinjan committed
558
559
560
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
            printc(f"------------- Flat Field correction time: {np.round(time.time() - start_time,3)} seconds ",bcolors.OKGREEN)
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
jonas's avatar
jonas committed
561
562
563
564

        except: 
          printc("ERROR, Unable to apply flat fields",color=bcolors.FAIL)

565
566
567
568
    else:
        print(" ")
        printc('-->>>>>>> No flat field correction mode',color=bcolors.WARNING)

jonas's avatar
jonas committed
569

jonas's avatar
jonas committed
570
571
572
573
574
    #-----------------
    # PREFILTER CORRECTION  
    #-----------------

    if prefilter_f is not None:
jonas's avatar
jonas committed
575
        print(" ")
jonas's avatar
jonas committed
576
577
578
579
580
581
582
583
584
585
586
587
        printc('-->>>>>>> Prefilter Correction ',color=bcolors.OKGREEN)

        start_time = time.time()

        prefilter_voltages = [-1300.00,-1234.53,-1169.06,-1103.59,-1038.12,-972.644,-907.173,-841.702,-776.231,-710.760,-645.289,
                            -579.818,-514.347,-448.876,-383.404,-317.933,-252.462,-186.991,-121.520,-56.0490,9.42212,74.8932,
                            140.364,205.835,271.307, 336.778,402.249,467.720,533.191,598.662,664.133,729.604,795.075,860.547,
                            926.018,991.489,1056.96,1122.43,1187.90,1253.37, 1318.84,1384.32,1449.79,1515.26,1580.73,1646.20,
                            1711.67,1777.14,1842.61]

        prefilter, _ = load_fits(prefilter_f)

jonas's avatar
jonas committed
588
        #prefilter = prefilter[:,652:1419,613:1380] #crop the helioseismology data
jonas's avatar
jonas committed
589

jonas's avatar
jonas committed
590
591
592
        def get_v1_index1(x):
            index1, v1 = min(enumerate([abs(i) for i in x]), key=itemgetter(1))
            return  v1, index1
jonas's avatar
jonas committed
593

jonas's avatar
jonas committed
594
        for scan in range(data_shape[-1]):
jonas's avatar
jonas committed
595

jonas's avatar
jonas committed
596
            voltage_list = voltagesData_arr[scan]
jonas's avatar
jonas committed
597
598
            
            for wv in range(6):
jonas's avatar
jonas committed
599

jonas's avatar
jonas committed
600
                v = voltage_list[wv]
jonas's avatar
jonas committed
601

jonas's avatar
jonas committed
602
603
604
605
606
607
608
609
610
611
612
613
                vdif = [v - pf for pf in prefilter_voltages]
                
                v1, index1 = get_v1_index1(vdif)
                
                if vdif[index1] >= 0:
                    v2 = vdif[index1 + 1]
                    index2 = index1 + 1
                    
                else:
                    v2 = vdif[index1-1]
                    index2 = index1 - 1
                    
jonas's avatar
jonas committed
614
                imprefilter = (prefilter[:,:, index1]*v1 + prefilter[:,:, index2]*v2)/(v1+v2)
jonas's avatar
jonas committed
615

jonas's avatar
jonas committed
616
                data[:,:,:,wv,scan] /= imprefilter[...,np.newaxis]
jonas's avatar
jonas committed
617
618
619
620
621
622


        printc('--------------------------------------------------------------',bcolors.OKGREEN)
        printc(f"------------- Prefilter correction time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
        printc('--------------------------------------------------------------',bcolors.OKGREEN)

623
624
625
    else:
        print(" ")
        printc('-->>>>>>> No prefilter mode',color=bcolors.WARNING)
jonas's avatar
jonas committed
626

627
628
629
630
    #-----------------
    # FIELD STOP 
    #-----------------

jonas's avatar
jonas committed
631
    if fs_c:
jonas's avatar
jonas committed
632
    
633
634
        print(" ")
        printc("-->>>>>>> Applying field stop",color=bcolors.OKGREEN)
635

636
637
638
        start_time = time.time()
        
        field_stop,_ = load_fits('./field_stop/HRT_field_stop.fits')
639

640
        field_stop = np.where(field_stop > 0,1,0)
641

642
643
644
645
646
647
648
649
650
        data *= field_stop[start_row:start_row + data_size[0],start_col:start_col + data_size[1],np.newaxis, np.newaxis, np.newaxis]

        printc('--------------------------------------------------------------',bcolors.OKGREEN)
        printc(f"------------- Field stop time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
        printc('--------------------------------------------------------------',bcolors.OKGREEN)

    else:
        print(" ")
        printc('-->>>>>>> No field stop mode',color=bcolors.WARNING)
651

jonas's avatar
jonas committed
652

jonas's avatar
jonas committed
653
654
655
656
    #-----------------
    # GHOST CORRECTION  
    #-----------------

657
658
    # if correct_ghost:
    #     printc('-->>>>>>> Correcting ghost image ',color=bcolors.OKGREEN)
jonas's avatar
jonas committed
659
660
661
662
663
664

  
    #-----------------
    # APPLY DEMODULATION 
    #-----------------

665
666
    if demod:

Jonas Sinjan's avatar
Jonas Sinjan committed
667
668
        print(" ")
        printc('-->>>>>>> Demodulating data         ',color=bcolors.OKGREEN)
jonas's avatar
jonas committed
669

jonas's avatar
jonas committed
670
671
        start_time = time.time()

jonas's avatar
jonas committed
672
        data,_ = demod_hrt(data, pmp_temp)
Jonas Sinjan's avatar
Jonas Sinjan committed
673
674
675
676
        
        printc('--------------------------------------------------------------',bcolors.OKGREEN)
        printc(f"------------- Demodulation time: {np.round(time.time() - start_time,3)} seconds ",bcolors.OKGREEN)
        printc('--------------------------------------------------------------',bcolors.OKGREEN)
jonas's avatar
jonas committed
677

678
679
680
    else:
        print(" ")
        printc('-->>>>>>> No demod mode',color=bcolors.WARNING)
jonas's avatar
jonas committed
681
682


jonas's avatar
jonas committed
683
684
685
686
    #-----------------
    # APPLY NORMALIZATION 
    #-----------------

687
    if norm_stokes:
Jonas Sinjan's avatar
Jonas Sinjan committed
688
689
690
691
        
        print(" ")
        printc('-->>>>>>> Normalising Stokes to Quiet Sun',color=bcolors.OKGREEN)
        
jonas's avatar
jonas committed
692
693
        start_time = time.time()

Jonas Sinjan's avatar
Jonas Sinjan committed
694
695
        for scan in range(data_shape[-1]):
            
jonas.sinjan's avatar
jonas.sinjan committed
696
            I_c = np.mean(data[512:1536,512:1536,0,cpos_arr[0],int(scan)]) #mean of central 1k x 1k of continuum stokes I
Jonas Sinjan's avatar
Jonas Sinjan committed
697
698
699
700
701
            data[:,:,:,:,scan] = data[:,:,:,:,scan]/I_c
       
        printc('--------------------------------------------------------------',bcolors.OKGREEN)
        printc(f"------------- Stokes Normalising time: {np.round(time.time() - start_time,3)} seconds ",bcolors.OKGREEN)
        printc('--------------------------------------------------------------',bcolors.OKGREEN)
jonas's avatar
jonas committed
702

703
704
705
    else:
        print(" ")
        printc('-->>>>>>> No normalising Stokes mode',color=bcolors.WARNING)
jonas's avatar
jonas committed
706
707


jonas's avatar
jonas committed
708
709
710
    #-----------------
    # CROSS-TALK CALCULATION 
    #-----------------
jonas's avatar
jonas committed
711

712
    if ItoQUV:
Jonas Sinjan's avatar
Jonas Sinjan committed
713
714
        
        print(" ")
jonas's avatar
jonas committed
715
        printc('-->>>>>>> Cross-talk correction I to Q,U,V ',color=bcolors.OKGREEN)
jonas's avatar
jonas committed
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773

        start_time = time.time()

        before_ctalk_data = np.copy(data)

        num_of_scans = data_shape[-1]

        try:
            assert ctalk_params.shape == (2,3)
        except AssertionError:
            print("ctalk_params is not in the required (2,3) shape, please reconcile")

        ctalk_params = np.repeat(ctalk_params[:,:,np.newaxis], num_of_scans, axis = 2)

        cont_stokes = np.mean(data[512:1536,512:1536,0,cpos_arr[0],:], axis = (0,1))

        for i in range(6):
            stokes_i_wv_avg = np.mean(data[512:1536,512:1536,0,i,:], axis = (0,1))

            if norm_stokes:
                #if normed, applies normalised offset to normed stokes

                tmp_param = ctalk_params*np.divide(stokes_i_wv_avg,cont_stokes)

                q_slope = tmp_param[0,0,:]
                u_slope = tmp_param[0,1,:]
                v_slope = tmp_param[0,2,:]

                q_int = tmp_param[1,0,:]
                u_int = tmp_param[1,1,:]
                v_int = tmp_param[1,2,:]

                data[:,:,1,i,:] = before_ctalk_data[:,:,1,i,:] - before_ctalk_data[:,:,0,i,:]*q_slope - q_int

                data[:,:,2,i,:] = before_ctalk_data[:,:,2,i,:] - before_ctalk_data[:,:,0,i,:]*u_slope - u_int

                data[:,:,3,i,:] = before_ctalk_data[:,:,3,i,:] - before_ctalk_data[:,:,0,i,:]*v_slope - v_int

            else:
                #if not normed, applies raw offset cross talk correction to raw stokes counts

                tmp_param = ctalk_params[0,:,:]*np.divide(stokes_i_wv_avg,cont_stokes)

                q_slope = tmp_param[0,:]
                u_slope = tmp_param[1,:]
                v_slope = tmp_param[2,:]

                q_int = ctalk_params[1,0,:]
                u_int = ctalk_params[1,1,:]
                v_int = ctalk_params[1,2,:]

                data[:,:,1,i,:] = before_ctalk_data[:,:,1,i,:] - before_ctalk_data[:,:,0,i,:]*q_slope - q_int*stokes_i_wv_avg 

                data[:,:,2,i,:] = before_ctalk_data[:,:,2,i,:] - before_ctalk_data[:,:,0,i,:]*u_slope - u_int*stokes_i_wv_avg 

                data[:,:,3,i,:] = before_ctalk_data[:,:,3,i,:] - before_ctalk_data[:,:,0,i,:]*v_slope - v_int*stokes_i_wv_avg 


jonas's avatar
jonas committed
774
775
776
        printc('--------------------------------------------------------------',bcolors.OKGREEN)
        printc(f"------------- I -> Q,U,V cross talk correction time: {np.round(time.time() - start_time,3)} seconds ",bcolors.OKGREEN)
        printc('--------------------------------------------------------------',bcolors.OKGREEN)
jonas's avatar
jonas committed
777
778
779
        
        data *= field_stop[start_row:start_row + data_size[0],start_col:start_col + data_size[1], np.newaxis, np.newaxis, np.newaxis]

780
781
782
783
    else:
        print(" ")
        printc('-->>>>>>> No ItoQUV mode',color=bcolors.WARNING)

jonas's avatar
jonas committed
784
785
786
787
788
789
790
791

    #-----------------
    #CHECK FOR INFs
    #-----------------

    data[np.isinf(data)] = 0
    data[np.isnan(data)] = 0

jonas.sinjan's avatar
jonas.sinjan committed
792
    
793
    if out_demod_file:
794
795
796
797
798

        #check if the output directory exists, if not, create it
        if not os.path.exists(out_dir): 
            print(f"{out_dir} does not exist, creating it now")
            os.makedirs(out_dir)
799
800
        
        if isinstance(data_f, list):
Jonas Sinjan's avatar
Jonas Sinjan committed
801
            print(" ")
jonas's avatar
jonas committed
802
            printc('Saving demodulated data to one _reduced.fits file per scan')
803

804
805
806
807
808
809
            #checking if the science scans have the same DID - this would cause an issue for naming the output demod files
            scan_name_list = [str(scan.split('.fits')[0][-10:]) for scan in data_f]

            seen = set()
            uniq_scan_DIDs = [x for x in scan_name_list if x in seen or seen.add(x)] #creates list of unique DIDs from the list

jonas's avatar
jonas committed
810
811
            #print(uniq_scan_DIDs)
            #print(scan_name_list)
jonas's avatar
jonas committed
812
            if uniq_scan_DIDs == []:
813
814
815
816
817
818
819
                print("The scan's DIDs are all unique")

            else:

                for x in uniq_scan_DIDs:
                    number = scan_name_list.count(x)
                    if number > 1: #if more than one
820
                        print(f"The DID: {x} is repeated {number} times")
821
822
823
824
825
826
827
828
829
                        i = 1
                        for index, name in enumerate(scan_name_list):
                            if name == x:
                                scan_name_list[index] = name + f"_{i}" #add _1, _2, etc to the file name, so that when written to output file not overwriting
                                i += 1

                print("The New DID list is: ", scan_name_list)
                        

Jonas Sinjan's avatar
Jonas Sinjan committed
830
            for count, scan in enumerate(data_f):
jonas's avatar
jonas committed
831

832
                with fits.open(scan) as hdu_list:
jonas's avatar
jonas committed
833

Jonas Sinjan's avatar
Jonas Sinjan committed
834
                    hdu_list[0].data = data[:,:,:,:,count]
835
                    hdu_list.writeto(out_dir + scan_name_list[count] + '_reduced.fits', overwrite=True)
836
837

        if isinstance(data_f, str):
Jonas Sinjan's avatar
Jonas Sinjan committed
838
            print(" ")
jonas's avatar
jonas committed
839
            printc('Saving demodulated data to a _reduced.fits file')
840
841
842
843

            with fits.open(data_f) as hdu_list:

                hdu_list[0].data = data
jonas.sinjan's avatar
jonas.sinjan committed
844
                hdu_list.writeto(out_dir + str(data_f.split('.fits')[0][-10:]) + '_reduced.fits', overwrite=True)
jonas's avatar
jonas committed
845

846
847
848
    else:
        print(" ")
        printc('-->>>>>>> No output demod file mode',color=bcolors.WARNING)
jonas's avatar
reorg    
jonas committed
849

jonas's avatar
jonas committed
850
851
852
853
854
    #-----------------
    # INVERSION OF DATA WITH CMILOS
    #-----------------

    if rte == 'RTE' or rte == 'CE' or rte == 'CE+RTE':
jonas's avatar
jonas committed
855

jonas's avatar
jonas committed
856
        print(" ")
jonas's avatar
jonas committed
857
        printc('-->>>>>>> RUNNING CMILOS ',color=bcolors.OKGREEN)
858
        
jonas's avatar
jonas committed
859
        try:
860
861
862
863
            CMILOS_LOC = os.path.realpath(__file__)

            CMILOS_LOC = CMILOS_LOC[:-11] + 'cmilos/' #11 as hrt_pipe.py is 11 characters

jonas's avatar
jonas committed
864
865
            if os.path.isfile(CMILOS_LOC+'milos'):
                printc("Cmilos executable located at:", CMILOS_LOC,color=bcolors.WARNING)
866

jonas's avatar
jonas committed
867
868
            else:
                raise ValueError('Cannot find cmilos:', CMILOS_LOC)
869

jonas's avatar
jonas committed
870
871
872
873
874
875
        except ValueError as err:
            printc(err.args[0],color=bcolors.FAIL)
            printc(err.args[1],color=bcolors.FAIL)
            return        

        wavelength = 6173.3356
jonas's avatar
jonas committed
876

jonas's avatar
jonas committed
877
        for scan in range(int(data_shape[-1])):
jonas's avatar
jonas committed
878

879
880
            start_time = time.time()

jonas's avatar
jonas committed
881
            if isinstance(data_f, str):
jonas's avatar
jonas committed
882

jonas's avatar
jonas committed
883
                file_path = data_f
jonas's avatar
jonas committed
884

jonas's avatar
jonas committed
885
            elif isinstance(data_f, list):
jonas's avatar
jonas committed
886

jonas's avatar
jonas committed
887
                file_path = data_f[scan]
jonas's avatar
jonas committed
888
                wave_axis = wve_axis_arr[scan]
jonas's avatar
jonas committed
889

jonas's avatar
jonas committed
890
891
            #must invert each scan independently, as cmilos only takes in one dataset at a time

jonas's avatar
jonas committed
892
             #get wave_axis from the header information of the science scans
jonas's avatar
jonas committed
893
894
895
896

            shift_w =  wave_axis[3] - wavelength
            wave_axis = wave_axis - shift_w

jonas's avatar
jonas committed
897
898
899
900
            print('It is assumed the wavelength array is given by the header')
            #print(wave_axis,color = bcolors.WARNING)
            print("Wave axis is: ", (wave_axis - wavelength)*1000.)
            print('Saving data into dummy_in.txt for RTE input')
jonas's avatar
jonas committed
901
902
903

            sdata = data[:,:,:,:,scan]
            y,x,p,l = sdata.shape
jonas's avatar
jonas committed
904
            #print(y,x,p,l)
jonas's avatar
jonas committed
905
906
907
908
909
910
911
912
913

            filename = 'dummy_in.txt'
            with open(filename,"w") as f:
                for i in range(x):
                    for j in range(y):
                        for k in range(l):
                            f.write('%e %e %e %e %e \n' % (wave_axis[k],sdata[j,i,0,k],sdata[j,i,1,k],sdata[j,i,2,k],sdata[j,i,3,k])) #wv, I, Q, U, V
            del sdata

jonas's avatar
jonas committed
914
            printc(f'  ---- >>>>> Inverting data scan number: {scan} .... ',color=bcolors.OKGREEN)
jonas's avatar
jonas committed
915
916
917

            cmd = CMILOS_LOC+"./milos"
            cmd = fix_path(cmd)
jonas's avatar
debug    
jonas committed
918

jonas's avatar
jonas committed
919
920
921
922
923
924
925
            if rte == 'RTE':
                rte_on = subprocess.call(cmd+" 6 15 0 0 dummy_in.txt  >  dummy_out.txt",shell=True)
            if rte == 'CE':
                rte_on = subprocess.call(cmd+" 6 15 2 0 dummy_in.txt  >  dummy_out.txt",shell=True)
            if rte == 'CE+RTE':
                rte_on = subprocess.call(cmd+" 6 15 1 0 dummy_in.txt  >  dummy_out.txt",shell=True)

jonas's avatar
jonas committed
926
            #print(rte_on)
jonas's avatar
debug    
jonas committed
927

jonas's avatar
jonas committed
928
929
            printc('  ---- >>>>> Reading results.... ',color=bcolors.OKGREEN)
            del_dummy = subprocess.call("rm dummy_in.txt",shell=True)
jonas's avatar
jonas committed
930
            #print(del_dummy)
jonas's avatar
jonas committed
931
932
933

            res = np.loadtxt('dummy_out.txt')
            npixels = res.shape[0]/12.
jonas's avatar
jonas committed
934
935
            #print(npixels)
            #print(npixels/x)
jonas's avatar
jonas committed
936
937
938
939
940
941
942
943
944
945
            result = np.zeros((12,y*x)).astype(float)
            rte_invs = np.zeros((12,y,x)).astype(float)
            for i in range(y*x):
                result[:,i] = res[i*12:(i+1)*12]
            result = result.reshape(12,y,x)
            result = np.einsum('ijk->ikj', result)
            rte_invs = result
            del result
            rte_invs_noth = np.copy(rte_invs)

jonas's avatar
jonas committed
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
            """
            From 0 to 11
            Counter (PX Id)
            Iterations
            Strength
            Inclination
            Azimuth
            Eta0 parameter
            Doppler width
            Damping
            Los velocity
            Constant source function
            Slope source function
            Minimum chisqr value
            """

            noise_in_V =  np.mean(data[:,:,3,cpos_arr[0],:])
jonas's avatar
jonas committed
963
            low_values_flags = np.max(np.abs(data[:,:,3,:,scan]),axis=-1) < noise_in_V  # Where values are low
jonas's avatar
jonas committed
964
965
966
967
            
            rte_invs[2,low_values_flags] = 0
            rte_invs[3,low_values_flags] = 0
            rte_invs[4,low_values_flags] = 0
jonas's avatar
jonas committed
968

jonas's avatar
jonas committed
969
            #np.savez_compressed(out_dir+'_RTE', rte_invs=rte_invs, rte_invs_noth=rte_invs_noth)
jonas's avatar
jonas committed
970
971
            
            del_dummy = subprocess.call("rm dummy_out.txt",shell=True)
jonas's avatar
jonas committed
972
            #print(del_dummy)
jonas's avatar
jonas committed
973

jonas's avatar
jonas committed
974
975
976
977
978
979
980
981
982
            rte_data_products = np.zeros((6,rte_invs_noth.shape[1],rte_invs_noth.shape[1]))

            rte_data_products[0,:,:] = rte_invs_noth[9,:,:] + rte_invs_noth[10,:,:] #continuum
            rte_data_products[1,:,:] = rte_invs_noth[2,:,:] #b mag strength
            rte_data_products[2,:,:] = rte_invs_noth[3,:,:] #inclination
            rte_data_products[3,:,:] = rte_invs_noth[4,:,:] #azimuth
            rte_data_products[4,:,:] = rte_invs_noth[8,:,:] #vlos
            rte_data_products[5,:,:] = rte_invs_noth[2,:,:]*np.cos(rte_invs_noth[3,:,:]*np.pi/180.) #blos

jonas's avatar
jonas committed
983
984
            rte_data_products *= field_stop[np.newaxis,start_row:start_row + data_size[0],start_col:start_col + data_size[1]] #field stop, set outside to 0

jonas's avatar
jonas committed
985
            with fits.open(file_path) as hdu_list:
jonas's avatar
jonas committed
986
                hdu_list[0].data = rte_data_products
jonas's avatar
jonas committed
987
                hdu_list.writeto(out_dir+str(file_path.split('.fits')[0][-10:])+'_rte_data_products.fits', overwrite=True)
jonas's avatar
jonas committed
988

jonas's avatar
jonas committed
989
            with fits.open(file_path) as hdu_list:
jonas's avatar
jonas committed
990
                hdu_list[0].data = rte_data_products[5,:,:]
jonas's avatar
jonas committed
991
                hdu_list.writeto(out_dir+str(file_path.split('.fits')[0][-10:])+'_blos_rte.fits', overwrite=True)
jonas's avatar
jonas committed
992

jonas's avatar
jonas committed
993
            with fits.open(file_path) as hdu_list:
jonas's avatar
jonas committed
994
                hdu_list[0].data = rte_data_products[4,:,:]
jonas's avatar
jonas committed
995
                hdu_list.writeto(out_dir+str(file_path.split('.fits')[0][-10:])+'_vlos_rte.fits', overwrite=True)
jonas's avatar
jonas committed
996

jonas's avatar
jonas committed
997
            with fits.open(file_path) as hdu_list:
jonas's avatar
jonas committed
998
                hdu_list[0].data = rte_data_products[0,:,:]
jonas's avatar
jonas committed
999
                hdu_list.writeto(out_dir+str(file_path.split('.fits')[0][-10:])+'_Icont_rte.fits', overwrite=True)
jonas's avatar
jonas committed
1000

jonas's avatar
jonas committed
1001
1002
1003
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
            printc(f"------------- CMILOS RTE Run Time: {np.round(time.time() - start_time,3)} seconds ",bcolors.OKGREEN)
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
jonas's avatar
jonas committed
1004

1005
1006
1007
1008
    else:
        print(" ")
        printc('-->>>>>>> No RTE Inversion mode',color=bcolors.WARNING)

Jonas Sinjan's avatar
Jonas Sinjan committed
1009
1010
1011
1012
    print(" ")
    printc('--------------------------------------------------------------',color=bcolors.OKGREEN)
    printc(f'------------ Reduction Complete: {np.round(time.time() - overall_time,3)} seconds',color=bcolors.OKGREEN)
    printc('--------------------------------------------------------------',color=bcolors.OKGREEN)
jonas's avatar
jonas committed
1013
1014


jonas's avatar
jonas committed
1015
    return data