File: wrappers.py

package info (click to toggle)
montage-wrapper 0.9.9-4
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, bullseye, sid
  • size: 1,092 kB
  • sloc: python: 8,325; makefile: 116; ansic: 88
file content (613 lines) | stat: -rw-r--r-- 21,399 bytes parent folder | download | duplicates (2)
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
import atexit
import os
import glob
import shutil as sh
import warnings
import tempfile

import numpy

from astropy.io import fits
from astropy import log

from . import commands as m
from .status import MontageError


def _finalize(cleanup, work_dir, silence=False):
    if cleanup:
        # Deleting work directory
        if not silence:
            log.info("Deleting work directory %s" % work_dir)
        sh.rmtree(work_dir)
    else:
        # Leave work directory as it is
        if not silence:
            log.info("Leaving work directory %s" % work_dir)


def reproject_hdu(in_hdu, **kwargs):
    '''
    Reproject an image (HDU version)

    Parameters
    ----------
    in_hdu : :class:`~astropy.io.fits.PrimaryHDU` or :class:`~astropy.io.fits.ImageHDU`
        Input FITS HDU to be reprojected.

    Notes
    -----
    Additional keyword arguments are passed to :func:`~montage_wrapper.wrappers.reproject`
    '''

    cleanup = kwargs.get('cleanup', True)
    silent_cleanup = kwargs.get('silent_cleanup', True)

    # Make work directory
    work_dir = tempfile.mkdtemp()

    try:
        in_image = os.path.join(work_dir, 'in.fits')
        out_image = os.path.join(work_dir, 'out.fits')

        fits.writeto(in_image, in_hdu.data, in_hdu.header)

        reproject(in_image, out_image, **kwargs)

        out_hdu = fits.open(out_image)[0]

        return out_hdu
    finally:
        _finalize(cleanup, work_dir, silent_cleanup)


def reproject_cube(in_image, out_image, header=None, bitpix=None,
                   north_aligned=False, system=None, equinox=None, factor=None,
                   common=False, cleanup=True, clobber=False,
                   silent_cleanup=True, hdu=0):
    '''
    Cube reprojection routine.

    If one input/output image is specified, and the header argument is set,
    the routine is equivalent to using mProject or mProjectPP. If header= is
    not set, a new header is computed by taking into account the
    north_aligned, system, and equinox arguments (if set).

    Parameters
    ----------

    in_image : str
        Path of input FITS file to be reprojected.

    out_image : str
        Path of output FITS file to be created.

    header : str, optional
        Path to the header file to use for re-projection.

    bitpix : int, optional
        BITPIX value for the ouput FITS file (default is -64). Possible
        values are: 8 (character or unsigned binary integer), 16 (16-bit
        integer), 32 (32-bit integer), -32 (single precision floating
        point), -64 (double precision floating point).

    north_aligned : bool, optional
        Align the pixel y-axis with North

    system : str, optional
        Specifies the coordinate system
        Possible values are: 'EQUJ', 'EQUB', 'ECLJ', 'ECLB', 'GAL', 'SGAL'

    equinox : str, optional
        If a coordinate system is specified, the equinox can also be given
        in the form 'YYYY'. Default is 'J2000'.

    factor : float, optional
        Drizzle factor (see `~montage_wrapper.commands.mProject`)

    clobber : bool, optional
        Overwrite the data cube if it already exists?

    silent_cleanup : bool, optional
        Hide messages related to tmp directory removal (there will be one
        for each plane of the cube if set to False)

    hdu: int
        Defaults to zero.  Selects which HDU to use from the FITS file.
    '''

    if header:
        if north_aligned or system or equinox:
            warnings.warn("header= is set, so north_aligned=, system=, and equinox= will be ignored")

    # Find path to input and output file
    in_image = os.path.abspath(in_image)
    out_image = os.path.abspath(out_image)

    if os.path.exists(out_image) and not clobber:
        raise IOError("File '%s' already exists and clobber=False." % out_image)

    # Make work directory
    work_dir = tempfile.mkdtemp()

    # Set paths

    raw_dir = os.path.join(work_dir, 'raw')
    final_dir = os.path.join(work_dir, 'final')

    if header:
        header_hdr = os.path.abspath(header)
    else:
        header_hdr = os.path.join(work_dir, 'header.hdr')

    images_raw_tbl = os.path.join(work_dir, 'images_raw.tbl')
    images_tmp_tbl = os.path.join(work_dir, 'images_tmp.tbl')

    # Create raw directory
    os.mkdir(raw_dir)
    os.mkdir(final_dir)

    # Make new header
    if not header:
        m.mMakeHdr(images_raw_tbl, header_hdr, north_aligned=north_aligned,
                   system=system, equinox=equinox)

    cubefile = fits.open(in_image)
    if len(cubefile[hdu].data.shape) != 3 or cubefile[hdu].header.get('NAXIS') != 3:
        raise Exception("Cube file must have 3 dimensions")

    # a temporary HDU that will be used to hold different data each time
    # and reproject each plane separately
    planefile = fits.PrimaryHDU(data=cubefile[hdu].data[0, :, :],
                                header=cubefile[hdu].header)

    # generate a blank HDU to store the eventual projected cube

    # first must remove END card from .hdr file
    header_temp = header_hdr.replace(".hdr", "_tmp.hdr")
    headerlines = open(header_hdr, 'r').readlines()[:-1]
    w = open(header_temp, 'w')
    w.writelines([line for line in headerlines])
    w.close()

    # when creating the new header, allow the 3rd axis to be
    # set by the input data cube
    newheader = fits.Header()
    newheader.fromTxtFile(header_temp)
    blank_data = numpy.zeros(
            [cubefile[hdu].header.get('NAXIS3'),
             newheader.get('NAXIS2'),
             newheader.get('NAXIS1')]
    )
    newcube = fits.PrimaryHDU(data=blank_data, header=newheader)

    for ii, plane in enumerate(cubefile[hdu].data):

        os.mkdir(os.path.join(final_dir, '%i' % ii))

        # reset the data plane within the temporary HDU
        planefile.data = plane

        # reproject the individual plane - exact size MUST be specified so that the
        # data can be put into the specified cube
        reprojected = reproject_hdu(planefile, header=header_hdr,
                                    exact_size=True, factor=factor, bitpix=bitpix,
                                    silent_cleanup=silent_cleanup)

        newcube.data[ii, :, :] = reprojected.data

    newcube.writeto(out_image, clobber=clobber)

    _finalize(cleanup, work_dir)

    return


def mProject_auto(*args, **kwargs):
    '''
    Run mProject, automatically selecting whether to run mProject or
    mProjectPP if possible (fast plane-to-plane projection). For details on
    required and optional arguments, see help(mProject).
    '''
    try:
        m.mProjectPP(*args, **kwargs)
    except MontageError:
        m.mProject(*args, **kwargs)


def reproject(in_images, out_images, header=None, bitpix=None,
              north_aligned=False, system=None, equinox=None, factor=None, common=False,
              exact_size=False, hdu=None, cleanup=True, silent_cleanup=False):
    '''
    General-purpose reprojection routine.

    If one input/output image is specified, and the header argument is set,
    the routine is equivalent to using mProject or mProjectPP. If header= is
    not set, a new header is computed by taking into account the
    north_aligned, system, and equinox arguments (if set).

    If tuples of input/output images are specified, the tuples need to have
    the same number of elements. If header= is specified, all images are
    projected to this common projection. If header= is not specified, then a
    new header is computed a new header is computed by taking into account the
    north_aligned, system, and equinox arguments (if set). If common=False,
    then a header is computed for each individual image, whereas if
    common=True, an optimal header is computed for all images.

    Parameters
    ----------

    in_images : str or list
        Path(s) of input FITS file(s) to be reprojected.

    out_images : str or list
        Path(s) of output FITS file(s) to be created.

    header : str, optional
        Path to the header file to use for re-projection.

    bitpix : int, optional
        BITPIX value for the ouput FITS file (default is -64). Possible
        values are: 8 (character or unsigned binary integer), 16 (16-bit
        integer), 32 (32-bit integer), -32 (single precision floating
        point), -64 (double precision floating point).

    north_aligned : bool, optional
        Align the pixel y-axis with North

    system : str, optional
        Specifies the coordinate system
        Possible values are: 'EQUJ', 'EQUB', 'ECLJ', 'ECLB', 'GAL', 'SGAL'

    equinox : str, optional
        If a coordinate system is specified, the equinox can also be given
        in the form 'YYYY'. Default is 'J2000'.

    factor : float, optional
        Drizzle factor (see `~montage_wrapper.commands.mProject`)

    exact_size : bool, optional
        Whether to reproject the image(s) to the exact header specified
        (i.e. whether cropping is unacceptable).

    hdu : int, optional
        The HDU to use in the file(s)

    silent_cleanup : bool, optional
        Hide messages related to tmp directory removal

    common : str, optional
        Compute a common optimal header for all images (only used if
        header=None and if multiple files are being reprojected)
    '''

    if type(in_images) == str and type(out_images) == str:
        in_images = (in_images, )
        out_images = (out_images, )
    elif type(in_images) in [tuple, list] and type(out_images) in [tuple, list]:
        pass
    else:
        raise Exception("Inconsistent type for in_images (%s) and out_images (%s)" % (type(in_images), type(out_images)))

    if header:
        if north_aligned or system or equinox:
            warnings.warn("header= is set, so north_aligned=, system=, and equinox= will be ignored")

    if common and len(in_images) == 1:
        warnings.warn("only one image is being reprojected, so common= will be ignored")

    # Find path to input and output file
    in_images = [os.path.abspath(in_image) for in_image in in_images]
    out_images = [os.path.abspath(out_image) for out_image in out_images]

    if len(in_images) > 1 and not header and not common:
        for i, in_image in enumerate(in_images):
            reproject(in_images[i], out_images[i], bitpix=bitpix,
                      north_aligned=north_aligned, system=system,
                      equinox=equinox, factor=factor,
                      exact_size=exact_size, cleanup=cleanup,
                      silent_cleanup=silent_cleanup)
        return

    # Make work directory
    work_dir = tempfile.mkdtemp()

    # Set paths

    raw_dir = os.path.join(work_dir, 'raw')
    final_dir = os.path.join(work_dir, 'final')

    if header:
        header_hdr = os.path.abspath(header)
    else:
        header_hdr = os.path.join(work_dir, 'header.hdr')

    images_raw_tbl = os.path.join(work_dir, 'images_raw.tbl')
    images_tmp_tbl = os.path.join(work_dir, 'images_tmp.tbl')

    # Create raw directory
    os.mkdir(raw_dir)
    os.mkdir(final_dir)

    # Link to images
    for i, in_image in enumerate(in_images):
        os.mkdir(os.path.join(raw_dir, '%i' % i))
        os.symlink(in_image, os.path.join(raw_dir, '%i' % i, 'image.fits'))

    # Make image table
    m.mImgtbl(raw_dir, images_raw_tbl, corners=True, recursive=True)

    # Make new header
    if not header:
        m.mMakeHdr(images_raw_tbl, header_hdr, north_aligned=north_aligned,
                   system=system, equinox=equinox)

    for i, in_image in enumerate(in_images):

        os.mkdir(os.path.join(final_dir, '%i' % i))

        mProject_auto(in_images[i], os.path.join(final_dir, '%i' % i, 'image_tmp.fits'),
                      header_hdr, hdu=hdu, factor=factor)

        if exact_size:
            m.mImgtbl(os.path.join(final_dir, '%i' % i), images_tmp_tbl, corners=True)
            m.mAdd(images_tmp_tbl, header_hdr,
                   os.path.join(final_dir, '%i' % i, 'image.fits'),
                   img_dir=os.path.join(final_dir, '%i' % i), exact=True)
        else:
            os.symlink(os.path.join(final_dir, '%i' % i, 'image_tmp.fits'),
                       os.path.join(final_dir, '%i' % i, 'image.fits'))

        m.mConvert(os.path.join(final_dir, '%i' % i, 'image.fits'), out_images[i],
                   bitpix=bitpix)

    _finalize(cleanup, work_dir, silence=silent_cleanup)

    return


def mosaic(input_dir, output_dir, header=None, image_table=None, mpi=False,
           n_proc=8, background_match=False, imglist=None, combine="mean",
           exact_size=False, cleanup=True, bitpix=-32, level_only=True,
           work_dir=None, background_n_iter=None, subset_fast=False,
           hdu=None):
    """
    Combine FITS files into a mosaic

    Parameters
    ----------

    input_dir : str
        The directory containing the input FITS files

    output_dir : str
        The output directory to create

    header : str, optional
        The header to project to. If this is not specified, then an optimal
        header is chosen.

    image_table : str, optional
        The table file containing the list of input images. This can be
        specified to avoid recomputing it every time a mosaic is made from
        the same set of input files.

    mpi : bool, optional
        Whether to use MPI whenever possible (requires the MPI-enabled
        Montage binaries to be installed).

    n_proc : int, optional
        The number of processes to use if `mpi` is set to `True`

    background_match : bool, optional
        Whether to include a background-matching step

    imglist : str, optional
        A list of images to use (useful if not all the files inside
        `input_dir` should be combined).

    combine : str, optional
        How to combine the images - this should be one of ``'mean'``,
        ``'median'``, or ``'count'``.

    exact_size : bool, optional
        Whether the output mosaic should match the input header exactly, or
        whether the mosaic should be trimmed if possible.

    cleanup : bool, optional
        Whether to remove any temporary directories used for mosaicking

    bitpix : int, optional
        BITPIX value for the ouput FITS file (default is -32). Possible
        values are: 8 (character or unsigned binary integer), 16 (16-bit
        integer), 32 (32-bit integer), -32 (single precision floating
        point), -64 (double precision floating point).

    level_only : bool, optional
        When doing background matching, whether to only allow changes in the
        level of frames, not the slope.

    work_dir : str, optional
        The temporary directory to use for the mosaicking. By default, this
        is a temporary directory in a system-defined location, but it can be
        useful to specify this manually for debugging purposes.

    background_n_iter : int, optional
        Number of iterations for the background matching (defaults to 5000). Can be
        between 1 and 32767.

    subset_fast : bool, optional
        Whether to use the fast mode for mSubset

    hdu: int, optional
        Which HDU to use when mosaicing
    """

    if not combine in ['mean', 'median', 'count']:
        raise Exception("combine should be one of mean/median/count")

    # Check that there are files in the input directory
    if len(glob.glob(os.path.join(input_dir, '*'))) == 0:
        raise Exception("No files in input directory")

    # Find path to input and output directory
    input_dir = os.path.abspath(input_dir)
    output_dir = os.path.abspath(output_dir)

    # Make work directory
    if work_dir:
        work_dir = os.path.abspath(work_dir)
        if os.path.exists(work_dir):
            raise Exception("Work directory already exists")
        os.mkdir(work_dir)
    else:
        # Use a temporary directory such as 'montage_mosaic_ngc2264_EdJyrj',
        # where 'ngc2264' would be here determined by the input directory.
        basename = os.path.basename(os.path.normpath(input_dir))
        prefix = "montage_mosaic_{0}_".format(basename)
        work_dir = tempfile.mkdtemp(prefix=prefix)

    # Make sure the working directory is cleaned up even when something goes
    # wrong (or Ctrl-C is pressed, raising the KeyboardInterrupt exception)
    @atexit.register
    def cleanup_workdir():
        """ Run _finalize() if necessary """
        if os.path.exists(work_dir):
            _finalize(cleanup, work_dir)

    images_raw_all_tbl = os.path.join(work_dir, 'images_raw_all.tbl')
    images_raw_tbl = os.path.join(work_dir, 'images_raw.tbl')
    images_projected_tbl = os.path.join(work_dir, 'images_projected.tbl')
    images_corrected_tbl = os.path.join(work_dir, 'images_corrected.tbl')
    corrections_tbl = os.path.join(work_dir, 'corrections.tbl')
    diffs_tbl = os.path.join(work_dir, 'diffs.tbl')
    stats_tbl = os.path.join(work_dir, 'stats.tbl')
    fits_tbl = os.path.join(work_dir, 'fits.tbl')

    raw_dir = os.path.join(work_dir, 'raw')
    projected_dir = os.path.join(work_dir, 'projected')
    corrected_dir = os.path.join(work_dir, 'corrected')
    diffs_dir = os.path.join(work_dir, 'diffs')

    header_hdr = os.path.join(work_dir, 'header.hdr')

    # Find path to header file if specified
    if header is not None:
        header = os.path.abspath(header)

    # Find path to image table if specified
    if image_table is not None:
        image_table = os.path.abspath(image_table)

    # Find path to image list if specified
    if imglist:
        imglist = os.path.abspath(imglist)

    # Create output dir
    if os.path.exists(output_dir):
        raise IOError("Output directory already exists")
    else:
        os.mkdir(output_dir)

    # Create symbolic links
    os.symlink(input_dir, raw_dir)

    if header:
        os.symlink(header, header_hdr)

    # Create temporary directories for Montage
    os.mkdir(projected_dir)
    if background_match:
        os.mkdir(diffs_dir)
        os.mkdir(corrected_dir)

    # List frames to mosaic
    if image_table is None:
        log.info("Listing raw frames")
        m.mImgtbl(raw_dir, images_raw_all_tbl, img_list=imglist, corners=True)
    else:
        sh.copy2(image_table, images_raw_all_tbl)

    # Compute header if needed
    if not header:
        log.info("Computing optimal header")
        m.mMakeHdr(images_raw_all_tbl, header_hdr)
        images_raw_tbl = images_raw_all_tbl
    else:
        log.info("Checking for coverage")
        s = m.mSubset(images_raw_all_tbl, header_hdr, images_raw_tbl, fast_mode=subset_fast)
        if s.nmatches == 0:
            raise MontageError("No images overlap with the requested header")

    # Projecting raw frames
    log.info("Projecting raw frames")

    if hdu is not None:
        from astropy.table import Table
        table = Table.read(images_raw_tbl, format='ascii.ipac')
        table_filtered = table[table['hdu'] == hdu]
        table_filtered.write(images_raw_tbl, format='ascii.ipac')

    m.mProjExec(images_raw_tbl, header_hdr, projected_dir, stats_tbl,
                raw_dir=raw_dir, mpi=mpi, n_proc=n_proc, exact=exact_size)

    # List projected frames
    s = m.mImgtbl(projected_dir, images_projected_tbl)
    if s.count == 0:
        raise MontageError("No images were successfully projected")

    if background_match:
        log.info("Determining overlaps")
        s = m.mOverlaps(images_projected_tbl, diffs_tbl)
        if s.count == 0:
            log.info("No overlapping frames, backgrounds will not be adjusted")
            background_match = False

    if background_match:

        # Modeling background

        log.info("Modeling background")
        m.mDiffExec(diffs_tbl, header_hdr, diffs_dir, proj_dir=projected_dir,
                    mpi=mpi, n_proc=n_proc)
        m.mFitExec(diffs_tbl, fits_tbl, diffs_dir, mpi=mpi, n_proc=n_proc)
        m.mBgModel(images_projected_tbl, fits_tbl, corrections_tbl,
                   n_iter=background_n_iter, level_only=level_only)

        # Matching background
        log.info("Matching background")
        m.mBgExec(images_projected_tbl, corrections_tbl, corrected_dir,
                  proj_dir=projected_dir, mpi=mpi, n_proc=n_proc)
        sh.copy(corrections_tbl, output_dir)

        # Mosaicking frames
        log.info("Mosaicking frames")

        m.mImgtbl(corrected_dir, images_corrected_tbl)
        m.mAdd(images_corrected_tbl, header_hdr,
               os.path.join(output_dir, 'mosaic64.fits'),
               img_dir=corrected_dir, type=combine, exact=exact_size)
        sh.copy(images_projected_tbl, output_dir)
        sh.copy(images_corrected_tbl, output_dir)

    else:

        # Mosaicking frames
        log.info("Mosaicking frames")

        m.mAdd(images_projected_tbl, header_hdr,
               os.path.join(output_dir, 'mosaic64.fits'),
               img_dir=projected_dir, type=combine, exact=exact_size)
        sh.copy(images_projected_tbl, output_dir)

    m.mConvert(os.path.join(output_dir, 'mosaic64.fits'),
               os.path.join(output_dir, 'mosaic.fits'),
               bitpix=bitpix)
    m.mConvert(os.path.join(output_dir, 'mosaic64_area.fits'),
               os.path.join(output_dir, 'mosaic_area.fits'),
               bitpix=bitpix)

    os.remove(os.path.join(output_dir, "mosaic64.fits"))
    os.remove(os.path.join(output_dir, "mosaic64_area.fits"))

    _finalize(cleanup, work_dir)