bbi_header.py 12.1 KB
Newer Older
1
2
3
import datetime
import re
from collections import OrderedDict
4
5
6
7
8
from glob import glob
from os.path import basename, isdir, join

import numpy as np
import tqdm
Sebastian Hoch's avatar
Sebastian Hoch committed
9
10
import yaml
from astropy.io import fits
11
12
from sunpy.coordinates.sun import earth_distance

13
from kis_tools.bbi.util import get_bbi_wl
14
from kis_tools.generic.settings import settings
15
from kis_tools.generic.settings.settings import get_settings
16
17
18
19
20
21
22
from kis_tools.interface import BasicArgParser
from kis_tools.util.util import setup_outdir
from .bbi_sort_observations import get_obs_info
from ..util.util import add_history, date_from_fn
from ..util.util import get_filter

settings = get_settings("bbi")
Sebastian Hoch's avatar
Sebastian Hoch committed
23
24
25
26
27
28
29
30


class OrderedLoader(yaml.Loader):
    """YAML loader for keeping order of items.
    This class is intended to be used like::
        >>> with open(filename) as f:
        >>> ... yaml.load(f, Loader=OrderedLoader)
    """
31

Sebastian Hoch's avatar
Sebastian Hoch committed
32
33
34
35
36
37
38
39
40
41
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        tag = yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG

        def constructor(loader, node):
            return OrderedDict(loader.construct_pairs(node))

        self.add_constructor(tag, constructor)

42

43
44
def get_separators(yaml_file):
    try:
Carl Schaffer's avatar
Carl Schaffer committed
45
        f = open(yaml_file, "rb")
46
47
48
        yamldata = f.read()
        f.close()
    except:
Carl Schaffer's avatar
Carl Schaffer committed
49
        print("Could not open file {:}".format(yaml_file))
50
51
52
53
54
55
        return None

    sep_data = yaml.load(yamldata, OrderedLoader)["separators"]
    return sep_data


56
def get_generic_header(filename=settings.header_template, version="reconstruction"):
Carl Schaffer's avatar
Carl Schaffer committed
57
58
59
60
61
62
63
    """
    Retrieve template for Header from a Yaml file
    :param filename: path to YAML file
    :param version: string 'reconstruction' or 'burst', differentiates header version
    :return: header: ordered dict of field configurations

    """
Sebastian Hoch's avatar
Sebastian Hoch committed
64
    try:
Carl Schaffer's avatar
Carl Schaffer committed
65
        f = open(filename, "rb")
Sebastian Hoch's avatar
Sebastian Hoch committed
66
67
68
        yamldata = f.read()
        f.close()
    except:
Carl Schaffer's avatar
Carl Schaffer committed
69
        print("Could not open file {:}".format(filename))
Sebastian Hoch's avatar
Sebastian Hoch committed
70
71
72
73
74
75
        return None

    headerdata = yaml.load(yamldata, OrderedLoader)
    try:
        header = fits.Header(headerdata[version].keys())
    except:
Carl Schaffer's avatar
Carl Schaffer committed
76
        print("Version {:} not found".format(version))
Sebastian Hoch's avatar
Sebastian Hoch committed
77
78
79
80
        return None

    headerversion = headerdata[version]
    for card in headerversion:
81
        try:
Sebastian Hoch's avatar
Sebastian Hoch committed
82
            gencard = headerversion.get(card)
Carl Schaffer's avatar
Carl Schaffer committed
83
84
            standard = gencard["DEFAULT"]
            comment = gencard.get("COMMENT", "")
85
            if comment is None:
Carl Schaffer's avatar
Carl Schaffer committed
86
                comment = ""
Sebastian Hoch's avatar
Sebastian Hoch committed
87
88
            header[card] = (standard, comment)
        except:
Carl Schaffer's avatar
Carl Schaffer committed
89
            print("Could not update header")
Sebastian Hoch's avatar
Sebastian Hoch committed
90
91
92
            print(card, gencard)
    return header

93

Sebastian Hoch's avatar
Sebastian Hoch committed
94
def get_generic_patches(filename, version):
95
    """
Carl Schaffer's avatar
Carl Schaffer committed
96
    Read list of generic 'easy' patches from a file on disk. Easy patches are renaming keywords and similar.
97
    :param filename: yaml file with generic patches
Carl Schaffer's avatar
Carl Schaffer committed
98
99
    :param version: string 'reconstruction' or 'burst', differentiates header version
    :return: patches ordered dict
100
    """
Sebastian Hoch's avatar
Sebastian Hoch committed
101
    try:
Carl Schaffer's avatar
Carl Schaffer committed
102
        f = open(filename, "rb")
Sebastian Hoch's avatar
Sebastian Hoch committed
103
104
105
        yamldata = f.read()
        f.close()
    except:
Carl Schaffer's avatar
Carl Schaffer committed
106
        print("Could not open file {:}".format(filename))
Sebastian Hoch's avatar
Sebastian Hoch committed
107
108
109
110
111
112
113
        return None

    patchesdata = yaml.load(yamldata, OrderedLoader)

    try:
        patches = patchesdata[version]
    except:
Carl Schaffer's avatar
Carl Schaffer committed
114
        print("Version {:} not found".format(version))
Sebastian Hoch's avatar
Sebastian Hoch committed
115
116
117
118
        return None

    return patches

119

120
def apply_patches(new_header, old_header, patches):
Carl Schaffer's avatar
Carl Schaffer committed
121
122
123
    """
    Loop over patches header and insert Patches into new header
    Args:
124
125
        new_header: astropy header, target
        old_header: astropy header, source
Carl Schaffer's avatar
Carl Schaffer committed
126
127
128
129
130
131
        patches: ordered dict of patch configurations

    Returns:
        newheader: modified astropy header

    """
Sebastian Hoch's avatar
Sebastian Hoch committed
132
133
    for patch in patches.items():
        try:
134
            new_header = apply_patch(new_header, old_header, patch)
Sebastian Hoch's avatar
Sebastian Hoch committed
135
        except:
Carl Schaffer's avatar
Carl Schaffer committed
136
            print("Could not apply patch {0}".format(patch))
Sebastian Hoch's avatar
Sebastian Hoch committed
137
            pass
138

139
    return new_header
140

Sebastian Hoch's avatar
Sebastian Hoch committed
141

142
def apply_patch(new_header, old_header, patch):
Carl Schaffer's avatar
Carl Schaffer committed
143
144
145
146
147
148
149
150
151
    """
    Apply Patch object to header
    A patch is a dictionary with 4 keys:
    "KIND": specify kind of transformation to perform
    "OLD": old keyword
    "NEW": new keyword
    "FUNCTION": expression to evaluate

    Args:
152
153
        new_header: astropy header, target
        old_header: astropy header, source
Carl Schaffer's avatar
Carl Schaffer committed
154
155
156
157
158
159
        patches: ordered dict of patch configurations

    Returns:

    """

Sebastian Hoch's avatar
Sebastian Hoch committed
160
    try:
Carl Schaffer's avatar
Carl Schaffer committed
161
162
163
164
165
166
167
168
        kind = patch[1]["KIND"]
        old = patch[1]["OLD"]
        new = patch[1]["NEW"]
        function = patch[1]["FUNCTION"]

        if kind == "UPDATEVALUE":
            function = function.replace("OLD", "'" + old + "'")
            function = function.replace("NEW", "'" + new + "'")
Sebastian Hoch's avatar
Sebastian Hoch committed
169
        else:
Carl Schaffer's avatar
Carl Schaffer committed
170
            print("Unknown Kind: {0}".format(kind))
Sebastian Hoch's avatar
Sebastian Hoch committed
171
172
        exec(function)
    except KeyError:
Carl Schaffer's avatar
Carl Schaffer committed
173
        print(f"Error Executing {function}")
Sebastian Hoch's avatar
Sebastian Hoch committed
174
    except:
Carl Schaffer's avatar
Carl Schaffer committed
175
        print("Could not evaluate patch: {0}".format(function))
176
        pass
177
    return new_header
Sebastian Hoch's avatar
Sebastian Hoch committed
178

179

Carl Schaffer's avatar
Carl Schaffer committed
180
181
182
183
184
185
186
187
188
189
def get_coords(filename):
    """
    Retrieve coordinates from observation list
    Args:
        filename: bbi filename

    Returns:
        hpc_x, hpc_y: helioprojective coordinates in arcseconds
    """
    res = get_obs_info(filename)
190
    hpc_x, hpc_y = res["HPC-X"], res["HPC-Y"]
Carl Schaffer's avatar
Carl Schaffer committed
191
192
193
194
195
196
197
198
199
200
201
202
203
    return hpc_x, hpc_y


def get_pointing_id(filename):
    """
    Match bbi filename to observation from list of known observations
    Args:
        filename: bbi filename

    Returns: point_id: string like '003'

    """
    res = get_obs_info(filename)
204
    point_id = res["POINT_ID"]
Carl Schaffer's avatar
Carl Schaffer committed
205
206
207
208

    return point_id


Carl Schaffer's avatar
Carl Schaffer committed
209
210
211
212
213
214
215
216
217
218
def insert_wcs_kws(header, filename):
    """
    Add WCS coordinate information Keywords into existing header
    Args:
        header: existing astropy Header instance to be updated
        filename: filename from where to extract the coordinate information

    Returns: new_header: updated header

    """
219
220
221
    # add wcs kws
    x, y = get_coords(filename)
    data = fits.getdata(filename)
Carl Schaffer's avatar
Carl Schaffer committed
222
    center = [int(ax / 2.0) for ax in data.shape]
223
    center.reverse()  # fix x and y order
Carl Schaffer's avatar
Carl Schaffer committed
224
    npix_x, npix_y = data.shape[1], data.shape[0]  # fix x and y order
225

Carl Schaffer's avatar
Carl Schaffer committed
226
    fov_x, fov_y = settings.fov_x, settings.fov_y
227
228
229
230
231
232
233
234
235
236

    step_width_x = fov_x / npix_x
    step_width_y = fov_y / npix_y

    wcs_values = dict(
        CTYPE1="HPLN-TAN",
        CUNIT1="arcsec",
        CRPIX1=center[0],
        CRVAL1=float(x),
        CDELT1=step_width_x,
Carl Schaffer's avatar
Carl Schaffer committed
237
238
239
        CSYER1=fov_x
        * 0.25
        * step_width_x,  # assume coordinate error of 25 percent of image
240
241
242
243
244
        CTYPE2="HPLT-TAN",
        CUNIT2="arcsec",
        CRPIX2=center[1],
        CRVAL2=float(y),
        CDELT2=step_width_y,
Carl Schaffer's avatar
Carl Schaffer committed
245
246
247
        CSYER2=fov_y
        * 0.25
        * step_width_y,  # assume coordinate error of 25 percent of image
Carl Schaffer's avatar
Carl Schaffer committed
248
249
250
        # Compatibility for SunPy Maps
        HGLN_OBS=0,
        HGLT_OBS=0,
Carl Schaffer's avatar
Carl Schaffer committed
251
        DSUN_OBS=earth_distance(header["DATE-BEG"]).meter,
252
    )
253
    wcs_values["DATE-OBS"] = str(date_from_fn(filename))
Carl Schaffer's avatar
Carl Schaffer committed
254

255
    for key, val in wcs_values.items():
Carl Schaffer's avatar
Carl Schaffer committed
256
257
        header[key] = val
    return header
258
259


260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
def insert_separators(new_header):
    """
    Insert separators into a given header, parses the separator list from header YAML

    Args:
        new_header: header

    Returns:
        new_header: modified header
    """

    sep_list = get_separators(settings.header_template)
    for sep in sep_list:
        for sep_contents, sep_config in sep.items():
            target = sep_config["after"]
Carl Schaffer's avatar
Carl Schaffer committed
275
            new_header.add_comment("-" * 70, after=target)
276
            if "additional_content" in sep_config:
Carl Schaffer's avatar
Carl Schaffer committed
277
                content = sep_config["additional_content"]
278

Carl Schaffer's avatar
Carl Schaffer committed
279
                lines = content.split("\n")
280
281
282
283
                lines.reverse()
                for line in lines:
                    new_header.add_comment(line, after=target)

Carl Schaffer's avatar
Carl Schaffer committed
284
285
            new_header.add_comment(sep_contents.center(70, "-"), after=target)
            new_header.add_comment("-" * 70, after=target)
286

287

Carl Schaffer's avatar
Carl Schaffer committed
288
def convert_bbi_header(filename):
Carl Schaffer's avatar
Carl Schaffer committed
289
290
291
292
293
294
295
296
    """
    Main Header conversion function
    Args:
        filename: path to bbi file

    Returns:

    """
Sebastian Hoch's avatar
Sebastian Hoch committed
297
    header = fits.getheader(filename)
Carl Schaffer's avatar
Carl Schaffer committed
298
    bbi_header = get_generic_header(settings.header_template, "reconstruction")
299
    bbi_fn_template = settings.fn_template
Carl Schaffer's avatar
Carl Schaffer committed
300
    patches = get_generic_patches(settings.header_patches_template, "reconstruction")
Sebastian Hoch's avatar
Sebastian Hoch committed
301

Carl Schaffer's avatar
Carl Schaffer committed
302
303
    # Check how much changes are necessary and inform user
    diff = fits.HeaderDiff(header, bbi_header)
Sebastian Hoch's avatar
Sebastian Hoch committed
304
305

    if diff.identical:
Carl Schaffer's avatar
Carl Schaffer committed
306
        print("Header is identical to generic header")
Sebastian Hoch's avatar
Sebastian Hoch committed
307
308
        return header

Carl Schaffer's avatar
Carl Schaffer committed
309
    new_header = bbi_header.copy()
Sebastian Hoch's avatar
Sebastian Hoch committed
310

Carl Schaffer's avatar
Carl Schaffer committed
311
    # add missing cards to new header
Sebastian Hoch's avatar
Sebastian Hoch committed
312
    for card in diff.diff_keyword_values.items():
313
314
        key = card[0]
        value = card[1][0][0]
Carl Schaffer's avatar
Carl Schaffer committed
315
        new_header[key] = value
Sebastian Hoch's avatar
Sebastian Hoch committed
316
317

    # Apply patches to header
318
    if patches is not None:
Carl Schaffer's avatar
Carl Schaffer committed
319
        new_header = apply_patches(new_header, header, patches)
Sebastian Hoch's avatar
Sebastian Hoch committed
320

321
    date_section = re.search(r"\d{8}[_\-]\d{6}", basename(filename)).group(0)
Carl Schaffer's avatar
Carl Schaffer committed
322
323
    timestamp = datetime.datetime.strptime(date_section, "%Y%m%d-%H%M%S")
    wavelength = new_header["WAVELNTH"]
324
    pointing_id = get_pointing_id(filename)
325
    new_header["POINT_ID"] = pointing_id
326
    new_header["EXTNAME"] = "_".join([new_header["POINT_ID"], f"WL{wavelength:03.0f}"])
327
    new_header["FILTER"] = get_filter(wavelength)
Sebastian Hoch's avatar
Sebastian Hoch committed
328

Carl Schaffer's avatar
Carl Schaffer committed
329
330
331
    file_mode = (
        "reco_bb" if (("Zyla-1" in filename) or ("narrow" in filename)) else "reco_nb"
    )
332
333
334
335
336
337

    fits_filename = bbi_fn_template.format(
        timestamp.strftime("%Y%m%d-%H%M%S"),
        pointing_id.split("_")[-1],
        int(wavelength),
        file_mode,
Carl Schaffer's avatar
Carl Schaffer committed
338
        "fits",
339
340
    )
    new_filename = fits_filename
341

342
343
    new_header["FILENAME"] = new_filename
    new_header["CREATOR"] = basename(__file__)
Sebastian Hoch's avatar
Sebastian Hoch committed
344

345
346
    # get additional plots and info from reconstruction results
    try:
347
348
        raise KeyError
    # speck = pipeline.speckle(filename + '.log')
349
    except:
Carl Schaffer's avatar
Carl Schaffer committed
350
        print("Could not find metadata")
351
352
353
354
        speck = None
    if speck is not None:
        estimated_alpha = speck.get_sub_alpha()
        estimated_alpha = estimated_alpha.mean((0, 1))[0]
Carl Schaffer's avatar
Carl Schaffer committed
355
356
357
358
        estimated_fried = (
            estimated_alpha * float(speck.parameter["telescopediameter_mm"]) / 10.0
        )
        new_header["ATMOS_R0"] = float(f"{estimated_fried:03.3f}")
359

360
    insert_wcs_kws(new_header, filename)
Carl Schaffer's avatar
Carl Schaffer committed
361
    add_history(new_header, __file__, "header_fix")
362
    insert_separators(new_header)
Carl Schaffer's avatar
Carl Schaffer committed
363
    return new_header
364
365


Carl Schaffer's avatar
Carl Schaffer committed
366
if __name__ == "__main__":
367
368
369
370
    infile = "/archive/bbi/new/20170519/Zyla1/RC/RC_data_andorZyla-1_20170519-123314144_HAlpha_656.fits"
    header_template = "/dat/schaffer/projects/bbi_header/generic_bbi_header.yaml"
    outpath = "/tmp/"
    patches_fn = "/dat/schaffer/projects/bbi_header/bbi_patches.yaml"
Carl Schaffer's avatar
Carl Schaffer committed
371
    res = convert_bbi_header(infile)
372
373
374
375
376
377
378
379
380


def run_bbi_header_conversion(args):
    parser = BasicArgParser()
    parser.add_argument(
        "-o",
        "--folder_out",
        dest="folder_out",
        default="/dat/sdc/bbi/",
Carl Schaffer's avatar
Carl Schaffer committed
381
        help="Output folder, all output files will be stored here, (default: %(default)s)",
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
    )
    args = parser.parse_args(args)
    in_folder = args.folder_in
    out_folder = args.folder_out

    assert isdir(in_folder)
    assert isdir(out_folder)

    patterns = [
        "RC*.fits",
        "RC/RC*.fits",
        "Zyla1/RC/RC*.fits",
    ]
    files = []
    i = 0
    for p in patterns:
        files = glob(join(in_folder, p))
        if files:
            break
    else:
        print(f"Could not find any BBI files needing conversion in {in_folder}")
        exit()

    message = f"Converting headers.\n\tSource:{in_folder}\n\tTarget:{out_folder}"
    print(message)
    progress_bar = tqdm.tqdm(files)
    for infile in progress_bar:
        progress_bar.set_description(basename(infile))
        header_template = settings.header_template
        patches_fn = settings.header_patches_template
        new_header = convert_bbi_header(
            infile,
        )
415
        wl = get_bbi_wl(new_header["FILENAME"])
416
417
418
        out_dir = join(out_folder, new_header["POINT_ID"] + f"_WL{wl}")

        setup_outdir(out_dir)
Carl Schaffer's avatar
Carl Schaffer committed
419
        out_file = join(out_dir, new_header["FILENAME"])
420
421
422
423
        data = fits.getdata(infile)

        flipped = np.flip(data, axis=0)

Carl Schaffer's avatar
Carl Schaffer committed
424
425
426
427
428
429
430
431
        fits.writeto(
            out_file,
            flipped,
            new_header,
            overwrite=True,
            checksum=True,
            output_verify="warn",
        )