-
Notifications
You must be signed in to change notification settings - Fork 3
/
tdose_extract_spectra.py
844 lines (723 loc) · 42.8 KB
/
tdose_extract_spectra.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
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
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
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
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
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
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
import numpy as np
import sys
import astropy.io.fits as afits
import collections
import tdose_utilities as tu
import tdose_extract_spectra as tes
import tdose_build_mock_cube as tbmc
import pdb
import scipy.ndimage.filters as snf
import matplotlib as mpl
mpl.use('Agg') # prevent pyplot from opening window; enables closing ssh session with detached screen running TDOSE
import matplotlib.pyplot as plt
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectra(model_cube_file,source_association_dictionary=None,nameext='tdose_spectrum',outputdir='./',clobber=False,
variance_cube_file=None,variance_cube_ext='ERROR',source_model_cube_file=None,source_cube_ext='DATA',
model_cube_ext='DATA',layer_scale_ext='WAVESCL',data_cube_file=None,verbose=True):
"""
Assemble the spectra determined by the wavelength layer scaling of the normalized models
when generating the source model cube
--- INPUT ---
model_cube_file Model cube to base extraction on (using header info and layer scales)
source_association_dictionary Source association dictionary defining what sources should be combined into
objects (individual spectra).
nameext The name extension to use for saved spectra
outputdir Directory to save spectra to
clobber Overwrite spectra if they already exists
variance_cube_file File containing variance cube of data to be used to estimate nois on 1D spectrum
variance_cube_ext Extension of variance cube to use
source_model_cube_file The source model cube defining the individual sources
source_cube_ext Extension of source model cube file that contins source models
model_cube_ext Extension of model cube file that contains model
layer_scale_ext Extension of model cube file that contains the layer scales
data_cube_file File containing original data cube used for extraction of aperture spectra
verbose
--- EXAMPLE OF USE ---
"""
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if verbose: print(' - Loading data needed for spectral assembly')
model_cube = afits.open(model_cube_file)[model_cube_ext].data
model_cube_hdr = afits.open(model_cube_file)[model_cube_ext].header
layer_scale_arr = afits.open(model_cube_file)[layer_scale_ext].data
if variance_cube_file is not None:
stddev_cube = np.sqrt(afits.open(variance_cube_file)[variance_cube_ext].data) # turn varinace into standard deviation
source_model_cube = afits.open(source_model_cube_file)[source_cube_ext].data
else:
stddev_cube = None
source_model_cube = None
Nsources = layer_scale_arr.shape[0]
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if data_cube_file is not None:
if verbose: print(' - Loading data cube ')
data_cube = afits.open(data_cube_file)[model_cube_ext].data
else:
data_cube = None
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if source_association_dictionary is None:
if verbose: print(' - Building default source association dictionary ' \
'(determining what sources are combined into objects), i.e., one source per object ')
sourcIDs_dic = collections.OrderedDict()
for oo in np.arange(int(Nsources)):
sourcIDs_dic[str(oo)] = [oo]
else:
sourcIDs_dic = source_association_dictionary
Nobj = len(list(sourcIDs_dic.keys()))
if verbose: print(' - Found '+str(Nobj)+' objects to generate spectra for in source_association_dictionary ')
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if verbose: print(' - Assembling wavelength vector for spectra ')
wavelengths = np.arange(model_cube_hdr['NAXIS3'])*model_cube_hdr['CD3_3']+model_cube_hdr['CRVAL3']
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
specfiles = []
for oo, key in enumerate(sourcIDs_dic.keys()):
obj_cube_hdr = model_cube_hdr.copy()
try:
specid = str("%.10d" % int(key))
except:
specid = str(key)
specname = outputdir+nameext+'_'+specid+'.fits'
specfiles.append(specname)
sourceIDs = sourcIDs_dic[key]
obj_cube_hdr.append(('OBJID ',specid ,'ID of object'),end=True)
obj_cube_hdr.append(('SRCIDS ',str(sourceIDs) ,'IDs of sources combined in object'),end=True)
if verbose:
infostr = ' - Extracting spectrum '+str("%6.f" % (oo+1))+' / '+str("%6.f" % Nobj)
sys.stdout.write("%s\r" % infostr)
sys.stdout.flush()
sourceoutput = tes.extract_spectrum(sourceIDs,layer_scale_arr,wavelengths,noise_cube=stddev_cube,
source_model_cube=source_model_cube, data_cube=data_cube,
specname=specname,obj_cube_hdr=obj_cube_hdr,clobber=clobber,verbose=True)
if verbose: print('\n - Done extracting spectra. Returning list of fits files generated')
return specfiles
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectrum(sourceIDs,layer_scale_arr,wavelengths,noise_cube=None,source_model_cube=None,
specname='tdose_extract_spectra_extractedspec.fits',obj_cube_hdr=None,data_cube=None,
clobber=False,verbose=True):
"""
Extracting a spectrum based on the layer scale image from the model cube provided a list of sources to combine.
Noise is estimated from the noise cube (of the data)
If all layer_scales are 1 a data_cube for the extractions is expected
--- INPUT ---
sourceIDs The source IDs to combine into spectrum
layer_scale_arr Layer scale array (or image) produced when generating the model cube
fractional flux belonging to the source in each pixel
wavelengths Wavelength vector to use for extracted 1D spectrum.
noise_cube Cube with uncertainties (sqrt(variance)) of data cube to be used for estimating 1D uncertainties
To estimate S/N and 1D noise, providing a source model cube is required
source_model_cube Source model cube containing the model cube for each individual source seperately
Needed in order to estimate noise from noise-cube
specname Name of file to save spectrum to
obj_cube_hdr Provide a template header to save the object cube (from combining the individual source cube)
as an extension to the extracted spectrum
data_cube In case all layers scales are 1, it is assumed that the source_model_cube contains a mask for the
spectral extraction, which will then be performed on this data_cube.
clobber To overwrite existing files set clobber=True
verbose Toggle verbosity
--- EXAMPLE OF USE ---
"""
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if verbose: print(' - Checking shape of wavelengths and layer_scale_arr')
if wavelengths.shape[0] != layer_scale_arr.shape[1]:
sys.exit(' ---> Shape of wavelength vector ('+str(wavelengths.shape)+
') and wavelength dimension of layer scale array ('+
layer_scale_arr.shape[1].shape+') do not match.')
else:
if verbose: print(' dimensions match; proceeding...')
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if verbose: print(' - Checking all sources have spectra in layer_scale_arr')
maxsource = np.max(sourceIDs)
if maxsource >= layer_scale_arr.shape[0]:
sys.exit(' ---> Sources in list '+str(str(sourceIDs))+
' not available among '+str(layer_scale_arr.shape[0])+' sources in layer_scale_arr.')
else:
if verbose: print(' All sources exist in layer_scale_arr; proceeding...')
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if verbose: print(' - Assembling object spectrum from source scaling')
source_ent = np.asarray(sourceIDs).astype(int)
if (layer_scale_arr == 1).all():
if verbose: print(' - All layer scales are 1; assuming source model cube contain mask for spectral extraction')
object_cube = np.sum(np.abs(source_model_cube[source_ent,:,:]),axis=0)
if data_cube is None:
sys.exit(' ---> Did not find a data cube to extrac spectra from as expected')
object_mask = (object_cube == 0) # masking all zeros in object mask
invalid_mask = np.ma.masked_invalid(data_cube).mask
comb_mask = (invalid_mask | object_mask)
spec_1D_masked = np.sum(np.sum( np.ma.array(data_cube,mask=comb_mask) ,axis=1),axis=1)
spec_1D = spec_1D_masked.filled(fill_value=0.0)
if noise_cube is not None:
if verbose: print(' Calculating noise as d_spec_k = sqrt( SUMij d_pix_ij**2 ), i.e., as the sqrt of variances summed')
invalid_mask_noise = np.ma.masked_invalid(noise_cube).mask
comb_mask = (comb_mask | invalid_mask_noise)
variance_1D_masked = np.ma.array(noise_cube,mask=comb_mask)**2
noise_1D_masked = np.sqrt( np.sum( np.sum( variance_1D_masked, axis=1), axis=1) )
noise_1D = noise_1D_masked.filled(fill_value=np.nan)
if verbose: print(' Generating S/N vector')
SN_1D = spec_1D / noise_1D
else:
if verbose: print(' - No "noise_cube" provided. Setting all errors and S/N values to NaN')
SN_1D = np.zeros(spec_1D.shape)*np.NaN
noise_1D = np.zeros(spec_1D.shape)*np.NaN
else:
if verbose: print(' - Some layer scales are different from 1; hence assembling spectra using layer scales')
if len(source_ent) < 1:
spec_1D = layer_scale_arr[source_ent,:]
else:
spec_1D = np.sum( layer_scale_arr[source_ent,:],axis=0)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if noise_cube is not None:
if verbose: print(' - Estimate S/N at each wavelength for 1D spectrum (see Eq. 16 of Kamann+2013)')
if verbose: print(' Estimating fraction of flux in each pixel wrt. total flux in each layer')
object_cube = np.sum((source_model_cube[source_ent,:,:,:]),axis=0) # summing source models for all source IDs
fluxfrac_cube_sents = np.zeros(source_model_cube.shape[1:])
for sent in source_ent:
object_cube_sent = np.sum((source_model_cube[[sent],:,:,:]),axis=0) # getting source model for model 'sent'
fluxscale1D_sent = layer_scale_arr[sent,:]
fluxfrac_cube_sent = object_cube_sent / fluxscale1D_sent[:,None,None]
fluxfrac_cube_sents = fluxfrac_cube_sents + fluxfrac_cube_sent
fluxfrac_cube = fluxfrac_cube_sents / len(source_ent) # renormalizing flux-fraction cube
if verbose: print(' Defining pixel mask (ignoring NaN pixels) ') #+\
# 'and pixels with <'+str(fluxfrac_min)+' of total pixel flux in model cube) '
# pix_mask = (fluxfrac_cube < fluxfrac_min)
invalid_mask1 = np.ma.masked_invalid(fluxfrac_cube).mask
invalid_mask2 = np.ma.masked_invalid(noise_cube).mask
# combining mask making sure all individual mask pixels have True for it to be true in combined mask
comb_mask = (invalid_mask1 | invalid_mask2) # | pix_mask
if verbose: print(' Calculating noise propogated as d_spec_k = 1/sqrt( SUMij (fluxfrac_ij**2 / d_pix_ij**2) )')
squared_ratio = np.ma.array(fluxfrac_cube,mask=comb_mask)**2 / np.ma.array(noise_cube,mask=comb_mask)**2
inv_noise_masked = np.sqrt( np.sum( np.sum( squared_ratio, axis=1), axis=1) )
noise_1D = (1.0/inv_noise_masked).filled(fill_value=0.0)
if verbose: print(' Generating S/N vector')
SN_1D = spec_1D / noise_1D
else:
if verbose: print(' - No "noise_cube" provided. Setting all errors and S/N values to NaN')
SN_1D = np.zeros(spec_1D.shape)*np.NaN
noise_1D = np.zeros(spec_1D.shape)*np.NaN
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if verbose: print(' - Saving extracted 1D spectrum and source cube to \n '+specname)
mainHDU = afits.PrimaryHDU() # primary HDU
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
c1 = afits.Column(name='wave', format='D', unit='ANGSTROMS', array=wavelengths)
c2 = afits.Column(name='flux', format='D', unit='', array=spec_1D)
c3 = afits.Column(name='fluxerror', format='D', unit='', array=noise_1D)
c4 = afits.Column(name='s2n', format='D', unit='', array=SN_1D)
coldefs = afits.ColDefs([c1,c2,c3,c4])
th = afits.BinTableHDU.from_columns(coldefs) # creating default header
# writing hdrkeys:'---KEY--', '----------------MAX LENGTH COMMENT-------------'
th.header.append(('EXTNAME ','SPEC1D' ,'cube containing source'),end=True)
head = th.header
tbHDU = afits.BinTableHDU.from_columns(coldefs, header=head)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if obj_cube_hdr is not None:
objHDU = afits.ImageHDU(object_cube)
for hdrkey in list(obj_cube_hdr.keys()):
if not hdrkey in list(objHDU.header.keys()):
objHDU.header.append((hdrkey,obj_cube_hdr[hdrkey],obj_cube_hdr.comments[hdrkey]),end=True)
try:
objHDU.header.append(('EXTNAMEC',objHDU.header['EXTNAME'] ,'EXTNAME of original source cube'),end=True)
del objHDU.header['EXTNAME']
except:
pass
objHDU.header.append(('EXTNAME ','SOURCECUBE' ,'cube containing source'),end=True)
hdulist = afits.HDUList([mainHDU,tbHDU,objHDU])
else:
hdulist = afits.HDUList([mainHDU,tbHDU])
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
hdulist.writeto(specname, overwrite=clobber)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
return wavelengths, spec_1D, noise_1D, object_cube
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectra_viasourcemodelcube(datacube,sourcemodelcube,wavelengths,speclist,specids='None',outputdir='./',
noisecube=False,sourcemodel_hdr='None',verbose=True):
"""
Wrapper for tes.extract_spectrum_viasourcemodelcube() to extract mutliple spectra
--- INPUT ----
datacube Datacube to extract spectra from
sourcemodelcube Cube containing the source models for each object used as "extraction cube"
Dimensions should be [Nsources,datacube.shape]
wavelengths Wavelength vector to use for extracted 1D spectrum.
speclist List of spectra to extract. Indexes corresponding to the source models in the
sourcemodlecube
specids List of IDs to use in naming of output for source models referred to in "speclist"
outputdir Directory to store spectra to
noisecube Cube with uncertainties (sqrt(variance)) of data cube to be used in extraction
souremodel_hdr If not 'None' provide a basic fits header for the source model cubes extracted
and they will be appended to the individual output fits file containing the extracted
spectra.
verbose Toggle verbosity
--- EXAMPLE OF USE ---
"""
if verbose: print(' - Check that source models indicated are present in source model cube ')
specnames = []
Nmodels = sourcemodelcube.shape[0]
maxobj = np.max(speclist)
if maxobj >= Nmodels:
sys.exit(' ---> Object model "'+str(maxobj)+'" is not included in source model cube (models start at 0)')
else:
if verbose: print(' All object models appear to be included in the '+str(Nmodels)+' source models found in cube')
if datacube.shape != sourcemodelcube[0].shape:
sys.exit(' ---> Shape of datacube ('+str(datacube.shape)+') and shape of source models ('+
sourcemodelcube[0].shape+') do not match.')
sourcemodel_sum = np.sum(sourcemodelcube,axis=0)
for ss, spec in enumerate(speclist):
if specids == 'None':
specid = spec
else:
specid = specids[ss]
specname = outputdir+'tdose_spectrum_'+str("%.12d" % specid)+'.fits'
specnames.append(specname)
sourcemodel = sourcemodelcube[spec,:,:,:]
sourceweights = sourcemodel/sourcemodel_sum # fractional flux of model for given source in each pixel
sourcemodel_hdr.append(('OBJMODEL',spec ,'Source model number in parent source model cube'),end=True)
sourcemodel_hdr.append(('OBJID ',specid ,'ID of source'),end=True)
if verbose:
infostr = ' - Extracting spectrum '+str("%6.f" % (spec+1))+' / '+str("%6.f" % len(speclist))
sys.stdout.write("%s\r" % infostr)
sys.stdout.flush()
sourceoutput = tes.extract_spectrum_viasourcemodelcube(datacube,sourceweights,wavelengths,specname=specname,
noisecube=noisecube,spec1Dmethod='sum',
sourcecube_hdr=sourcemodel_hdr,verbose=verbose)
if verbose: print('\n - Done extracting spectra. Returning list of fits files containing spectra')
return specnames
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def extract_spectrum_viasourcemodelcube(datacube,sourceweights,wavelengths,
specname='tdose_extract_spectra_extractedspec.fits',
noisecube=None,spec1Dmethod='sum',sourcecube_hdr='None',verbose=True):
"""
Extracting a spectrum from a data cube given a source model (cube) to be used as 'extraction cube'
--- INPUT ---
datacube Datacube to extract spectra from
sourceweights Weights from source model to use as "extraction cube". The weights should contain the
fractional flux belonging to the source in each pixel
wavelengths Wavelength vector to use for extracted 1D spectrum.
specname Name of spectrum to generate
noisecube Cube with uncertainties (sqrt(variance)) of data cube to be used in extraction
spec1Dmethod Method used to extract 1D spectrum from source cube with
sourcecube_hdr If not 'None' provide a fits header for the source cube and it ill be appended to the
output fits file.
verbose Toggle verbosity
--- EXAMPLE OF USE ---
"""
if verbose: print(' - Checking shape of data and source model cubes')
if datacube.shape != sourceweights.shape:
sys.exit(' ---> Shape of datacube ('+str(datacube.shape)+') and source weights ('+
sourceweights.shape+') do not match.')
else:
if verbose: print(' dimensions match; proceeding with extraction ')
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if verbose: print(' - Applying weights to "datacube" to obtain source cube ')
sourcecube = datacube*sourceweights
if noisecube is not None:
if verbose: print(' - Using "noisecube" for error propagation ')
datanoise = noisecube
else:
if verbose: print(' - No "noisecube" provided. Setting all errors to 1')
datanoise = np.ones(datacube.shape)
if verbose: print(' - Assuming uncertainty on source weights equals the datanoise when propgating errors')
sourceweights_err = datanoise
sourcecube_err = sourcecube * np.sqrt( (datanoise/datacube)**2 + (sourceweights_err/sourceweights)**2 )
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if verbose: print(' - Generating 1D spectrum from source cube via:')
spec_wave = wavelengths
maskinvalid = np.ma.masked_invalid(sourcecube * sourcecube_err).mask
if spec1Dmethod == 'sum':
if verbose: print(' Simple summation of fluxes in sourcecube.')
spec_flux = np.sum(np.sum(np.ma.array(sourcecube,mask=maskinvalid),axis=1),axis=1).filled()
if verbose: print(' Errors are propagated as sum of squares.')
spec_err = np.sqrt( np.sum( np.sum(np.ma.array(sourcecube_err,mask=maskinvalid)**2,axis=1),axis=1) ).filled()
elif spec1Dmethod == 'sum_SNweight':
pdb.set_trace()
else:
sys.exit(' ---> The chosen spec1Dmethod ('+str(spec1Dmethod)+') is invalid')
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if verbose: print(' - Saving extracted 1D spectrum and source cube to \n '+specname)
mainHDU = afits.PrimaryHDU() # primary HDU
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
c1 = afits.Column(name='wave', format='D', unit='ANGSTROMS', array=spec_wave)
c2 = afits.Column(name='flux', format='D', unit='', array=spec_flux)
c3 = afits.Column(name='fluxerror', format='D', unit='', array=spec_err)
coldefs = afits.ColDefs([c1,c2,c3])
th = afits.BinTableHDU.from_columns(coldefs) # creating default header
# writing hdrkeys:'---KEY--', '----------------MAX LENGTH COMMENT-------------'
th.header.append(('EXTNAME ','SPEC1D' ,'cube containing source'),end=True)
th.header.append(('SPECMETH' , spec1Dmethod ,'Method used for spectral extraction'),end=True)
head = th.header
tbHDU = afits.BinTableHDU.from_columns(coldefs, header=head)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if sourcecube_hdr != 'None':
sourceHDU = afits.ImageHDU(sourcecube) # default HDU with default minimal header
for hdrkey in list(sourcecube_hdr.keys()):
if not hdrkey in list(sourceHDU.header.keys()):
sourceHDU.header.append((hdrkey,sourcecube_hdr[hdrkey],sourcecube_hdr.comments[hdrkey]),end=True)
sourceHDU.header.append(('EXTNAME ','SOURCECUBE' ,'cube containing source'),end=True)
hdulist = afits.HDUList([mainHDU,tbHDU,sourceHDU])
else:
hdulist = afits.HDUList([mainHDU,tbHDU])
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
hdulist.writeto(specname, overwrite=True)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
return sourcecube, sourcecube_err, spec_wave, spec_flux, spec_err
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def plot_1Dspecs(filelist,plotname='./tdose_1Dspectra.pdf',colors=None,labels=None,plotSNcurve=False,
tdose_wavecol='wave',tdose_fluxcol='flux',tdose_errcol='fluxerror',
simsources=None,simsourcefile='/Users/kschmidt/work/TDOSE/mock_cube_sourcecat161213_all.fits',
sim_cube_dim=None,comparisonspecs=None,comp_colors=['blue'],comp_labels=None,
comp_wavecol='WAVE_AIR',comp_fluxcol='FLUX',comp_errcol='FLUXERR',
xrange=None,yrange=None,showspecs=False,shownoise=True,
skyspecs=None,sky_colors=['red'],sky_labels=['sky'],
sky_wavecol='lambda',sky_fluxcol='data',sky_errcol='stat',
showlinelists=None,linelistcolors=['gray'],smooth=0,ylog=False,
plotratio=False,
verbose=True,pubversion=False):
"""
Plots of multiple 1D spectra
--- INPUT ---
filelist List of spectra filenames to plot
plotname Name of plot to generate
colors Colors of the spectra in filelist to use
labels Labels of the spectra in filelist to use
plotSNcurve Show signal-to-noise curve instead of flux spectra
tdose_wavecol Wavelength column of the spectra in filelist
tdose_fluxcol Flux column of the spectra in filelist
tdose_errcol Flux error column of the spectra in filelist
simsources To plot simulated sources provide ids here
simsourcefile Source file with simulated sources to plot
sim_cube_dim Dimensions of simulated cubes
comparisonspecs To plot comparison spectra provide the filenames of those here
comp_colors Colors of the spectra in comparisonspecs list to use
comp_labels Labels of the spectra in comparisonspecs list to use
comp_wavecol Wavelength column of the spectra in comparisonspecs list
comp_fluxcol Flux column of the spectra in comparisonspecs list
comp_errcol Flux error column of the spectra in comparisonspecs list
xrange Xrange of plot
yrange Yrange of plot
showspecs To show plot instead of storing it to disk set showspecs=True
shownoise To add noise envelope around spectrum set shownoise=True
skyspecs To plot sky spectra provide the filenames of those here
sky_colors Colors of the spectra in skyspecs list to use
sky_labels Labels of the spectra in skyspecs list to use
sky_wavecol Wavelength column of the spectra in skyspecs list
sky_fluxcol Flux column of the spectra in skyspecs list
sky_errcol Flux error column of the spectra in skyspecs list
showlinelists To show line lists provide a list of arrays of dimension (Nlines,2) where each row in the
arrays contains [waveobs, name], where 'waveobs' is the observed wavelengths and 'name' is
a string with the name of each of the Nlines postions to mark on the spectrum.
linelistcolors List of colors for line lists provided in showlinelists
smooth To smooth the spectra, provide sigma of the 1D gaussian smoothing kernel to apply.
For smooth = 0, no smoothing is performed.
ylog To plot y-axis in log scale set to true
plotratio To plot the ratio between the main spectrum and the comparison spectra instead of the actual
spectra, set this keyword to true.
verbose Toggle verbosity
pubversion Generate more publication friendly version of figure
"""
if len(filelist) == 1:
if verbose: print(' - Plotting data from '+filelist[0])
else:
if verbose: print(' - Plotting data from filelist ')
if pubversion:
fig = plt.figure(figsize=(6, 3))
fig.subplots_adjust(wspace=0.1, hspace=0.1,left=0.15, right=0.95, bottom=0.18, top=0.83)
Fsize = 12
else:
fig = plt.figure(figsize=(10, 3))
fig.subplots_adjust(wspace=0.1, hspace=0.1,left=0.06, right=0.81, bottom=0.15, top=0.95)
Fsize = 10
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# Looking for flux units in spectra
bunit = 'BUNIT FLUX' # Default BUNIT
for unitspec in filelist:
if bunit == 'BUNIT FLUX':
try:
sourcecubehdr = afits.open(unitspec)['SOURCECUBE'].header
bunit = sourcecubehdr['BUNIT']
except:
try: # Backwards compatibility to TDOSE v2.0 extractions
sourcecubehdr = afits.open(unitspec)[2].header
bunit = sourcecubehdr['BUNIT']
except:
pass
if bunit == 'BUNIT FLUX':
if verbose: print(' - Did not find BUNIT in SOURCECUBE header for any spectra in filelist - are they not from TDOSE?')
if bunit == '10**(-20)*erg/s/cm**2/Angstrom': # Making bunit LaTeXy for MUSE-Wide BUNIT format
bunit = '1e-20 erg/s/cm$^2$/\AA'
else:
bunit = '$'+bunit+'$' # minimizing pronlems with LaTeXing plot axes
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
lthick = 1
plt.rc('text', usetex=True)
plt.rc('font', family='serif',size=Fsize)
plt.rc('xtick', labelsize=Fsize)
plt.rc('ytick', labelsize=Fsize)
plt.clf()
plt.ioff()
#plt.title(plotname.split('TDOSE 1D spectra'),fontsize=Fsize)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
for ff, specfile in enumerate(filelist):
specdat = afits.open(specfile)[1].data
if colors is None:
spec_color = None
else:
spec_color = colors[ff]
if labels is None:
spec_label = specfile
else:
spec_label = labels[ff]
if xrange is not None:
goodent = np.where((specdat[tdose_wavecol] > xrange[0]) & (specdat[tdose_wavecol] < xrange[1]))[0]
if goodent == []:
if verbose: print(' - The chosen xrange is not covered by the input spectrum. Plotting full spectrum')
goodent = np.arange(len(specdat[tdose_wavecol]))
else:
goodent = np.arange(len(specdat[tdose_wavecol]))
if plotSNcurve:
try:
s2ndat = specdat['s2n'][goodent]
except:
s2ndat = specdat[tdose_fluxcol][goodent]/specdat[tdose_errcol][goodent]
if smooth > 0:
s2ndat = snf.gaussian_filter(s2ndat, smooth)
if not plotratio:
plt.plot(specdat[tdose_wavecol][goodent],s2ndat,color=spec_color,lw=lthick, label=spec_label)
ylabel = 'S/N'
else:
plt.plot(specdat[tdose_wavecol][goodent],s2ndat/s2ndat,color=spec_color,lw=lthick, label=None)
ylabel = 'S/N ratio'
#plotname = plotname.replace('.pdf','_S2N.pdf')
else:
fillalpha = 0.30
fluxdat = specdat[tdose_fluxcol][goodent]
errlow = specdat[tdose_fluxcol][goodent]-specdat[tdose_errcol][goodent]
errhigh = specdat[tdose_fluxcol][goodent]+specdat[tdose_errcol][goodent]
if smooth > 0:
fluxdat = snf.gaussian_filter(fluxdat, smooth)
errlow = snf.gaussian_filter(errlow, smooth)
errhigh = snf.gaussian_filter(errhigh, smooth)
if smooth > 0:
fluxdat = snf.gaussian_filter(fluxdat, smooth)
if not plotratio:
if shownoise:
plt.fill_between(specdat[tdose_wavecol][goodent],errlow,errhigh,
alpha=fillalpha,color=spec_color)
plt.plot(specdat[tdose_wavecol][goodent],fluxdat,
color=spec_color,lw=lthick, label=spec_label)
ylabel = tdose_fluxcol
else:
plt.plot(specdat[tdose_wavecol][goodent],fluxdat/fluxdat,
color=spec_color,lw=lthick, label=None)
ylabel = tdose_fluxcol+' ratio '
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if simsources is not None:
sim_total = np.zeros(len(specdat[tdose_wavecol]))
for sourcenumber in simsources:
sourcedat = afits.open(simsourcefile)[1].data
xpos = sourcedat['xpos'][sourcenumber]
ypos = sourcedat['ypos'][sourcenumber]
fluxscale = sourcedat['fluxscale'][sourcenumber]
sourcetype = sourcedat['sourcetype'][sourcenumber]
spectype = sourcedat['spectype'][sourcenumber]
sourcecube = tbmc.gen_source_cube([ypos,xpos],fluxscale,sourcetype,spectype,cube_dim=sim_cube_dim,
verbose=verbose,showsourceimgs=False)
simspec = np.sum( np.sum(sourcecube, axis=1), axis=1)
sim_total = sim_total + simspec
if smooth > 0:
simspec = snf.gaussian_filter(simspec, smooth)
plt.plot(specdat[tdose_wavecol],simspec,'--',color='black',lw=lthick)
plt.plot(specdat[tdose_wavecol],sim_total,'--',color='black',lw=lthick,
label='Sim. spectrum: \nsimsource='+str(simsources))
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if comparisonspecs is not None:
for cc, comparisonspec in enumerate(comparisonspecs):
compdat = afits.open(comparisonspec)[1].data
if xrange is not None:
goodent = np.where((compdat[comp_wavecol] > xrange[0]) & (compdat[comp_wavecol] < xrange[1]))[0]
if goodent == []:
if verbose: print(' - The chosen xrange is not covered by the comparison spectrum. Plotting full spectrum')
goodent = np.arange(len(compdat[comp_wavecol]))
else:
goodent = np.arange(len(compdat[comp_wavecol]))
if comp_colors is None:
comp_color = None
else:
comp_color = comp_colors[cc]
if comp_labels is None:
comp_label = comparisonspec
else:
comp_label = comp_labels[cc]
if plotSNcurve:
s2ncompdat = compdat[comp_fluxcol][goodent]/compdat[comp_errcol][goodent]
if smooth > 0:
s2ncompdat = snf.gaussian_filter(s2ncompdat, smooth)
if not plotratio:
plt.plot(compdat[comp_wavecol][goodent],s2ncompdat,
color=comp_color,lw=lthick, label=comp_label)
else:
plt.plot(compdat[comp_wavecol][goodent],s2ndat/s2ncompdat,
color=comp_color,lw=lthick, label=comp_label)
else:
fillalpha = 0.30
fluxcompdat = compdat[comp_fluxcol][goodent]
errlow = compdat[comp_fluxcol][goodent]-compdat[comp_errcol][goodent]
errhigh = compdat[comp_fluxcol][goodent]+compdat[comp_errcol][goodent]
if smooth > 0:
fluxcompdat = snf.gaussian_filter(fluxcompdat, smooth)
errlow = snf.gaussian_filter(errlow, smooth)
errhigh = snf.gaussian_filter(errhigh, smooth)
if not plotratio:
if shownoise:
plt.fill_between(compdat[comp_wavecol][goodent],errlow,errhigh,
alpha=fillalpha,color=comp_color)
plt.plot(compdat[comp_wavecol][goodent],fluxcompdat,
color=comp_color,lw=lthick, label=comp_label)
else:
plt.plot(compdat[comp_wavecol][goodent],fluxdat/fluxcompdat,
color=comp_color,lw=lthick, label=comp_label)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if skyspecs is not None:
for ss, skyspec in enumerate(skyspecs):
skydat = afits.open(skyspec)[1].data
if xrange is not None:
goodent = np.where((skydat[sky_wavecol] > xrange[0]) & (skydat[sky_wavecol] < xrange[1]))[0]
if goodent == []:
if verbose: print(' - The chosen xrange is not covered by the sky spectrum. Plotting full spectrum')
goodent = np.arange(len(skydat[sky_wavecol]))
else:
goodent = np.arange(len(skydat[sky_wavecol]))
if sky_colors is None:
sky_color = None
else:
sky_color = sky_colors[ss]
if sky_labels is None:
sky_label = skyspec
else:
sky_label = sky_labels[ss]
if plotSNcurve:
s2nsky = skydat[sky_fluxcol][goodent]/skydat[sky_errcol][goodent]
if smooth > 0:
s2nsky = snf.gaussian_filter(s2nsky, smooth)
plt.plot(skydat[sky_wavecol][goodent],s2nsky,
color=sky_color,lw=lthick, label=sky_label)
else:
fillalpha = 0.30
fluxsky = skydat[sky_fluxcol][goodent]
errlow = skydat[sky_fluxcol][goodent]-skydat[sky_errcol][goodent]
errhigh = skydat[sky_fluxcol][goodent]+skydat[sky_errcol][goodent]
if smooth > 0:
fluxsky = snf.gaussian_filter(fluxsky, smooth)
errlow = snf.gaussian_filter(errlow, smooth)
errhigh = snf.gaussian_filter(errhigh, smooth)
if shownoise:
plt.fill_between(skydat[sky_wavecol][goodent],errlow,errhigh,
alpha=fillalpha,color=sky_color)
plt.plot(skydat[sky_wavecol][goodent],fluxsky,
color=sky_color,lw=lthick, label=sky_label)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if xrange is None:
xvals = [4800,9300]
else:
xvals = xrange
plt.plot(xvals,[0,0],'--k',lw=lthick)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
plt.xlabel('Wavelength [\AA]', fontsize=Fsize)
if pubversion:
if plotSNcurve:
ylabel = 'Signal-to-Noise'
else:
ylabel = 'Flux ['+str(bunit)+']'
if plotratio:
ylabel = ylabel+' ratio'
plt.ylabel(ylabel, fontsize=Fsize)
if ylog:
plt.yscale('log')
if yrange is not None:
plt.ylim(yrange)
if xrange is not None:
plt.xlim(xrange)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if showlinelists is not None:
for sl, showlinelist in enumerate(showlinelists):
ymin, ymax = plt.ylim()
xmin, xmax = plt.xlim()
for ww, wave in enumerate(showlinelist[:,0]):
wave = float(wave)
if (wave < xmax) & (wave > xmin):
plt.plot([wave,wave],[ymin,ymax],linestyle='--',color=linelistcolors[sl],lw=lthick)
plt.text(wave,ymin+1.03*np.abs([ymax-ymin]),showlinelist[:,1][ww],color=linelistcolors[sl], fontsize=Fsize-2., ha='center')
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if pubversion:
leg = plt.legend(fancybox=True, loc='upper center',prop={'size':Fsize-2},ncol=4,numpoints=1,
bbox_to_anchor=(0.44, 1.27)) # add the legend
else:
leg = plt.legend(fancybox=True, loc='upper right',prop={'size':Fsize},ncol=1,numpoints=1,
bbox_to_anchor=(1.25, 1.03)) # add the legend
leg.get_frame().set_alpha(0.7)
if showspecs:
if verbose: print(' Showing plot (not saving to file)')
plt.show()
else:
if verbose: print(' Saving plot to',plotname)
plt.savefig(plotname)
plt.clf()
plt.close('all')
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =
def plot_histograms(datavectors,plotname='./tdose_cubehist.pdf',colors=None,labels=None,bins=None,
xrange=None,yrange=None,verbose=True,norm=True,ylog=True):
"""
Plot histograms of a set of data vectors.
--- INPUT ---
datavectors Set of data vectors to plot histograms of
plotname Name of plot to generate
colors Colors to use for histograms
labels Labels for the data vectors
bins Bins to use for histograms. Can be generated with np.arange(minval,maxval+binwidth,binwidth)
xrange Xrange of plot
yrange Yrange of plot
verbose Toggle verbosity
norm Noramlize the histograms
ylog Use a logarithmic y-axes when plotting
"""
Ndat = len(datavectors)
if verbose: print(' - Plotting histograms of N = '+str(Ndat)+' data vectors')
if colors is None:
colors = ['blue']*Ndat
if labels is None:
labels = ['data vector no. '+str(ii+1) for ii in np.arange(Ndat)]
if bins is None:
bins = np.arange(-100,102,2)
fig = plt.figure(figsize=(10, 3))
fig.subplots_adjust(wspace=0.1, hspace=0.1,left=0.08, right=0.81, bottom=0.1, top=0.95)
Fsize = 10
lthick = 1
plt.rc('text', usetex=True)
plt.rc('font', family='serif',size=Fsize)
plt.rc('xtick', labelsize=Fsize)
plt.rc('ytick', labelsize=Fsize)
plt.clf()
plt.ioff()
#plt.title(plotname.split('TDOSE 1D spectra'),fontsize=Fsize)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
for dd, datavec in enumerate(datavectors):
hist = plt.hist(datavec[~np.isnan(datavec)],color=colors[dd],bins=bins,histtype="step",lw=lthick,
label=labels[dd],normed=norm)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
if yrange is None:
yvals = [1e-5,1e8]
else:
yvals = yrange
plt.plot([0,0],yvals,'--k',lw=lthick)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
plt.xlabel('', fontsize=Fsize)
plt.ylabel('\#', fontsize=Fsize)
if yrange is not None:
plt.ylim(yrange)
if xrange is not None:
plt.xlim(xrange)
if ylog:
plt.yscale('log')
leg = plt.legend(fancybox=True, loc='upper right',prop={'size':Fsize},ncol=1,numpoints=1,
bbox_to_anchor=(1.25, 1.03)) # add the legend
leg.get_frame().set_alpha(0.7)
if verbose: print(' Saving plot to',plotname)
plt.savefig(plotname)
plt.clf()
plt.close('all')
# = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =