<filename>spexxy/tools/spectrum/extract.py<gh_stars>1-10
import glob
import logging
import argparse
import os
from typing import List

from spexxy.data import FitsSpectrum, SpectrumFitsHDU

log = logging.getLogger(__name__)


def add_parser(subparsers):
    # init parser
    parser = subparsers.add_parser('extract', help='Extract wavelength range from spectrum',
                                   formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('start', type=float, help='Start of wavelength range')
    parser.add_argument('end', type=float, help='End of wavelength range')
    parser.add_argument('input', type=str, help='Input spectrum', nargs='+')
    parser.add_argument('-p', '--prefix', type=str, help='Output file prefix', default='extracted_')

    # argparse wrapper for create_grid
    parser.set_defaults(func=lambda args: extract_spectra(**vars(args)))


def extract_spectra(input: List[str], prefix: str, start: float, end: float, **kwargs):
    # get list of filenames
    files = []
    for filename in input:
        if '*' in input or '?' in input:
            files.extend(glob.glob(filename))
        else:
            files.append(filename)

    # make unique
    files = sorted(list(set(files)))

    # loop files
    for i, filename in enumerate(files, 1):
        print('(%d/%d) %s' % (i, len(files), filename))
        extract_spectrum(filename, prefix + os.path.basename(filename), start, end)


def extract_spectrum(input: str, output: str, start: float, end: float):
    """Extracts the given wavelength range from input and store it as output.

    Args:
        input: Input filename
        output: Output filename
        start: Wavelength start
        end: Wavelength end
    """

    # load spectrum
    with FitsSpectrum(input) as fs_in:
        # open spectrum to write
        with FitsSpectrum(output, 'w') as fs_out:
            # copy extracted spectrum
            fs_out.spectrum = fs_in.spectrum.extract(start, end)

            # loop all extensions:
            for ext in fs_in.hdu_names():
                # skip no name
                if ext.strip() == '':
                    continue

                # try to get hdu
                try:
                    # get hdu
                    hdu = fs_in[ext]

                    # need to create new hdu with primary=False, since that gets lost on extract
                    fs_out[ext] = SpectrumFitsHDU(spec=hdu.extract(start, end), primary=False)

                except ValueError:
                    pass
<filename>spexxy/tools/spectrum/extract.py<gh_stars>1-10
import glob
import logging
import argparse
import os
from typing import List

from spexxy.data import FitsSpectrum, SpectrumFitsHDU

log = logging.getLogger(__name__)


def add_parser(subparsers):
    # init parser
    parser = subparsers.add_parser('extract', help='Extract wavelength range from spectrum',
                                   formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('start', type=float, help='Start of wavelength range')
    parser.add_argument('end', type=float, help='End of wavelength range')
    parser.add_argument('input', type=str, help='Input spectrum', nargs='+')
    parser.add_argument('-p', '--prefix', type=str, help='Output file prefix', default='extracted_')

    # argparse wrapper for create_grid
    parser.set_defaults(func=lambda args: extract_spectra(**vars(args)))


def extract_spectra(input: List[str], prefix: str, start: float, end: float, **kwargs):
    # get list of filenames
    files = []
    for filename in input:
        if '*' in input or '?' in input:
            files.extend(glob.glob(filename))
        else:
            files.append(filename)

    # make unique
    files = sorted(list(set(files)))

    # loop files
    for i, filename in enumerate(files, 1):
        print('(%d/%d) %s' % (i, len(files), filename))
        extract_spectrum(filename, prefix + os.path.basename(filename), start, end)


def extract_spectrum(input: str, output: str, start: float, end: float):
    """Extracts the given wavelength range from input and store it as output.

    Args:
        input: Input filename
        output: Output filename
        start: Wavelength start
        end: Wavelength end
    """

    # load spectrum
    with FitsSpectrum(input) as fs_in:
        # open spectrum to write
        with FitsSpectrum(output, 'w') as fs_out:
            # copy extracted spectrum
            fs_out.spectrum = fs_in.spectrum.extract(start, end)

            # loop all extensions:
            for ext in fs_in.hdu_names():
                # skip no name
                if ext.strip() == '':
                    continue

                # try to get hdu
                try:
                    # get hdu
                    hdu = fs_in[ext]

                    # need to create new hdu with primary=False, since that gets lost on extract
                    fs_out[ext] = SpectrumFitsHDU(spec=hdu.extract(start, end), primary=False)

                except ValueError:
                    pass
# init parser # argparse wrapper for create_grid # get list of filenames # make unique # loop files Extracts the given wavelength range from input and store it as output.      Args:         input: Input filename         output: Output filename         start: Wavelength start         end: Wavelength end # load spectrum # open spectrum to write # copy extracted spectrum # loop all extensions: # skip no name # try to get hdu # get hdu # need to create new hdu with primary=False, since that gets lost on extract