Plotten eines Intervalls um eine Spektrallinie herum (aus FITS-Datei)

from astropy.io import fits
import numpy as np
import matplotlib.pyplot as plt
import sys

def main():
    with fits.open(sys.argv[1]) as hdul:
        hdul.info()
        hdu = hdul[0]
        hdu.verify('fix')
        data = hdu.data
        header = hdu.header
        plot_intensity(header, data)
        plt.tight_layout()
        plt.savefig(sys.argv[4])


def plot_intensity(header, data):
    wavelength_of_interest = float(sys.argv[2])
    interval_of_interest = float(sys.argv[3])
    start_wavelength = header['CRVAL1']
    increment_wavelength = header['CDELT1']
    pixel_number = len(data)
    wavelengths = start_wavelength + np.arange(0, pixel_number) * increment_wavelength
    condition = np.logical_and(wavelengths >= (wavelength_of_interest - interval_of_interest), wavelengths <= (wavelength_of_interest + interval_of_interest))
    wavelengths = np.extract(condition, wavelengths)
    intensities = np.extract(condition, data)
    plt.gcf().set_size_inches(7,4.5)
    plt.xlabel('Wellenlänge $[\mathrm{\AA}]$')
    plt.ylabel('Intensität')
    plt.plot(wavelengths, intensities)

if __name__ == "__main__":
    main()