hrt_pipe.py 23.5 KB
Newer Older
jonas's avatar
jonas committed
1
2
3
4
5
6
7
import numpy as np 
import os.path
from astropy.io import fits
import random, statistics
import subprocess
from scipy.ndimage import gaussian_filter
import matplotlib.pyplot as plt
8
import time
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
14
15
16
17
18
19
20
21
22
23
24


def get_data(path):
    """load science data from path"""
    try:
        data, header = load_fits(path)
      
        data /=  256. #conversion from 24.8bit to 32bit

        accu = header['ACCACCUM']*header['ACCROWIT']*header['ACCCOLIT'] #getting the number of accu from header

        data /= accu

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

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

Jonas Sinjan's avatar
Jonas Sinjan committed
31
   
jonas's avatar
jonas committed
32

Jonas Sinjan's avatar
Jonas Sinjan committed
33
def demod_hrt(data,pmp_temp):
jonas's avatar
jonas committed
34
    '''
Jonas Sinjan's avatar
Jonas Sinjan committed
35
    Use constant demodulation matrices to demodulate data
jonas's avatar
jonas committed
36
    '''
Jonas Sinjan's avatar
Jonas Sinjan committed
37
38
39
40
41
42
 
    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
43
        
Jonas Sinjan's avatar
Jonas Sinjan committed
44
45
46
47
48
    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
49
    
jonas's avatar
jonas committed
50
    else:
Jonas Sinjan's avatar
Jonas Sinjan committed
51
        printc("Demodulation Matrix for PMP TEMP of {pmp_temp} deg is not available", color = bcolors.FAIL)
jonas's avatar
jonas committed
52

Jonas Sinjan's avatar
Jonas Sinjan committed
53
    printc(f'Using a constant demodulation matrix for a PMP TEMP of {pmp_temp} deg',color = bcolors.OKGREEN)
jonas's avatar
jonas committed
54
    
Jonas Sinjan's avatar
Jonas Sinjan committed
55
56
57
    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
58
59
60
61

    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
62

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


jonas.sinjan's avatar
jonas.sinjan committed
74
def phihrt_pipe(data_f,dark_f,flat_f,norm_f = True, clean_f = False, flat_states = 24, 
Jonas Sinjan's avatar
Jonas Sinjan committed
75
                pmp_temp = '50',flat_c = True,dark_c = True, demod = True, continuum_wavelength = 0, norm_stokes = True, 
76
                out_dir = './',  out_demod_file = False,  correct_ghost = False, 
jonas's avatar
jonas committed
77
                ItoQUV = False, rte = False):
jonas's avatar
jonas committed
78
79
80
81
82

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

    Parameters
    ----------
        Input:
    data_f : list or string
        list containing paths to fits files of the raw HRT data OR string of path to one file  
    dark_f : string
        Fits file of a dark file (ONLY ONE FILE)
    flat_f : string
        Fits file of a HRT flatfield (ONLY ONE FILE)

    ** Options:
jonas's avatar
jonas committed
108
109
110
    norm_f: bool, DEFAULT: True
        to normalise the flat fields before applying
    clean_f: bool, DEFAULT: False
jonas.sinjan's avatar
jonas.sinjan committed
111
        clean the flat field with unsharp masking   
jonas's avatar
jonas committed
112
113
114
115
116
117
118
119
120
121
    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)
    pmp_temp: str, DEFAULT: '50'
        temperature of the PMP, to select the correct demodulation matrix
    flat_c: bool, DEFAULT: True
        apply flat field correction
    dark_c: bool, DEFAULT: True
        apply dark field correction
    demod: bool, DEFAULT: True
        apply demodulate to the stokes
Jonas Sinjan's avatar
Jonas Sinjan committed
122
123
    continuum_wavelength: int, DEFAULT: 0
        index of the continuum wavelength in the science data, assumes that the flat field has the same
jonas's avatar
jonas committed
124
125
126
127
128
129
130
131
132
133
134
135
136
137
    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
    rte: bool, DEFAULT: False 
        invert using cmilos
    out_rte_file: bool, DEFAULT: False
        output of rte result to fits file
jonas's avatar
jonas committed
138
139
140
    
    Returns
    -------
jonas's avatar
jonas committed
141
142
    data: nump array
        stokes vector 
jonas's avatar
jonas committed
143
144
145
146
147
148
149
150
151
152
153

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

    '''

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

Jonas Sinjan's avatar
Jonas Sinjan committed
154
    overall_time = time.time()
jonas's avatar
jonas committed
155
156
157
    #-----------------
    # READ DATA
    #-----------------
Jonas Sinjan's avatar
Jonas Sinjan committed
158
159
    
    print(" ")
Jonas Sinjan's avatar
Jonas Sinjan committed
160
    printc('-->>>>>>> Reading Data              ',color=bcolors.OKGREEN) 
jonas's avatar
jonas committed
161

162
    start_time = time.time()
jonas's avatar
jonas committed
163
164

    if isinstance(data_f, list):
165
        #if the data_f contains several scans
Jonas Sinjan's avatar
Jonas Sinjan committed
166
        printc(f'Input contains {len(data_f)} scan(s)',color=bcolors.OKGREEN)
167
168
169
170
171
172
173
174
175
176
177
        
        number_of_scans = len(data_f)

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

        for scan in range(number_of_scans):
            data_arr[scan], hdr_arr[scan] = get_data(data_f[scan])

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

Jonas Sinjan's avatar
Jonas Sinjan committed
178
                print(f"This scan: {data_f[scan]} has a bits per pixel of: 16 \nPerforming the extra scaling")
179
180
181
182
183

                data_arr[scan] *= 81920/127 #conversion factor if 16 bits

        #test if the scans have different sizes
        first_shape = data_arr[scan].shape
jonas's avatar
jonas committed
184

185
186
        result = all(element.shape == first_shape for element in data_arr)
        if (result):
Jonas Sinjan's avatar
Jonas Sinjan committed
187
            print("All the scans have the same dimension")
jonas's avatar
jonas committed
188

189
190
191
192
193
194
195
        else:
            print("The scans have different dimensions! \n Ending process")

            exit()

        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
196

197
        print(f"Data shape is {data.shape}")
jonas's avatar
jonas committed
198
199

    elif isinstance(data_f, str):
200
201
202
203
204
205
206
207
208
209
210
        #case when data f is just one file
        data, header = get_data(data_f)
        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

jonas's avatar
jonas committed
211
  
212
213
        print(f"Data shape is {data.shape}")

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

Jonas Sinjan's avatar
Jonas Sinjan committed
218
    
jonas's avatar
jonas committed
219
      
220
221
222
223
224
225
226
227
228
229
    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
230
    print("Data reshaped to: ", data.shape)
231

Jonas Sinjan's avatar
Jonas Sinjan committed
232
    diff = 2048-data_size[0] #handling 0/2 errors
233
234
235
236
237
238
239
240
    
    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
241
242
243
244
        
    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
245
246
247
248

    #-----------------
    # TODO: Could check data dimensions? As an extra fail safe before progressing?
    #-----------------
jonas's avatar
jonas committed
249
250
251
252
253
254
255
    
    
    #-----------------
    # READ FLAT FIELDS
    #-----------------

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

259
260
        start_time = time.time()
    
jonas's avatar
jonas committed
261
        try:
Jonas Sinjan's avatar
Jonas Sinjan committed
262
            flat, header_flat = get_data(flat_f)
263
264
265

            print(f"Flat field shape is {flat.shape}")
            
Jonas Sinjan's avatar
Jonas Sinjan committed
266
            if header_flat['BITPIX'] == 16:
267
268
269
270
271
272
    
                print("Number of bits per pixel is: 16")

                flat *= 614400/128

            flat = np.moveaxis(flat, 0,-1) #so that it is [y,x,24]
Jonas Sinjan's avatar
Jonas Sinjan committed
273
274
275
            flat = flat.reshape(data_size[0],data_size[1],6,4) #separate 24 images, into 6 wavelengths, with each 4 pol states
            flat = np.moveaxis(flat, 2,-1)
            
jonas.sinjan's avatar
jonas.sinjan committed
276
            print(flat.shape)
Jonas Sinjan's avatar
Jonas Sinjan committed
277
278
279
280
281
282
283
            
            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)
                
284

Jonas Sinjan's avatar
Jonas Sinjan committed
285
286
287
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
            printc(f"------------ Load flats time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
288

jonas's avatar
jonas committed
289
290
291
292
293
        except Exception:
            printc("ERROR, Unable to open flats file: {}",flat_f,color=bcolors.FAIL)


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

jonas's avatar
jonas committed
297

jonas.sinjan's avatar
jonas.sinjan committed
298
299
300
301
302
303
304
305
306
307
308
    #-----------------
    # OPTIONAL Unsharp Masking clean the flat field stokes V images
    #-----------------

    if clean_f:
        print(" ")
        printc('-->>>>>>> Cleaning flats with Unsharp Masking',color=bcolors.OKGREEN)
        
        #call the clean function (not yet built)


jonas's avatar
jonas committed
309
310
311
    #-----------------
    # READ AND CORRECT DARK FIELD
    #-----------------
312

jonas's avatar
jonas committed
313
    if dark_c:
Jonas Sinjan's avatar
Jonas Sinjan committed
314
        print(" ")
jonas's avatar
jonas committed
315
        printc('-->>>>>>> Reading Darks                   ',color=bcolors.OKGREEN)
316
317
318

        start_time = time.time()

jonas's avatar
jonas committed
319
        try:
jonas's avatar
jonas committed
320
321
322
323
324
325
326
327
328
329
330
331
332
333
            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:,:]
                
334
                except Exception:
jonas's avatar
jonas committed
335
336
                    printc("ERROR, Unable to correct shape of dark field data: {}",dark_f,color=bcolors.FAIL)

Jonas Sinjan's avatar
Jonas Sinjan committed
337
338
339
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
            printc(f"------------ Load darks time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
340

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

jonas's avatar
jonas committed
344

345
346
347
        #-----------------
        # APPLY DARK CORRECTION 
        #-----------------    
Jonas Sinjan's avatar
Jonas Sinjan committed
348
        print(" ")
jonas.sinjan's avatar
jonas.sinjan committed
349
        print("'-->>>>>>> Subtracting dark field")
Jonas Sinjan's avatar
Jonas Sinjan committed
350
        
351
352
        start_time = time.time()

Jonas Sinjan's avatar
Jonas Sinjan committed
353
        flat -= dark[..., np.newaxis, np.newaxis]
354
        
Jonas Sinjan's avatar
Jonas Sinjan committed
355
        data -= dark[start_row:start_row + data_size[0],start_col:start_col + data_size[1], np.newaxis, np.newaxis, np.newaxis] 
356

Jonas Sinjan's avatar
Jonas Sinjan committed
357
358
359
        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
360
361


362
363
364
365
366
    #-----------------
    # NORM FLAT FIELDS
    #-----------------

    if norm_f and flat_c:
Jonas Sinjan's avatar
Jonas Sinjan committed
367
368
        
        print(" ")
369
370
371
372
373
        printc('-->>>>>>> Normalising Flats',color=bcolors.OKGREEN)

        start_time = time.time()

        try:
Jonas Sinjan's avatar
Jonas Sinjan committed
374
375
            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, ...]
376

Jonas Sinjan's avatar
Jonas Sinjan committed
377
378
379
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
            printc(f"------------- Normalising flat time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
380
381
382
383

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

jonas's avatar
jonas committed
384

jonas's avatar
jonas committed
385
386
387
388
389
    #-----------------
    # APPLY FLAT CORRECTION 
    #-----------------

    if flat_c:
Jonas Sinjan's avatar
Jonas Sinjan committed
390
        print(" ")
jonas's avatar
jonas committed
391
        printc('-->>>>>>> Correcting Flatfield',color=bcolors.OKGREEN)
392
393
394

        start_time = time.time()

jonas's avatar
jonas committed
395
        try:
396
397

            if flat_states == 6:
jonas's avatar
jonas committed
398
        
399
400
401
402
                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
403
                data /= tmp[start_row:start_row + data_size[0],start_col:start_col + data_size[1], np.newaxis, np.newaxis]
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418


            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
419
        
Jonas Sinjan's avatar
Jonas Sinjan committed
420
421
422
            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
423
424
425
426

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

jonas's avatar
jonas committed
427

428
429
430
    #-----------------
    # FIELD STOP 
    #-----------------
Jonas Sinjan's avatar
Jonas Sinjan committed
431
432
433
    
    print(" ")
    printc("-->>>>>>>  Applying field stop",color=bcolors.OKGREEN)
434
435

    start_time = time.time()
jonas's avatar
jonas committed
436
    
Jonas Sinjan's avatar
Jonas Sinjan committed
437
    field_stop,_ = load_fits('./field_stop/HRT_field_stop.fits')
438
439
440
441
442

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

    data *= field_stop[start_row:start_row + data_size[0],start_col:start_col + data_size[1],np.newaxis, np.newaxis, np.newaxis]

Jonas Sinjan's avatar
Jonas Sinjan committed
443
444
445
    printc('--------------------------------------------------------------',bcolors.OKGREEN)
    printc(f"------------- Field stop time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
    printc('--------------------------------------------------------------',bcolors.OKGREEN)
446

jonas's avatar
jonas committed
447

jonas's avatar
jonas committed
448
449
450
451
    #-----------------
    # GHOST CORRECTION  
    #-----------------

452
453
    # if correct_ghost:
    #     printc('-->>>>>>> Correcting ghost image ',color=bcolors.OKGREEN)
jonas's avatar
jonas committed
454
455
456
457
458
459

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

460
461
    if demod:

Jonas Sinjan's avatar
Jonas Sinjan committed
462
463
        print(" ")
        printc('-->>>>>>> Demodulating data         ',color=bcolors.OKGREEN)
jonas's avatar
jonas committed
464

jonas's avatar
jonas committed
465
466
        start_time = time.time()

jonas.sinjan's avatar
jonas.sinjan committed
467
        data = demod_hrt(data, pmp_temp)
Jonas Sinjan's avatar
Jonas Sinjan committed
468
469
470
471
        
        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
472

jonas's avatar
jonas committed
473

jonas's avatar
jonas committed
474
475
476
477
    #-----------------
    # APPLY NORMALIZATION 
    #-----------------

478
    if norm_stokes:
Jonas Sinjan's avatar
Jonas Sinjan committed
479
480
481
482
        
        print(" ")
        printc('-->>>>>>> Normalising Stokes to Quiet Sun',color=bcolors.OKGREEN)
        
jonas's avatar
jonas committed
483
484
        start_time = time.time()

Jonas Sinjan's avatar
Jonas Sinjan committed
485
486
487
488
489
490
491
492
        for scan in range(data_shape[-1]):
            
            I_c = np.mean(data[512:1536,512:1536,0,continuum_wavelength,scan], axis = (0,1)) #mean of central 1k x 1k of continuum stokes I
            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
493
494
495
496

    #-----------------
    # CROSS-TALK CALCULATION 
    #-----------------
497
    if ItoQUV:
Jonas Sinjan's avatar
Jonas Sinjan committed
498
499
500
        
        print(" ")
        printc('-->>>>>>> Cross-talk correction from Stokes I to Stokes Q,U,V ',color=bcolors.OKGREEN)
jonas's avatar
jonas committed
501
502
503
504
505
506
507
508
509
   

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

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

jonas.sinjan's avatar
jonas.sinjan committed
510
    
511
512
513
    if out_demod_file:
        
        if isinstance(data_f, list):
Jonas Sinjan's avatar
Jonas Sinjan committed
514
515
            print(" ")
            printc('Saving demodulated data to one file per scan: ')
516

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

519
                with fits.open(scan) as hdu_list:
jonas's avatar
jonas committed
520

Jonas Sinjan's avatar
Jonas Sinjan committed
521
                    hdu_list[0].data = data[:,:,:,:,count]
jonas.sinjan's avatar
jonas.sinjan committed
522
                    hdu_list.writeto(out_dir + str(scan.split('.fits')[0][-10:]) + '_reduced.fits', overwrite=True)
523
524

        if isinstance(data_f, str):
Jonas Sinjan's avatar
Jonas Sinjan committed
525
526
            print(" ")
            printc('Saving demodulated data to one file: ')
527
528
529
530

            with fits.open(data_f) as hdu_list:

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

jonas's avatar
reorg    
jonas committed
533

jonas's avatar
jonas committed
534
535
536
537
538
    #-----------------
    # INVERSION OF DATA WITH CMILOS
    #-----------------

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

jonas's avatar
jonas committed
540
        print(" ")
jonas's avatar
jonas committed
541
542
        printc('---------------------RUNNING CMILOS --------------------------',color=bcolors.OKGREEN)

jonas's avatar
debug    
jonas committed
543
        start_time = time.time()
jonas's avatar
jonas committed
544
        try:
545
546
547
548
            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
549
550
            if os.path.isfile(CMILOS_LOC+'milos'):
                printc("Cmilos executable located at:", CMILOS_LOC,color=bcolors.WARNING)
551

jonas's avatar
jonas committed
552
553
            else:
                raise ValueError('Cannot find cmilos:', CMILOS_LOC)
554

jonas's avatar
jonas committed
555
556
557
558
559
560
        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
561

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

jonas's avatar
jonas committed
564
565
566
567
568
            if isinstance(data_f, str):
                file_path = data_f
            elif isinstance(data_f, list):
                file_path = data_f[scan]

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

jonas's avatar
jonas committed
571
            wave_axis, voltagesData, tuning_constant, cpos = fits_get_sampling(file_path,verbose = True) #get wave_axis from the header information of the science scans
jonas's avatar
jonas committed
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597

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

            printc('   It is assumed the wavelength is given by the header info ')
            printc(wave_axis,color = bcolors.WARNING)
            printc((wave_axis - wavelength)*1000.,color = bcolors.WARNING)
            printc('   saving data into dummy_in.txt for RTE input')

            sdata = data[:,:,:,:,scan]
            y,x,p,l = sdata.shape
            print(y,x,p,l)

            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

            printc('  ---- >>>>> Inverting data.... ',color=bcolors.OKGREEN)
            umbral = 3.

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

            print(cmd)

jonas's avatar
jonas committed
601
602
603
604
605
606
607
608
            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)

            print(rte_on)
jonas's avatar
debug    
jonas committed
609
610
611
612
613

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

jonas's avatar
jonas committed
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
            printc('  ---- >>>>> Reading results.... ',color=bcolors.OKGREEN)
            del_dummy = subprocess.call("rm dummy_in.txt",shell=True)
            print(del_dummy)

            res = np.loadtxt('dummy_out.txt')
            npixels = res.shape[0]/12.
            print(npixels)
            print(npixels/x)
            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)

            noise_in_V =  np.mean(data[:,:,0,3])
            low_values_flags = np.max(np.abs(data[:,:,3,:]),axis=0) < noise_in_V*umbral  # Where values are low
            
            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
638

jonas's avatar
jonas committed
639
            np.savez_compressed(out_dir+'_RTE', rte_invs=rte_invs, rte_invs_noth=rte_invs_noth)
jonas's avatar
jonas committed
640
641
642
            
            del_dummy = subprocess.call("rm dummy_out.txt",shell=True)
            print(del_dummy)
jonas's avatar
jonas committed
643

jonas's avatar
jonas committed
644
645
            b_los = rte_invs_noth[2,:,:]*np.cos(rte_invs_noth[3,:,:]*np.pi/180.)
            v_los = rte_invs_noth[8,:,:]
jonas's avatar
jonas committed
646

jonas's avatar
jonas committed
647
648
649
            
            with fits.open(file_path) as hdu_list:
                hdu_list[0].data = b_los
jonas's avatar
jonas committed
650
                hdu_list.writeto(out_dir+str(file_path.split('.fits')[0][-10:])+'_blos_rte.fits', overwrite=True)
jonas's avatar
jonas committed
651

jonas's avatar
jonas committed
652
653
            with fits.open(file_path) as hdu_list:
                hdu_list[0].data = v_los
jonas's avatar
jonas committed
654
                hdu_list.writeto(out_dir+str(file_path.split('.fits')[0][-10:])+'_vlos_rte.fits', overwrite=True)
jonas's avatar
jonas committed
655

jonas's avatar
jonas committed
656
657
            with fits.open(file_path) as hdu_list:
                hdu_list[0].data = rte_invs[9,:,:]+rte_invs[10,:,:]
jonas's avatar
jonas committed
658
                hdu_list.writeto(out_dir+str(file_path.split('.fits')[0][-10:])+'_Icont_rte.fits', overwrite=True)
jonas's avatar
jonas committed
659

jonas's avatar
jonas committed
660
661
662
663

            printc('--------------------- RTE END ----------------------------',color=bcolors.FAIL)

        
Jonas Sinjan's avatar
Jonas Sinjan committed
664
665
666
667
    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
668
669


Jonas Sinjan's avatar
Jonas Sinjan committed
670
    return data
jonas's avatar
rm end    
jonas committed
671

672