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

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


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

jonas's avatar
jonas committed
17
18
    if bit_convert_scale: #conversion from 24.8bit to 32bit
        data /=  256.
jonas's avatar
jonas committed
19

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

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

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

jonas's avatar
jonas committed
30
31
    #except Exception:
    #    printc("ERROR, Unable to open fits file: {}",path,color=bcolors.FAIL)
jonas's avatar
jonas committed
32

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
    


jonas's avatar
jonas committed
76
def phihrt_pipe(data_f, dark_f = '', flat_f = '', scale_data = True, bit_flat = 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
                demod = True, norm_stokes = True, out_dir = './',  out_demod_file = False,  out_demod_filename = None,
jonas's avatar
jonas committed
79
                ItoQUV = False, ctalk_params = None, rte = False, out_rte_filename = None, p_milos = True, config_file = True):
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 pmilos/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
115
116
    scale_data: bool, DEFAULT True
        performs the accumulation scaling + conversion for flat and science (only FALSE for commissioning data)
    bit_flat: bool, DEFAULT True
        divides the scan + flat by 256 to convert from 24.8bit to 32bits
jonas's avatar
jonas committed
117
118
119
    norm_f: bool, DEFAULT: True
        to normalise the flat fields before applying
    clean_f: bool, DEFAULT: False
jonas's avatar
jonas committed
120
121
122
        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
123
124
    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
125
126
    prefilter_f: str, DEFAULT None
        file path location to prefilter fits file, apply prefilter correction
jonas's avatar
jonas committed
127
128
129
130
    flat_c: bool, DEFAULT: True
        apply flat field correction
    dark_c: bool, DEFAULT: True
        apply dark field correction
jonas's avatar
jonas committed
131
    fs_c: bool, DEFAULT True
132
        apply HRT field stop
jonas's avatar
jonas committed
133
134
135
136
137
138
139
140
    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
141
142
    out_demod_filename: str, DEFAULT = None
        if None, takes last 10 characters of input scan filename (assumes its a DID), change if want other name
jonas's avatar
jonas committed
143
144
    ItoQUV: bool, DEFAULT: False 
        apply I -> Q,U,V correction
jonas's avatar
jonas committed
145
146
    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
147
148
    rte: str, DEFAULT: False 
        invert using cmilos, options: 'RTE' for Milne Eddington Inversion, 'CE' for Classical Estimates, 'CE+RTE' for combined
149
150
    out_rte_filename: str, DEFAULT = ''
        if '', takes last 10 characters of input scan filename (assumes its a DID), change if want other name
151
    p_milos: bool, DEFAULT = True
jonas's avatar
jonas committed
152
        if True, will execute the RTE inversion using the parallel version of the CMILOS code on 16 processors
jonas's avatar
jonas committed
153
154
    config_file: bool, DEFAULT = True
        if True, will generate config.txt file that writes the reduction process steps done
jonas's avatar
jonas committed
155
156
157
    
    Returns
    -------
jonas's avatar
jonas committed
158
159
    data: nump array
        stokes vector 
jonas's avatar
jonas committed
160
161
162
163
164
165
166
167
168
169
170

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

    '''

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

Jonas Sinjan's avatar
Jonas Sinjan committed
171
    overall_time = time.time()
jonas's avatar
jonas committed
172
    saved_args = locals()
jonas's avatar
jonas committed
173
    saved_args['ctalk_params'] = ctalk_params.tolist()
jonas's avatar
jonas committed
174
175
176
    #-----------------
    # READ DATA
    #-----------------
177

Jonas Sinjan's avatar
Jonas Sinjan committed
178
    print(" ")
Jonas Sinjan's avatar
Jonas Sinjan committed
179
    printc('-->>>>>>> Reading Data              ',color=bcolors.OKGREEN) 
jonas's avatar
jonas committed
180

181
    start_time = time.time()
jonas's avatar
jonas committed
182

183
184
185
    if isinstance(data_f, str):
        data_f = [data_f]

jonas's avatar
jonas committed
186
    if isinstance(data_f, list):
187
        #if the data_f contains several scans
Jonas Sinjan's avatar
Jonas Sinjan committed
188
        printc(f'Input contains {len(data_f)} scan(s)',color=bcolors.OKGREEN)
189
190
191
192
193
194
        
        number_of_scans = len(data_f)

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

jonas's avatar
jonas committed
195
196
        wve_axis_arr = [0]*number_of_scans
        cpos_arr = [0]*number_of_scans
jonas's avatar
jonas committed
197
198
        voltagesData_arr = [0]*number_of_scans
        tuning_constant_arr = [0]*number_of_scans
jonas's avatar
jonas committed
199

200
        for scan in range(number_of_scans):
201
            data_arr[scan], hdr_arr[scan] = get_data(data_f[scan], scaling = scale_data)
202

203
            wve_axis_arr[scan], voltagesData_arr[scan], tuning_constant_arr[scan], cpos_arr[scan] = fits_get_sampling(data_f[scan],verbose = True)
204

jonas's avatar
jonas committed
205
            if scale_data: #not for commissioning data
206

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

jonas's avatar
jonas committed
209
            if 'IMGDIRX' in hdr_arr[scan] and hdr_arr[scan]['IMGDIRX'] == 'YES':
210
                print(f"This scan has been flipped in the Y axis to conform to orientation standards. \n File: {data_f[scan]}")
jonas's avatar
jonas committed
211

jonas's avatar
jonas committed
212

jonas's avatar
jonas committed
213
        #--------
jonas's avatar
jonas committed
214
        # test if the scans have different sizes
jonas's avatar
jonas committed
215
216
        #--------

217
218
219
        first_shape = data_arr[scan].shape
        result = all(element.shape == first_shape for element in data_arr)
        if (result):
220
            print("All the scan(s) have the same dimension")
jonas's avatar
jonas committed
221

222
223
224
225
226
        else:
            print("The scans have different dimensions! \n Ending process")

            exit()

jonas's avatar
jonas committed
227
228

        #--------
jonas's avatar
jonas committed
229
        # test if the scans have different continuum wavelength_positions
jonas's avatar
jonas committed
230
231
232
233
234
        #--------

        first_cpos = cpos_arr[0]
        result = all(c_position == first_cpos for c_position in cpos_arr)
        if (result):
235
            print("All the scan(s) have the same continuum wavelength position")
jonas's avatar
jonas committed
236
237

        else:
238
239
240
            print("The scans have different continuum_wavelength postitions! Please fix \n Ending Process")

            exit()
jonas's avatar
jonas committed
241

jonas's avatar
jonas committed
242
        #--------
jonas's avatar
jonas committed
243
        # test if the scans have different pmp temperatures
jonas's avatar
jonas committed
244
245
246
247
248
        #--------

        first_pmp_temp = hdr_arr[0]['HPMPTSP1']
        result = all(hdr['HPMPTSP1'] == first_pmp_temp for hdr in hdr_arr)
        if (result):
249
            print(f"All the scan(s) have the same PMP Temperature Set Point: {first_pmp_temp}")
jonas's avatar
jonas committed
250
251
252
253
254
255
256
            pmp_temp = str(first_pmp_temp)

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

            exit()

257
258
        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
259

260
        print(f"Data shape is {data.shape}")
jonas's avatar
jonas committed
261

jonas's avatar
jonas committed
262
263
264
265
        #--------
        # test if the scans have same IMGDIRX keyword
        #--------
        if all('IMGDIRX' in hdr for hdr in hdr_arr):
jonas's avatar
jonas committed
266
            header_imgdirx_exists = True
jonas's avatar
jonas committed
267
268
269
            first_imgdirx = hdr_arr[0]['IMGDIRX']
            result = all(hdr['IMGDIRX'] == first_imgdirx for hdr in hdr_arr)
            if (result):
270
                print(f"All the scan(s) have the same IMGDIRX keyword: {first_imgdirx}")
jonas's avatar
jonas committed
271
272
273
274
275
                imgdirx_flipped = str(first_imgdirx)
            else:
                print("The scans have different IMGDIRX keywords! Please fix \n Ending Process")
                exit()
        else:
jonas's avatar
jonas committed
276
            header_imgdirx_exists = False
277
            print("Not all the scan(s) contain the 'IMGDIRX' keyword! Assuming all not flipped - Proceed with caution")
jonas's avatar
jonas committed
278

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

jonas's avatar
jonas committed
283
284
    saved_args['science_cpos'] = int(cpos_arr[0])

285
286
287
288
289
290
291
292
293
294
    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
295
    print("Data reshaped to: ", data.shape)
296

Jonas Sinjan's avatar
Jonas Sinjan committed
297
    diff = 2048-data_size[0] #handling 0/2 errors
298
299
300
301
302
303
304
305
    
    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
306
307
308
309
        
    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
310

jonas's avatar
jonas committed
311
312
313
314
315
    #-----------------
    # READ FLAT FIELDS
    #-----------------

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

319
320
        start_time = time.time()
    
jonas's avatar
jonas committed
321
322
        #try:
        flat, header_flat = get_data(flat_f, bit_convert_scale=bit_flat)
323

jonas's avatar
jonas committed
324
325
326
        print(f"Flat field shape is {flat.shape}")
        
        if header_flat['BITPIX'] == 16:
327

jonas's avatar
jonas committed
328
            print("Number of bits per pixel is: 16")
329

jonas's avatar
jonas committed
330
            flat *= 614400/128
jonas's avatar
jonas committed
331

jonas's avatar
jonas committed
332
333
334
335
336
337
338
        flat = np.moveaxis(flat, 0,-1) #so that it is [y,x,24]
        flat = flat.reshape(2048,2048,6,4) #separate 24 images, into 6 wavelengths, with each 4 pol states
        flat = np.moveaxis(flat, 2,-1)
        
        print(flat.shape)

        _, _, _, cpos_f = fits_get_sampling(flat_f,verbose = True)
jonas's avatar
jonas committed
339

jonas's avatar
jonas committed
340
        print(f"The continuum position of the flat field is at {cpos_f} index position")
jonas's avatar
jonas committed
341
342

        saved_args['flat_cpos'] = int(cpos_f)
jonas's avatar
jonas committed
343
        
jonas's avatar
jonas committed
344
345
        if cpos_f != cpos_arr[0]:
            print("The flat field continuum position is not the same as the data, trying to correct.")
346

jonas's avatar
jonas committed
347
            if cpos_f == 5 and cpos_arr[0] == 0:
jonas's avatar
jonas committed
348

jonas's avatar
jonas committed
349
                flat = np.roll(flat, 1, axis = -1)
jonas's avatar
jonas committed
350

jonas's avatar
jonas committed
351
352
353
354
355
356
357
358
            elif cpos_f == 0 and cpos_arr[0] == 5:

                flat = np.roll(flat, -1, axis = -1)

            else:
                print("Cannot reconcile the different continuum positions. \n Ending Process.")

                exit()
jonas's avatar
jonas committed
359

jonas's avatar
jonas committed
360
        flat_pmp_temp = str(header_flat['HPMPTSP1'])
jonas's avatar
jonas committed
361

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

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


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

jonas's avatar
jonas committed
376
377
378
379

    #-----------------
    # READ AND CORRECT DARK FIELD
    #-----------------
380

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

        start_time = time.time()

jonas's avatar
jonas committed
387
        try:
jonas's avatar
jonas committed
388
389
390
391
392
393
394
395
396
397
398
399
400
401
            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:,:]
                
402
                except Exception:
jonas's avatar
jonas committed
403
404
                    printc("ERROR, Unable to correct shape of dark field data: {}",dark_f,color=bcolors.FAIL)

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

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

jonas's avatar
jonas committed
412

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

Jonas Sinjan's avatar
Jonas Sinjan committed
421
        flat -= dark[..., np.newaxis, np.newaxis]
jonas's avatar
jonas committed
422

423
424
425
426
        if header_imgdirx_exists:
            if imgdirx_flipped == 'YES':
                dark_copy = np.copy(dark)
                dark_copy = dark_copy[:,::-1]
jonas's avatar
jonas committed
427

428
                data -= dark_copy[start_row:start_row + data_size[0],start_col:start_col + data_size[1], np.newaxis, np.newaxis, np.newaxis] 
429
        
jonas's avatar
jonas committed
430
431
        else:
            data -= dark[start_row:start_row + data_size[0],start_col:start_col + data_size[1], np.newaxis, np.newaxis, np.newaxis] 
432

Jonas Sinjan's avatar
Jonas Sinjan committed
433
434
435
        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
436

437
438
439
    else:
        print(" ")
        printc('-->>>>>>> No dark mode',color=bcolors.WARNING)
jonas's avatar
jonas committed
440

jonas's avatar
jonas committed
441

jonas's avatar
jonas committed
442
443
444
445
446
447
448
449
450
451
452
453
454
    #-----------------
    # 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
455
456
            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
457
458
459

        #demod the flats

jonas's avatar
jonas committed
460
        flat_demod, demodM = demod_hrt(flat, flat_pmp_temp)
jonas's avatar
jonas committed
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475

        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)

jonas's avatar
jonas committed
476
        for pol in range(3,4):
jonas's avatar
jonas committed
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495

            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)

496
497
498
    else:
        print(" ")
        printc('-->>>>>>> No clean flats mode',color=bcolors.WARNING)
jonas's avatar
jonas committed
499
500


501
502
503
504
505
    #-----------------
    # NORM FLAT FIELDS
    #-----------------

    if norm_f and flat_c:
Jonas Sinjan's avatar
Jonas Sinjan committed
506
507
        
        print(" ")
508
509
510
511
512
        printc('-->>>>>>> Normalising Flats',color=bcolors.OKGREEN)

        start_time = time.time()

        try:
Jonas Sinjan's avatar
Jonas Sinjan committed
513
514
            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, ...]
515

Jonas Sinjan's avatar
Jonas Sinjan committed
516
517
518
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
            printc(f"------------- Normalising flat time: {np.round(time.time() - start_time,3)} seconds",bcolors.OKGREEN)
            printc('--------------------------------------------------------------',bcolors.OKGREEN)
519
520
521
522

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

523
524
525
526
    else:
        print(" ")
        printc('-->>>>>>> No normalising flats mode',color=bcolors.WARNING)

jonas's avatar
jonas committed
527

jonas's avatar
jonas committed
528
529
530
531
532
    #-----------------
    # APPLY FLAT CORRECTION 
    #-----------------

    if flat_c:
Jonas Sinjan's avatar
Jonas Sinjan committed
533
        print(" ")
jonas's avatar
jonas committed
534
        printc('-->>>>>>> Correcting Flatfield',color=bcolors.OKGREEN)
535
536
537

        start_time = time.time()

jonas's avatar
jonas committed
538
539
540
541
        if header_imgdirx_exists:
            if imgdirx_flipped == 'YES': #should be YES for any L1 data, but mistake in processing software
                flat = flat[:,::-1, :, :]

jonas's avatar
jonas committed
542
        try:
543
544

            if flat_states == 6:
jonas's avatar
jonas committed
545
        
546
547
548
549
                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's avatar
jonas committed
550
                data /= tmp[start_row:start_row + data_size[0],start_col:start_col + data_size[1], np.newaxis, :, np.newaxis]
551
552
553
554
555
556
557
558
559
560
561
562
563
564


            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

jonas's avatar
jonas committed
565
                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
566
        
Jonas Sinjan's avatar
Jonas Sinjan committed
567
568
569
            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
570
571
572
573

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

574
575
576
577
    else:
        print(" ")
        printc('-->>>>>>> No flat field correction mode',color=bcolors.WARNING)

jonas's avatar
jonas committed
578

jonas's avatar
jonas committed
579
580
581
582
583
    #-----------------
    # PREFILTER CORRECTION  
    #-----------------

    if prefilter_f is not None:
jonas's avatar
jonas committed
584
        print(" ")
jonas's avatar
jonas committed
585
586
587
588
589
590
591
592
593
594
595
596
        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
597
        #prefilter = prefilter[:,652:1419,613:1380] #crop the helioseismology data
jonas's avatar
jonas committed
598

jonas's avatar
jonas committed
599
600
601
        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
602

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

jonas's avatar
jonas committed
605
            voltage_list = voltagesData_arr[scan]
jonas's avatar
jonas committed
606
607
            
            for wv in range(6):
jonas's avatar
jonas committed
608

jonas's avatar
jonas committed
609
                v = voltage_list[wv]
jonas's avatar
jonas committed
610

jonas's avatar
jonas committed
611
612
613
614
615
616
617
618
619
620
621
622
                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
623
                imprefilter = (prefilter[:,:, index1]*v1 + prefilter[:,:, index2]*v2)/(v1+v2)
jonas's avatar
jonas committed
624

jonas's avatar
jonas committed
625
                data[:,:,:,wv,scan] /= imprefilter[...,np.newaxis]
jonas's avatar
jonas committed
626
627
628
629
630
631


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

632
633
634
    else:
        print(" ")
        printc('-->>>>>>> No prefilter mode',color=bcolors.WARNING)
jonas's avatar
jonas committed
635

636
637
638
639
    #-----------------
    # FIELD STOP 
    #-----------------

jonas's avatar
jonas committed
640
    if fs_c:
jonas's avatar
jonas committed
641
    
642
643
        print(" ")
        printc("-->>>>>>> Applying field stop",color=bcolors.OKGREEN)
644

645
646
647
        start_time = time.time()
        
        field_stop,_ = load_fits('./field_stop/HRT_field_stop.fits')
648

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

jonas's avatar
jonas committed
651
652
653
        if header_imgdirx_exists:
            if imgdirx_flipped == 'YES': #should be YES for any L1 data, but mistake in processing software
                field_stop = field_stop[:,::-1] #also need to flip the flat data after dark correction
jonas's avatar
jonas committed
654

655
656
657
658
659
660
661
662
663
        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)
664

jonas's avatar
jonas committed
665

jonas's avatar
jonas committed
666
667
668
669
    #-----------------
    # APPLY DEMODULATION 
    #-----------------

670
671
    if demod:

Jonas Sinjan's avatar
Jonas Sinjan committed
672
673
        print(" ")
        printc('-->>>>>>> Demodulating data         ',color=bcolors.OKGREEN)
jonas's avatar
jonas committed
674

jonas's avatar
jonas committed
675
676
        start_time = time.time()

jonas's avatar
jonas committed
677
        data,_ = demod_hrt(data, pmp_temp)
Jonas Sinjan's avatar
Jonas Sinjan committed
678
679
680
681
        
        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
682

683
684
685
    else:
        print(" ")
        printc('-->>>>>>> No demod mode',color=bcolors.WARNING)
jonas's avatar
jonas committed
686
687


jonas's avatar
jonas committed
688
689
690
691
    #-----------------
    # APPLY NORMALIZATION 
    #-----------------

692
    if norm_stokes:
Jonas Sinjan's avatar
Jonas Sinjan committed
693
694
695
696
        
        print(" ")
        printc('-->>>>>>> Normalising Stokes to Quiet Sun',color=bcolors.OKGREEN)
        
jonas's avatar
jonas committed
697
698
        start_time = time.time()

Jonas Sinjan's avatar
Jonas Sinjan committed
699
700
        for scan in range(data_shape[-1]):
            
jonas.sinjan's avatar
jonas.sinjan committed
701
            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
702
703
704
705
706
            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
707

708
709
710
    else:
        print(" ")
        printc('-->>>>>>> No normalising Stokes mode',color=bcolors.WARNING)
jonas's avatar
jonas committed
711
712


jonas's avatar
jonas committed
713
714
715
    #-----------------
    # CROSS-TALK CALCULATION 
    #-----------------
jonas's avatar
jonas committed
716

717
    if ItoQUV:
Jonas Sinjan's avatar
Jonas Sinjan committed
718
719
        
        print(" ")
jonas's avatar
jonas committed
720
        printc('-->>>>>>> Cross-talk correction I to Q,U,V ',color=bcolors.OKGREEN)
jonas's avatar
jonas committed
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
774
775
776
777
778

        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
779
780
781
        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
782
783
784
        
        data *= field_stop[start_row:start_row + data_size[0],start_col:start_col + data_size[1], np.newaxis, np.newaxis, np.newaxis]

785
786
787
788
    else:
        print(" ")
        printc('-->>>>>>> No ItoQUV mode',color=bcolors.WARNING)

jonas's avatar
jonas committed
789
790
791
792
793
794
795
796

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

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

jonas.sinjan's avatar
jonas.sinjan committed
797
    
798
    if out_demod_file:
799
800
801

        #check if the output directory exists, if not, create it
        if not os.path.exists(out_dir): 
jonas's avatar
jonas committed
802
            print(f"{out_dir} does not exist -->>>>>>> Creating it")
803
            os.makedirs(out_dir)
804
        
805
806
        print(" ")
        printc('Saving demodulated data to one _reduced.fits file per scan')
807

808
        def check_filenames(data_f):
809
810
811
812
813
814
            #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
815
816
            #print(uniq_scan_DIDs)
            #print(scan_name_list)
jonas's avatar
jonas committed
817
            if uniq_scan_DIDs == []:
jonas's avatar
jonas committed
818
                print("The scans' DIDs are all unique")
819
820
821
822
823
824

            else:

                for x in uniq_scan_DIDs:
                    number = scan_name_list.count(x)
                    if number > 1: #if more than one
825
                        print(f"The DID: {x} is repeated {number} times")
826
827
828
829
830
831
832
833
                        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)

834
            return scan_name_list
jonas's avatar
jonas committed
835

836
        if out_demod_filename is not None:
jonas's avatar
jonas committed
837
838
839
840
841

            if isinstance(out_demod_filename,str):
                out_demod_filename = [out_demod_filename]

            if int(len(out_demod_filename)) == int(data_shape[-1]):
jonas's avatar
jonas committed
842
                scan_name_list = out_demod_filename
843
844
845
846
847
848
849
850
851
852
                scan_name_defined = True
            else:
                print("Input demod filenames do not match the number of input arrays, reverting to default naming")
                scan_name_defined = False
        else:
            scan_name_defined = False

        if not scan_name_defined: #check if already defined by input, otherwise generate
            scan_name_list = check_filenames(data_f)
        
jonas's avatar
jonas committed
853

854
        for count, scan in enumerate(data_f):
855

856
            with fits.open(scan) as hdu_list:
jonas's avatar
jonas committed
857
                print(f"Writing out demod file as: {scan_name_list[count]}_reduced.fits")
858
859
860
861
862
863
864
865
866
867
868
869
                hdu_list[0].data = data[:,:,:,:,count]
                hdu_list.writeto(out_dir + scan_name_list[count] + '_reduced.fits', overwrite=True)

        # if isinstance(data_f, str):
        #     print(" ")
        #     printc('Saving demodulated data to a _reduced.fits file')

        #     if out_demod_filename is not None:
        #         scan_name = str(out_demod_filename)
            
        #     if 'scan_name' not in locals(): #check if already defined by input, otherwise generate
        #         scan_name = str(data_f.split('.fits')[0][-10:])
870

871
        #     with fits.open(data_f) as hdu_list:
872

873
874
        #         hdu_list[0].data = data
        #         hdu_list.writeto(out_dir + scan_name + '_reduced.fits', overwrite=True)
jonas's avatar
jonas committed
875

876
877
878
    else:
        print(" ")
        printc('-->>>>>>> No output demod file mode',color=bcolors.WARNING)
jonas's avatar
reorg    
jonas committed
879

jonas's avatar
jonas committed
880
881
882
883
884
    #-----------------
    # INVERSION OF DATA WITH CMILOS
    #-----------------

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

886
        if p_milos:
jonas's avatar
jonas committed
887

jonas's avatar
jonas committed
888
889
890
891
892
            try:
                pmilos(data_f, wve_axis_arr, data_shape, cpos_arr, data, rte, field_stop, start_row, start_col, out_rte_filename, out_dir)
                    
            except ValueError:
                print("Running CMILOS instead!")
jonas's avatar
jonas committed
893
                cmilos(data_f, hdr_arr, wve_axis_arr, data_shape, cpos_arr, data, rte, field_stop, start_row, start_col, out_rte_filename, out_dir)
jonas's avatar
jonas committed
894

jonas's avatar
jonas committed
895
        else:
jonas's avatar
jonas committed
896
            cmilos(data_f, hdr_arr, wve_axis_arr, data_shape, cpos_arr, data, rte, field_stop, start_row, start_col, out_rte_filename, out_dir)
jonas's avatar
jonas committed
897

898
899
900
901
    else:
        print(" ")
        printc('-->>>>>>> No RTE Inversion mode',color=bcolors.WARNING)

jonas's avatar
jonas committed
902
903
904
905
906
907
908
909
910
    
    #-----------------
    # CONFIG FILE
    #-----------------

    if config_file:
        print(" ")
        printc('-->>>>>>> Writing out config file ',color=bcolors.OKGREEN)

jonas's avatar
jonas committed
911
        json.dump(saved_args, open(f"{out_dir + scan_name_list[count]}" + '_hrt_pipeline_config.txt', "w"))
jonas's avatar
jonas committed
912

Jonas Sinjan's avatar
Jonas Sinjan committed
913
914
915
916
    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
917
918


jonas's avatar
jonas committed
919
    return data