Beispiel #1
0
    def __init__(self):

        #
        # Create the base wavelengths and flux
        #

        self.wavelengths_um = np.linspace(0.4, 1.05, 100)

        g1 = models.Gaussian1D(amplitude=2000, mean=0.56, stddev=0.01)
        g2 = models.Gaussian1D(amplitude=500, mean=0.62, stddev=0.02)
        g3 = models.Gaussian1D(amplitude=-400, mean=0.80, stddev=0.02)
        g4 = models.Gaussian1D(amplitude=-350, mean=0.52, stddev=0.01)
        ramp = models.Linear1D(slope=300, intercept=0.0)

        self.base_flux = g1(self.wavelengths_um) + g2(self.wavelengths_um) + \
                         g3(self.wavelengths_um) + g4(self.wavelengths_um) + \
                         ramp(self.wavelengths_um) + 1000

        #
        # Initialize the seed so the random numbers are not quite as random
        #

        np.random.seed(42)

        #
        # Create two spectra with the only difference in the instance of noise
        #

        self._flux_e1 = self.base_flux + 400 * np.random.random(
            self.base_flux.shape)
        self._s1_um_mJy_e1 = Spectrum1D(spectral_axis=self.wavelengths_um *
                                        u.um,
                                        flux=self._flux_e1 * u.mJy)

        self._flux_e2 = self.base_flux + 400 * np.random.random(
            self.base_flux.shape)
        self._s1_um_mJy_e2 = Spectrum1D(spectral_axis=self.wavelengths_um *
                                        u.um,
                                        flux=self._flux_e2 * u.mJy)

        #
        # Create on spectrum with the same flux but in angstrom units
        #

        self.wavelengths_AA = self.wavelengths_um * 10000
        self._s1_AA_mJy_e3 = Spectrum1D(spectral_axis=self.wavelengths_AA *
                                        u.AA,
                                        flux=self._flux_e1 * u.mJy)

        #
        # Create on spectrum with the same flux but in angstrom units and nJy
        #

        self._flux_e4 = (self.base_flux + 400 *
                         np.random.random(self.base_flux.shape)) * 1000000
        self._s1_AA_nJy_e4 = Spectrum1D(spectral_axis=self.wavelengths_AA *
                                        u.AA,
                                        flux=self._flux_e4 * u.nJy)
def _fit_1D(initial_model, spectrum, run_fitter):
    """
    Fits an astropy CompoundModel to a Spectrum1D instance.

    Parameters
    ----------
    spectrum : :class:`specutils.spectrum.Spectrum1D`
        The spectrum to be fitted.
    initial_model : :class: `astropy.modeling.CompoundModel`
        Initial guess for the model to be fitted.
    run_fitter : bool
        When False (the default), the function composes the compound
        model and returns it without fitting.

    Returns
    -------
    :class: `astropy.modeling.CompoundModel`
        The model resulting from the fit.
    :class:`specutils.spectrum.Spectrum1D`
        The realization of the fitted model as a spectrum.

    """
    if run_fitter:
        output_model = fit_lines(spectrum, initial_model)
        output_values = output_model(spectrum.spectral_axis)
    else:
        # Return without fitting.
        output_model = initial_model
        output_values = initial_model(spectrum.spectral_axis)

    # Build return spectrum
    output_spectrum = Spectrum1D(spectral_axis=spectrum.spectral_axis,
                                 flux=output_values)

    return output_model, output_spectrum
Beispiel #3
0
def test_fitting_backend():
    np.random.seed(42)

    x, y = build_spectrum()

    spectrum = Spectrum1D(flux=y*u.Jy, spectral_axis=x*u.um)

    g1f = models.Gaussian1D(0.7*u.Jy, 4.65*u.um, 0.3*u.um, name='g1')
    g2f = models.Gaussian1D(2.0*u.Jy, 5.55*u.um, 0.3*u.um, name='g2')
    g3f = models.Gaussian1D(-2.*u.Jy, 8.15*u.um, 0.2*u.um, name='g3')
    zero_level = models.Const1D(1.*u.Jy, name='const1d')

    model_list = [g1f, g2f, g3f, zero_level]
    expression = "g1 + g2 + g3 + const1d"

    # Returns the initial model
    fm, fitted_spectrum = fb.fit_model_to_spectrum(spectrum, model_list, expression,
                                                   run_fitter=False)

    parameters_expected = np.array([0.7, 4.65, 0.3, 2., 5.55, 0.3, -2.,
                                    8.15, 0.2, 1.])
    assert np.allclose(fm.parameters, parameters_expected, atol=1e-5)

    # Returns the fitted model
    fm, fitted_spectrum = fb.fit_model_to_spectrum(spectrum, model_list, expression,
                                                   run_fitter=True)

    parameters_expected = np.array([1.0104705, 4.58956282, 0.19590464, 2.39892026,
                                    5.49867754, 0.10834472, -1.66902953, 8.19714439,
                                    0.09535613, 3.99125545])
    assert np.allclose(fm.parameters, parameters_expected, atol=1e-5)
Beispiel #4
0
def test_invert():
    sr = (SpectralRegion(0.15 * u.um, 0.2 * u.um) +
          SpectralRegion(0.3 * u.um, 0.4 * u.um) +
          SpectralRegion(0.45 * u.um, 0.6 * u.um) +
          SpectralRegion(0.8 * u.um, 0.9 * u.um) +
          SpectralRegion(1.0 * u.um, 1.2 * u.um) +
          SpectralRegion(1.3 * u.um, 1.5 * u.um))

    sr_inverted_expected = [
        (0.05 * u.um, 0.15 * u.um), (0.2 * u.um, 0.3 * u.um),
        (0.4 * u.um, 0.45 * u.um), (0.6 * u.um, 0.8 * u.um),
        (0.9 * u.um, 1.0 * u.um), (1.2 * u.um, 1.3 * u.um),
        (1.5 * u.um, 3.0 * u.um)
    ]

    # Invert from range.
    sr_inverted = sr.invert(0.05 * u.um, 3 * u.um)

    for ii, expected in enumerate(sr_inverted_expected):
        assert sr_inverted.subregions[ii] == sr_inverted_expected[ii]

    # Invert from spectrum.
    spectrum = Spectrum1D(spectral_axis=np.linspace(0.05, 3, 20) * u.um,
                          flux=np.random.random(20) * u.Jy)
    sr_inverted = sr.invert_from_spectrum(spectrum)
    for ii, expected in enumerate(sr_inverted_expected):
        assert sr_inverted.subregions[ii] == sr_inverted_expected[ii]
    def __call__(self):
        results = {'x': [], 'y': [], 'fitted_model': [], 'fitted_values': []}
        for parameters in self.param_set:
            x = parameters[0]
            y = parameters[1]

            # Calling the Spectrum1D constructor for every spaxel
            # turned out to be less expensive than expected. Experiments
            # show that the cost amounts to a couple percent additional
            # running time in comparison with a version that uses a 3D
            # spectrum as input. Besides, letting an externally-created
            # spectrum reference into the callable somehow prevents it
            # to execute. This behavior was seen also with other functions
            # passed to the callable.
            flux = self.cube[x, y, :]  # transposed!
            sp = Spectrum1D(spectral_axis=self.wave, flux=flux)
            import sys, traceback

            try:
                fitted_model = fit_lines(sp, self.model)
                fitted_values = fitted_model(self.wave)
                results['x'].append(x)
                results['y'].append(y)
                results['fitted_model'].append(fitted_model.unitless_model)
                results['fitted_values'].append(fitted_values.value)
            except:
                pass
                # print("Exception on [{},{}]:".format(x, y), sys.exc_info()[0])
        # print(results)
        return results
Beispiel #6
0
def generic_fits(file_name, **kwargs):
    name = os.path.basename(file_name.rstrip(os.sep)).rsplit('.', 1)[0]

    with fits.open(file_name, **kwargs) as hdulist:
        header = hdulist[0].header
        data3 = hdulist[0].data
        wcs = WCS(header)
        shape = data3.shape

        # take the reference pixel if the pos= was not supplied by the reader
        if 'pos' in kwargs:
            ix = kwargs['pos'][0]
            iy = kwargs['pos'][1]
        else:
            ix = int(wcs.wcs.crpix[0])
            iy = int(wcs.wcs.crpix[1])

        # grab a spectrum from the cube
        if len(shape) == 3:
            data = data3[:, iy, ix]
        elif len(shape) == 4:
            data = data3[:, :, iy, ix].squeeze()
            # make sure this is a 1D array
            # if len(data.shape) != 1:
            #    raise Exception,"not a true cube"
        else:
            print("Unexpected shape", shape)
            #

        # store some meta data
        meta = {'header': header}
        meta['xpos'] = ix
        meta['ypos'] = iy

        # attach units (get it from header['BUNIT'] - what about 'JY/BEAM '
        #  NOTE:  astropy doesn't support beam, but see comments in radio_beam
        data = data * Unit("Jy")

        # now figure out the frequency axis....
        sp_axis = 3
        naxis3 = header['NAXIS%d' % sp_axis]
        cunit3 = wcs.wcs.cunit[sp_axis - 1]
        crval3 = wcs.wcs.crval[sp_axis - 1]
        cdelt3 = wcs.wcs.cdelt[sp_axis - 1]
        crpix3 = wcs.wcs.crpix[sp_axis - 1]

        freqs = np.arange(naxis3) + 1
        freqs = (freqs - crpix3) * cdelt3 + crval3

        freqs = freqs * cunit3

        # should wcs be transformed to a 1D case ?

    return Spectrum1D(flux=data, wcs=wcs, meta=meta, spectral_axis=freqs)
Beispiel #7
0
def test_model_fitting(specviz_gui):
    hub = Hub(workspace=specviz_gui.current_workspace)

    # Generate fake data
    np.random.seed(42)
    g1 = models.Gaussian1D(1, 0, 0.2)
    g2 = models.Gaussian1D(2.5, 0.5, 0.1)

    x = np.linspace(-1, 1, 200)
    y = g1(x) + g2(x) + np.random.normal(0., 0.2, x.shape)

    # Regular fitting
    gg_init = models.Gaussian1D(1.3, 0, 0.1) + models.Gaussian1D(1.8, 0.5, 0.1)
    fitter = fitting.LevMarLSQFitter()
    gg_fit = fitter(gg_init, x, y)

    # SpecViz fitting
    spectral_axis_unit = u.Unit(hub.plot_window.plot_widget.spectral_axis_unit
                                or "")
    data_units = u.Unit(hub.plot_window.plot_widget.data_unit or "")
    s1d = Spectrum1D(flux=y * data_units, spectral_axis=x * spectral_axis_unit)
    hub.workspace.model.add_data(s1d, name="fitting_data")
    model_editor = specviz_gui.current_workspace._plugin_bars['Model Editor']

    model_editor._on_create_new_model()
    model_editor._add_fittable_model(models.Gaussian1D)
    model_editor._add_fittable_model(models.Gaussian1D)

    index = model_editor.data_selection_combo.findText("fitting_data")
    if index >= 0:
        model_editor.data_selection_combo.setCurrentIndex(index)
    value_dict = {
        'Gaussian1D': {
            'amplitude': '1.3',
            'mean': '0',
            'stddev': '0.1'
        },
        'Gaussian1D1': {
            'amplitude': '1.8',
            'mean': '0.5',
            'stddev': '0.1'
        }
    }
    plot_data_item = hub.plot_item
    model_editor_model = plot_data_item.data_item.model_editor_model

    fill_in_models(model_editor_model, value_dict)

    model_editor._on_fit_clicked(eq_pop_up=False)

    model_editor_model = plot_data_item.data_item.model_editor_model
    result = model_editor_model.evaluate()

    np.testing.assert_allclose(result.parameters, gg_fit.parameters)
Beispiel #8
0
    def _on_fit_clicked(self, model_plot_data_item):
        fit_mod = fit_lines(self.hub.data_item.spectrum, result)
        flux = fit_mod(self.hub.data_item.spectrum.spectral_axis)

        new_spec = Spectrum1D(
            flux=flux, spectral_axis=self.hub.data_item.spectrum.spectral_axis)
        # self.hub.model.add_data(new_spec, "Fitted Model Spectrum")

        # Update the stored plot data item object for this model editor model
        # self._model_editor_model.plot_data_item.data_item.set_data(new_spec)

        # Fitted quantity models do not preserve the names of the sub models
        # which are used to relate the fitted sub models back to the displayed
        # models in the model editor. Go through and hope that their order is
        # preserved.
        if result.n_submodels() > 1:
            for i, x in enumerate(result):
                fit_mod.unitless_model._submodels[i].name = x.name
            sub_mods = [x for x in fit_mod.unitless_model]
        else:
            fit_mod.unitless_model.name = result.name
            sub_mods = [fit_mod.unitless_model]

        disp_mods = {item.text(): item for item in model_editor_model.items}

        for i, sub_mod in enumerate(sub_mods):
            # Get the base astropy model object
            model_item = disp_mods.get(sub_mod.name)

            # For each of the children `StandardItem`s, parse out their
            # individual stored values
            for cidx in range(model_item.rowCount()):
                param_name = model_item.child(cidx, 0).data()

                if result.n_submodels() > 1:
                    parameter = getattr(fit_mod,
                                        "{0}_{1}".format(param_name, i))
                else:
                    parameter = getattr(fit_mod, param_name)

                model_item.child(cidx,
                                 1).setText("{:.4g}".format(parameter.value))
                model_item.child(cidx, 1).setData(parameter.value,
                                                  Qt.UserRole + 1)

                model_item.child(cidx, 3).setData(parameter.fixed,
                                                  Qt.UserRole + 1)

        for i in range(0, 3):
            self.model_tree_view.resizeColumnToContents(i)
Beispiel #9
0
    def _on_create_new_model(self):
        if self.hub.data_item is None:
            self.new_message_box(text="No item selected, cannot create model.",
                                 info="There is currently no item selected. "
                                 "Please select an item before attempting"
                                 " to create a new model.")

        # Grab the currently selected plot data item
        new_spec = Spectrum1D(
            flux=np.zeros(self.hub.data_item.spectral_axis.size) *
            self.hub.data_item.flux.unit,
            spectral_axis=self.hub.data_item.spectral_axis)

        self.create_model_data_item(new_spec, data_item=self.hub.data_item)
Beispiel #10
0
    def _on_create_new_model(self):
        if self.hub.data_item is None:
            QMessageBox.warning(
                self, "No item selected, cannot create model.",
                "There is currently no item selected. Please "
                "select an item before attempting to create "
                "a new model.")
            return

        # Grab the currently selected plot data item
        new_spec = Spectrum1D(
            flux=np.zeros(self.hub.data_item.spectral_axis.size) *
            self.hub.data_item.flux.unit,
            spectral_axis=self.hub.data_item.spectral_axis)

        self.create_model_data_item(new_spec, data_item=self.hub.data_item)
Beispiel #11
0
def test_spectral_axis_conversion():
    np.random.seed(42)

    x, y = build_spectrum()

    spectrum = Spectrum1D(flux=y*u.Jy, spectral_axis=x*u.um)
    new_spectral_axis = "micron"

    converted_spectrum = uc.UnitConversion.process_unit_conversion(uc.UnitConversion,
                                                                   spectrum=spectrum,
                                                                   new_spectral_axis=new_spectral_axis)

    result_spectral_axis = [ 0, 1.11111111, 2.22222222, 3.33333333,
                             4.44444444, 5.55555556, 6.66666667, 7.77777778, 8.88888889, 10]


    assert np.allclose(converted_spectrum.spectral_axis.value, result_spectral_axis, atol=1e-5)
Beispiel #12
0
def test_flux_conversion():

    np.random.seed(42)

    x, y = build_spectrum()

    spectrum = Spectrum1D(flux=y*u.Jy, spectral_axis=x*u.um)
    new_flux = "erg / (s cm2 um)"

    converted_spectrum = uc.UnitConversion.process_unit_conversion(uc.UnitConversion,
                                                                   spectrum=spectrum,
                                                                   new_flux=new_flux)

    # result_flux = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    result_flux = [9.67970066e-09, 2.46763877e-09, 1.12034621e-09, 7.15682642e-10, 5.94364037e-10, 2.80465524e-10, 2.02021996e-10, 1.49988629e-10, 1.21543537e-10]

    # TODO: first element is `inf`, why is this?
    assert np.allclose(converted_spectrum.flux.value[1:], result_flux, atol=1e-5)
Beispiel #13
0
def test_from_list_list():
    g1 = Gaussian1D(1, 4.6, 0.2)
    g2 = Gaussian1D(2.5, 5.5, 0.1)
    g3 = Gaussian1D(-1.7, 8.2, 0.1)

    x = np.linspace(0, 10, 200)
    y = g1(x) + g2(x) + g3(x)

    spectrum = Spectrum1D(flux=y * u.Jy, spectral_axis=x * u.um)

    lines = find_lines_derivative(spectrum, flux_threshold=0.01)

    spec_reg = SpectralRegion.from_line_list(lines)
    expected = [(4.072864321608041 * u.um, 5.072864321608041 * u.um),
                (4.977386934673367 * u.um, 5.977386934673367 * u.um),
                (7.690954773869347 * u.um, 8.690954773869347 * u.um)]

    for i, reg in enumerate(expected):
        assert_quantity_allclose(reg, (spec_reg[i].lower, spec_reg[i].upper))
Beispiel #14
0
def noise_region_uncertainty(spectrum, spectral_region, noise_func=np.std):
    """
    Generates a new spectrum with an uncertainty from the noise in a particular
    region of the spectrum.

    Parameters
    ----------

    spectrum: `~specutils.spectra.Spectrum1D
        The spectrum to which we want to set the uncertainty.

    spectral_region: `~specutils.spectra.SpectralRegion`
        The region to use to calculate the standard deviation.

    noise_func: callable
        A function which takes the flux in the ``spectral_region`` and yields a
        *single* value for the noise to use in the result spectrum.

    Return
    ------
    spectrum_uncertainty: `~specutils.spectra.Spectrum1D
        The ``spectrum``, but with a constant uncertainty set by the result of
        the noise region calculation

    """

    # Extract the sub spectrum based on the region
    sub_spectrum = spectral_region.extract(spectrum)

    # Compute the standard deviation of the flux.
    noise = noise_func(sub_spectrum.flux)

    uncertainty = StdDevUncertainty(noise * np.ones(spectrum.flux.shape))

    # Return new specturm with uncertainty set.
    return Spectrum1D(flux=spectrum.flux,
                      spectral_axis=spectrum.spectral_axis,
                      uncertainty=uncertainty,
                      wcs=spectrum.wcs,
                      unit=spectrum.unit,
                      spectral_axis_unit=spectrum.spectral_axis_unit,
                      velocity_convention=spectrum.velocity_convention,
                      rest_value=spectrum.rest_value)
Beispiel #15
0
    def __call__(self, parameters):
        x = parameters[0]
        y = parameters[1]

        # Calling the Spectrum1D constructor for every spaxel
        # turned out to be less expensive than expected. Experiments
        # show that the cost amounts to a couple percent additional
        # running time in comparison with a version that uses a 3D
        # spectrum as input. Besides, letting an externally-created
        # spectrum reference into the callable somehow prevents it
        # to execute. This behavior was seen also with other functions
        # passed to the callable.
        flux = self.cube[x, y, :] # transposed!
        sp = Spectrum1D(spectral_axis=self.wave, flux=flux)

        fitted_model = fit_lines(sp, self.model)

        fitted_values = fitted_model(self.wave)

        return (x, y, fitted_model, fitted_values)
Beispiel #16
0
def test_both_conversion():
    np.random.seed(42)

    x, y = build_spectrum()

    spectrum = Spectrum1D(flux=y*u.Jy, spectral_axis=x*u.um)
    new_spectral_axis = "micron"
    new_flux = "erg / (s cm2 um)"

    converted_spectrum = uc.UnitConversion.process_unit_conversion(uc.UnitConversion,
                                                                   spectrum=spectrum,
                                                                   new_flux=new_flux,
                                                                   new_spectral_axis=new_spectral_axis)

    result_spectral_axis = [ 0, 1.11111111, 2.22222222, 3.33333333,
                             4.44444444, 5.55555556, 6.66666667, 7.77777778, 8.88888889, 10]
    result_flux = [9.67970066e-09, 2.46763877e-09, 1.12034621e-09, 7.15682642e-10, 5.94364037e-10, 2.80465524e-10, 2.02021996e-10, 1.49988629e-10, 1.21543537e-10]

    # TODO: first element is `inf`, why is this?
    assert np.allclose(converted_spectrum.flux.value[1:], result_flux, atol=1e-5)
    assert np.allclose(converted_spectrum.spectral_axis.value, result_spectral_axis, atol=1e-5)
Beispiel #17
0
    def _on_create_new_model(self):
        if self.hub.data_item is None:
            message_box = QMessageBox()
            message_box.setText("No item selected, cannot create model.")
            message_box.setIcon(QMessageBox.Warning)
            message_box.setInformativeText(
                "There is currently no item selected. Please select an item "
                "before attempting to create a new model.")

            message_box.exec()
            return

        # Set the currently displayed plugin panel widget to the model editor
        self.hub.set_active_plugin_bar(name="Model Editor")

        # Grab the currently selected plot data item
        new_spec = Spectrum1D(
            flux=np.zeros(self.hub.data_item.spectral_axis.size) *
            self.hub.data_item.flux.unit,
            spectral_axis=self.hub.data_item.spectral_axis)

        model_data_item = ModelDataItem(model=ModelFittingModel(),
                                        name="Fittable Model Spectrum",
                                        identifier=uuid.uuid4(),
                                        data=new_spec)

        self.hub.append_data_item(model_data_item)

        plot_data_item = self.hub.plot_data_item_from_data_item(
            model_data_item)

        # Connect data change signals so that the plot updates when the user
        # changes a parameter in the model view model
        model_data_item.model_editor_model.dataChanged.connect(
            lambda tl, br, r, pi=plot_data_item: self._on_model_data_changed(
                tl, br, pi))
Beispiel #18
0
    print('...')

wl, fl = np.genfromtxt(INPUT_spec, unpack=True, skip_header=1)

# TRIM LIMITS
lamb1 = 8405
lamb2 = 8700

interactive_mode = False

#trim spectrum
mask = (wl > lamb1) & (wl < lamb2)
wl = wl[mask]
fl = fl[mask]

spectrum = Spectrum1D(flux=fl * u.Jy, spectral_axis=wl * u.AA)

#         start   end    start    end
regions = [8489.0, 8563.0, 8642.0, 8697.0]

if interactive_mode == False:
    g1_fit = fit_generic_continuum(spectrum,
                                   exclude_regions=[
                                       SpectralRegion(regions[0] * u.AA,
                                                      regions[1] * u.AA),
                                       SpectralRegion(regions[2] * u.AA,
                                                      regions[3] * u.AA)
                                   ])
    y_continuum_fitted = g1_fit(wl * u.AA)

    spec_normalized = spectrum / y_continuum_fitted
Beispiel #19
0
print(len(wavelength))

#STACKED DATA IMPORT
# input_file = '/home/zak/Documents/ProjectData/DAP_Stacks/z_0.5_0.6/CB_C3_LRG/z0506_CBC3LRG_V3_DAP.fits'
#
# hdu = fits.open(input_file)
#
# wavelength = hdu[1].data
# print(wavelength)
# fluxy = hdu[2].data
# len(fluxy)
# #error = flux*0.1
# error = hdu[3].data

#SPECUTILIS ATTEMPT
spectrum = Spectrum1D(flux=galaxy * (1 * u.erg / u.cm**2 / u.s / u.AA),
                      spectral_axis=wavelength * u.AA)
g1_fit = fit_continuum(spectrum)
y_continuum_fitted = g1_fit(wavelength * u.AA)
print(y_continuum_fitted)
spec_normalized = spectrum / y_continuum_fitted

# plt.subplot(2,1,1)
# plt.plot(wavelength, galaxy, label='Spectra',c='C0')
# plt.plot(wavelength,y_continuum_fitted,label='Continuum Fit',c='C1')
# #plt.plot(wavelength, y_continuum_fitted)
# plt.title('Specutilis Continuum Removal')
# plt.legend()
# plt.grid(True)
# spec_normalized = spectrum / y_continuum_fitted
#
# plt.subplot(2,1,2)
Beispiel #20
0
def test_cube_fitting_backend():
    np.random.seed(42)

    SIGMA = 0.1  # noise in data
    TOL = 0.4  # test tolerance

    # Flux cube oriented as in JWST data. To build a Spectrum1D
    # instance with this, one need to transpose it so the spectral
    # axis direction corresponds to the last index.
    flux_cube = np.zeros((SPECTRUM_SIZE, IMAGE_SIZE, IMAGE_SIZE))

    # Generate list of all spaxels to be fitted
    _spx = [[(x, y) for x in range(IMAGE_SIZE)] for y in range(IMAGE_SIZE)]
    spaxels = [item for sublist in _spx for item in sublist]

    # Fill cube spaxels with spectra that differ from
    # each other only by their noise component.
    x, _ = build_spectrum()
    for spx in spaxels:
        flux_cube[:, spx[0], spx[1]] = build_spectrum(sigma=SIGMA)[1]

    # Transpose so it can be packed in a Spectrum1D instance.
    flux_cube = flux_cube.transpose(1, 2, 0)

    spectrum = Spectrum1D(flux=flux_cube*u.Jy, spectral_axis=x*u.um)

    # Initial model for fit.
    g1f = models.Gaussian1D(0.7*u.Jy, 4.65*u.um, 0.3*u.um, name='g1')
    g2f = models.Gaussian1D(2.0*u.Jy, 5.55*u.um, 0.3*u.um, name='g2')
    g3f = models.Gaussian1D(-2.*u.Jy, 8.15*u.um, 0.2*u.um, name='g3')
    zero_level = models.Const1D(1.*u.Jy, name='const1d')

    model_list = [g1f, g2f, g3f, zero_level]
    expression = "g1 + g2 + g3 + const1d"

    # Fit to all spaxels.
    fitted_parameters, fitted_spectrum = fb.fit_model_to_spectrum(
        spectrum, model_list, expression)

    # Check that parameter results are formatted as expected.
    assert type(fitted_parameters) == list
    assert len(fitted_parameters) == 225

    for m in fitted_parameters:
        if m['x'] == 3 and m['y'] == 2:
            fitted_model = m['model']

    assert type(fitted_model[0].amplitude.value) == np.float64
    assert fitted_model[0].amplitude.unit == u.Jy

    assert type(fitted_model[0] == params.Parameter)
    assert type(fitted_model[0].mean.value) == np.float64
    assert fitted_model[0].mean.unit == u.um

    # Check that spectrum result is formatted as expected.
    assert type(fitted_spectrum) == Spectrum1D
    assert len(fitted_spectrum.shape) == 3
    assert fitted_spectrum.shape == (IMAGE_SIZE, IMAGE_SIZE, SPECTRUM_SIZE)
    assert fitted_spectrum.flux.unit == u.Jy

    # The important point here isn't to check the accuracy of the
    # fit, which was already tested elsewhere. We are mostly
    # interested here in checking the correctness of the data
    # packaging into the output products.

    assert np.allclose(fitted_model[0].amplitude.value, 1.09, atol=TOL)
    assert np.allclose(fitted_model[1].amplitude.value, 2.4, atol=TOL)
    assert np.allclose(fitted_model[2].amplitude.value, -1.7, atol=TOL)

    assert np.allclose(fitted_model[0].mean.value, 4.6, atol=TOL)
    assert np.allclose(fitted_model[1].mean.value, 5.5, atol=TOL)
    assert np.allclose(fitted_model[2].mean.value, 8.2, atol=TOL)

    assert np.allclose(fitted_model[0].stddev.value, 0.2, atol=TOL)
    assert np.allclose(fitted_model[1].stddev.value, 0.1, atol=TOL)
    assert np.allclose(fitted_model[2].stddev.value, 0.1, atol=TOL)

    assert np.allclose(fitted_model[3].amplitude.value, 4.0, atol=TOL)
Beispiel #21
0
#  20-jun-2017  PJT             summer project

import os, sys
import numpy as np

from specutils.spectra import Spectrum1D
from astropy.units import Quantity

if len(sys.argv) == 1:
    print("Usage: %s fitsfile [xpos ypos]" % sys.argv[0])
    sys.exit(1)

if len(sys.argv) > 3:
    fitsfile = sys.argv[1]
    pos = [int(sys.argv[2]), int(sys.argv[3])]
    s3 = Spectrum1D.read(fitsfile, format="cubetest1", pos=pos)
elif len(sys.argv) == 2:
    fitsfile = sys.argv[1]
    s3 = Spectrum1D.read(fitsfile, format="cubetest1")

nfreq = len(s3.data)
restfrq = Quantity(s3.wcs.wcs.restfrq, "Hz")
print("RESTFREQ", restfrq)
s3.rest_value = restfrq

### nooo!!!!   needs to be done in reader since we're immutable
s3.velocity_convention = 'relativistic'
#print("RESTFREQ",s3.rest_value(restfrq))

#  even though native units in Hz, it needs the rest
#s3.to_dispersion("MHz",  rest = restfrq)
Beispiel #22
0
def _load_fits(filename):
    from astropy.io import fits
    import astropy.units as u
    import astropy.wcs as fitswcs
    from specutils.spectra import Spectrum1D

    hdu = fits.getheader(filename)
    flux = fits.getdata(filename)

    # Handle multi-array flux from fits by picking that with
    # larger values (presumably smaller is flux error)
    if len(flux.shape) > 1:
        f1, f2 = flux
        if f2[0][0] < f1[0][0]:
            flux = f1[0]
        else:
            flux = f2[0]

    if not ('WAT0_001' in hdu) or hdu['WAT0_001'] == 'system=equispec':
        uflux = u.erg / (u.cm ** 2 * u.s)

        try:
            if 'CDELT1' in hdu:
                cdelt1 = hdu['CDELT1']
            else:
                cdelt1 = hdu['CD1_1']
        except KeyError as e:
            raise e

        if 'CUNIT1' in hdu:
            cunit1 = hdu['CUNIT1']
        else:
            cunit1 = 'Angstrom'

        crval1 = hdu['CRVAL1']
        ctype1 = hdu['CTYPE1']
        crpix1 = hdu['CRPIX1']
        my_wcs = fitswcs.WCS(header={
            'CDELT1': cdelt1,
            'CRVAL1': crval1,
            'CUNIT1': cunit1,
            'CTYPE1': ctype1,
            'CRPIX1': crpix1
        })

        spec = Spectrum1D(flux=flux * uflux, wcs=my_wcs)

        wavel = np.array(spec.wavelength)
        flux = np.array(spec.flux)

    elif hdu['WAT0_001'] == 'system=multispec':
        # This is a terrible .fits system that doesn't contain
        # delta-wavelength info, so I improvised.
        wavel = ''
        for line in (x for x in hdu.keys() if x[:4] == 'WAT2'):
            wavel += hdu[line]
        wavel = wavel.split()
        wavel[-1] = wavel[-1][:-1]   # get rid of end quote
        wavel = wavel[16:]

        # Split error values due to errors in header key-value reading
        wave_to_add = []
        for w in reversed(wavel):
            if w.count('.') == 2:
                w_all = w.split('.')
                # Assumes wavelength is between 1000-9999
                wave_to_add.append('%s.%s' % (w_all[0], w_all[1][:-4]))
                wave_to_add.append('%s.%s' % (w_all[1][-4:], w_all[2]))
                wavel.remove(w)

        # For some reason the wavelength is sometimes reversed
        if float(wavel[0]) > float(wavel[-1]):
            flux = list(reversed(flux))

        wavel += wave_to_add
        wavel = [float(x) for x in wavel]
        wavel = sorted(wavel)

        wavel = np.array(wavel)
        flux = np.array(flux)

    flux_err = np.zeros(len(flux))

    return np.c_[wavel, flux, flux_err]
Beispiel #23
0
def _fit_3D(initial_model, spectrum):
    """
    Fits an astropy CompoundModel to every spaxel in a cube
    using a multiprocessor pool running in parallel. Computes
    realizations of the models over each spaxel.

    Parameters
    ----------
    spectrum : :class:`specutils.spectrum.Spectrum1D`
        The spectrum that stores the cube in its 'flux' attribute.

    Returns
    -------
    :list: a list that stores 2D arrays. Each array contains one
        parameter from `astropy.modeling.CompoundModel` instances
        fitted to every spaxel in the input cube.
    :class:`specutils.spectrum.Spectrum1D`
        The spectrum that stores the fitted model values in its 'flux'
        attribute.
    """

    # Worker for the multiprocess pool.
    worker = SpaxelWorker(spectrum.flux,
                          spectrum.spectral_axis,
                          initial_model)

    # Generate list of all spaxels to be fitted
    spaxels = _generate_spaxel_list(spectrum)

    # Build cube with empty slabs, one per model parameter. These
    # will store only parameter values for now, so a cube suffices.
    parameters_cube = np.zeros(shape=(len(initial_model.parameters),
                                      spectrum.flux.shape[0],
                                      spectrum.flux.shape[1]))

    # Build cube with empty arrays, one per input spaxel. These
    # will store the flux values corresponding to the fitted
    # model realization over each spaxel.
    output_flux_cube = np.zeros(shape=spectrum.flux.shape)

    # Callback to collect results from workers into the cubes
    def collect_result(result):
        x = result[0]
        y = result[1]
        model = result[2]
        fitted_values = result[3]

        # Store fitted model parameters
        for index, name in enumerate(model.param_names):
            param = getattr(model, name)
            parameters_cube[index, x, y] = param.value

        # Store fitted values
        output_flux_cube[x, y, :] = fitted_values

    # Run multiprocessor pool to fit each spaxel and
    # compute model values on that same spaxel.
    results = []
    pool = Pool(mp.cpu_count() - 1)

    for spx in spaxels:
        r = pool.apply_async(worker, (spx,), callback=collect_result)
        results.append(r)
    for r in results:
        r.wait()

    # Collect units from all parameters
    param_units = []
    for name in initial_model.param_names:
        param = getattr(initial_model, name)
        param_units.append(param.unit)

    # Re-format parameters cube to a list of 2D Quantity arrays.
    fitted_parameters = _handle_parameter_units(initial_model,
                                                parameters_cube,
                                                param_units)

    # Build output 3D spectrum
    funit = spectrum.flux.unit
    output_spectrum = Spectrum1D(spectral_axis=spectrum.spectral_axis,
                                 flux=output_flux_cube * funit)

    return fitted_parameters, output_spectrum
Beispiel #24
0
def test_model_fitting(specviz_gui, monkeypatch):
    # Monkeypatch the QMessageBox widget so that it doesn't block the test
    # progression. In this case, accept the information dialog indicating that
    # a loader has been saved.
    monkeypatch.setattr(QMessageBox, "information",
                        lambda *args: QMessageBox.Ok)
    monkeypatch.setattr(QMessageBox, "warning", lambda *args: QMessageBox.Ok)

    hub = Hub(workspace=specviz_gui.current_workspace)

    # Generate fake data
    np.random.seed(42)
    g1 = models.Gaussian1D(1, 0, 0.2)
    g2 = models.Gaussian1D(2.5, 0.5, 0.1)

    x = np.linspace(-1, 1, 200)
    y = g1(x) + g2(x) + np.random.normal(0., 0.2, x.shape)

    # Regular fitting
    gg_init = models.Gaussian1D(1.3, 0, 0.1) + models.Gaussian1D(1.8, 0.5, 0.1)
    fitter = fitting.LevMarLSQFitter()
    gg_fit = fitter(gg_init, x, y)

    # SpecViz fitting
    spectral_axis_unit = u.Unit(hub.plot_window.plot_widget.spectral_axis_unit
                                or "")
    data_units = u.Unit(hub.plot_window.plot_widget.data_unit or "")
    s1d = Spectrum1D(flux=y * data_units, spectral_axis=x * spectral_axis_unit)
    hub.workspace.model.add_data(s1d, name="fitting_data")
    model_editor = specviz_gui.current_workspace._plugin_bars['Model Editor']

    model_editor._on_create_new_model()
    model_editor._add_fittable_model(models.Gaussian1D)
    model_editor._add_fittable_model(models.Gaussian1D)

    index = model_editor.data_selection_combo.findText("fitting_data")
    if index >= 0:
        model_editor.data_selection_combo.setCurrentIndex(index)
    value_dict = {
        'Gaussian1D': {
            'amplitude': '1.3',
            'mean': '0',
            'stddev': '0.1'
        },
        'Gaussian1D1': {
            'amplitude': '1.8',
            'mean': '0.5',
            'stddev': '0.1'
        }
    }
    plot_data_item = hub.plot_item
    model_editor_model = plot_data_item.data_item.model_editor_model

    fill_in_models(model_editor_model, value_dict)

    model_editor._on_fit_clicked(eq_pop_up=False)

    model_editor_model = plot_data_item.data_item.model_editor_model
    result = model_editor_model.evaluate()

    np.testing.assert_allclose(result.parameters, gg_fit.parameters, rtol=1e-4)
Beispiel #25
0
    def __init__(self):

        #
        # Create the base wavelengths and flux
        #

        self.wavelengths_um = np.linspace(0.4, 1.05, 100)

        g1 = models.Gaussian1D(amplitude=2000, mean=0.56, stddev=0.01)
        g2 = models.Gaussian1D(amplitude=500, mean=0.62, stddev=0.02)
        g3 = models.Gaussian1D(amplitude=-400, mean=0.80, stddev=0.02)
        g4 = models.Gaussian1D(amplitude=-350, mean=0.52, stddev=0.01)
        ramp = models.Linear1D(slope=300, intercept=0.0)

        self.base_flux = (g1(self.wavelengths_um) + g2(self.wavelengths_um) +
                          g3(self.wavelengths_um) + g4(self.wavelengths_um) +
                          ramp(self.wavelengths_um) + 1000)

        #
        # Initialize the seed so the random numbers are not quite as random
        #

        np.random.seed(42)

        #
        # Create two spectra with the only difference in the instance of noise
        #

        self._flux_e1 = self.base_flux + 400 * np.random.random(
            self.base_flux.shape)
        self._s1_um_mJy_e1 = Spectrum1D(spectral_axis=self.wavelengths_um *
                                        u.um,
                                        flux=self._flux_e1 * u.mJy)

        self._flux_e2 = self.base_flux + 400 * np.random.random(
            self.base_flux.shape)
        self._s1_um_mJy_e2 = Spectrum1D(spectral_axis=self.wavelengths_um *
                                        u.um,
                                        flux=self._flux_e2 * u.mJy)

        #
        # Create one spectrum with the same flux but in angstrom units
        #

        self.wavelengths_AA = self.wavelengths_um * 10000
        self._s1_AA_mJy_e3 = Spectrum1D(spectral_axis=self.wavelengths_AA *
                                        u.AA,
                                        flux=self._flux_e1 * u.mJy)

        #
        # Create one spectrum with the same flux but in angstrom units and nJy
        #

        self._flux_e4 = (self.base_flux + 400 *
                         np.random.random(self.base_flux.shape)) * 1000000
        self._s1_AA_nJy_e4 = Spectrum1D(spectral_axis=self.wavelengths_AA *
                                        u.AA,
                                        flux=self._flux_e4 * u.nJy)

        #
        # Create one spectrum like 1 but with a mask
        #
        self._s1_um_mJy_e1_masked = copy(
            self._s1_um_mJy_e1
        )  # SHALLOW copy - the data are shared with the above non-masked case  # noqa
        self._s1_um_mJy_e1_masked.mask = (
            np.random.randn(*self.base_flux.shape) + 1) > 0

        # Create a spectrum like 1, but with descending spectral axis
        self._s1_um_mJy_e1_desc = Spectrum1D(
            spectral_axis=self.wavelengths_um[::-1] * u.um,
            flux=self._flux_e1[::-1] * u.mJy)
Beispiel #26
0
def generic_spectrum_from_table(table, wcs=None, **kwargs):
    """
    Load spectrum from an Astropy table into a Spectrum1D object.
    Uses the following logic to figure out which column is which:

     * Spectral axis (dispersion) is the first column with units
     compatible with u.spectral() or with length units such as 'pix'.

     * Flux is taken from the first column with units compatible with
     u.spectral_density(), or with other likely culprits such as
     'adu' or 'cts/s'.

     * Uncertainty comes from the next column with the same units as flux.

    Parameters
    ----------
    file_name: str
        The path to the ECSV file
    wcs : :class:`~astropy.wcs.WCS`
        A FITS WCS object. If this is present, the machinery will fall back
        to using the wcs to find the dispersion information.

    Returns
    -------
    data: Spectrum1D
        The spectrum that is represented by the data in this table.

    Raises
    ------
    Warns if uncertainty has zeros or negative numbers.
    Raises IOError if it can't figure out the columns.

    """

    # Local function to find the wavelength or frequency column
    def _find_spectral_axis_column(table, columns_to_search):
        """
        Figure out which column in a table holds the spectral axis (dispersion).
        Take the first column that has units compatible with u.spectral()
        equivalencies. If none meet that criterion, look for other likely
        length units such as 'pix'.
        """
        additional_valid_units = [u.Unit('pix')]
        found_column = None

        # First, search for a column with units compatible with Angstroms
        for c in columns_to_search:
            try:
                table[c].to("AA", equivalencies=u.spectral())
                found_column = c
                break
            except:
                continue

        # If no success there, check for other possible length units
        if found_column is None:
            for c in columns_to_search:
                if table[c].unit in additional_valid_units:
                    found_column = c
                    break

        return found_column

    # Local function to find the flux column
    def _find_spectral_column(table, columns_to_search, spectral_axis):
        """
        Figure out which column in a table holds the fluxes or uncertainties.
        Take the first column that has units compatible with
        u.spectral_density() equivalencies. If none meet that criterion,
        look for other likely length units such as 'adu' or 'cts/s'.
        """
        additional_valid_units = [u.Unit('adu'), u.Unit('ct/s')]
        found_column = None

        # First, search for a column with units compatible with Janskies
        for c in columns_to_search:
            try:
                # Check for multi-D flux columns
                if table[c].ndim == 1:
                    spec_ax = spectral_axis
                else:
                    # Assume leading dimension corresponds to spectral_axis
                    spec_shape = np.ones(table[c].ndim, dtype=np.int)
                    spec_shape[0] = -1
                    spec_ax = spectral_axis.reshape(spec_shape)
                table[c].to("Jy", equivalencies=u.spectral_density(spec_ax))
                found_column = c
                break
            except:
                continue

        # If no success there, check for other possible flux units
        if found_column is None:
            for c in columns_to_search:
                if table[c].unit in additional_valid_units:
                    found_column = c
                    break

        return found_column

    # Make a copy of the column names so we can remove them as they are found
    colnames = table.colnames.copy()

    # Use the first column that has spectral unit as the dispersion axis
    spectral_axis_column = _find_spectral_axis_column(table, colnames)

    if spectral_axis_column is None and wcs is None:
        raise IOError(
            "Could not identify column containing the wavelength, frequency or energy"
        )
    elif wcs is not None:
        spectral_axis = None
    else:
        spectral_axis = table[spectral_axis_column].to(
            table[spectral_axis_column].unit)
        colnames.remove(spectral_axis_column)

    # Use the first column that has a spectral_density equivalence as the flux
    flux_column = _find_spectral_column(table, colnames, spectral_axis)
    if flux_column is None:
        raise IOError("Could not identify column containing the flux")
    flux = table[flux_column].to(table[flux_column].unit)
    colnames.remove(flux_column)
    # For > 1D data transpose to row-major format
    if flux.ndim > 1:
        flux = flux.T

    # Use the next column with the same units as flux as the uncertainty
    # Interpret it as a standard deviation and check if it has zeros or negative values
    err_column = None
    for c in colnames:
        if table[c].unit == table[flux_column].unit:
            err_column = c
            break
    if err_column is not None:
        if table[err_column].ndim > 1:
            err = table[err_column].T
        elif flux.ndim > 1:  # Repeat uncertainties over all flux columns
            err = np.tile(table[err_column], flux.shape[0], 1)
        else:
            err = table[err_column]
        err = StdDevUncertainty(err.to(err.unit))
        if np.min(table[err_column]) <= 0.:
            warnings.warn("Standard Deviation has values of 0 or less",
                          AstropyUserWarning)

    # Create the Spectrum1D object and return it
    if wcs is not None or spectral_axis_column is not None and flux_column is not None:
        if err_column is not None:
            spectrum = Spectrum1D(flux=flux,
                                  spectral_axis=spectral_axis,
                                  uncertainty=err,
                                  meta=table.meta,
                                  wcs=wcs)
        else:
            spectrum = Spectrum1D(flux=flux,
                                  spectral_axis=spectral_axis,
                                  meta=table.meta,
                                  wcs=wcs)

    return spectrum
Beispiel #27
0
def spectrum_from_column_mapping(table, column_mapping, wcs=None):
    """
    Given a table and a mapping of the table column names to attributes
    on the Spectrum1D object, parse the information into a Spectrum1D.

    Parameters
    ----------
    table : :class:`~astropy.table.Table`
        The table object (e.g. returned from ``Table.read('data_file')``).
    column_mapping : dict
        A dictionary describing the relation between the table columns
        and the arguments of the `Spectrum1D` class, along with unit
        information. The dictionary keys should be the table column names
        while the values should be a two-tuple where the first element is the
        associated `Spectrum1D` keyword argument, and the second element is the
        unit for the file column (or ``None`` to take unit from the table)::

            column_mapping = {'FLUX': ('flux', 'Jy'),
                              'WAVE': ('spectral_axis', 'um')}

    wcs : :class:`~astropy.wcs.WCS` or :class:`gwcs.WCS`
        WCS object passed to the Spectrum1D initializer.
    """
    spec_kwargs = {}

    # Associate columns of the file with the appropriate spectrum1d arguments
    for col_name, (kwarg_name, cm_unit) in column_mapping.items():
        # If the table object couldn't parse any unit information,
        # fallback to the column mapper defined unit
        tab_unit = table[col_name].unit

        if tab_unit and cm_unit is not None:
            # If the table unit is defined, retrieve the quantity array for
            # the column
            kwarg_val = u.Quantity(table[col_name], tab_unit)

            # Attempt to convert the table unit to the user-defined unit.
            logging.debug(
                "Attempting auto-convert of table unit '%s' to "
                "user-provided unit '%s'.", tab_unit, cm_unit)

            if not isinstance(cm_unit, u.Unit):
                cm_unit = u.Unit(cm_unit)
            if cm_unit.physical_type in ('length', 'frequency'):
                # Spectral axis column information
                kwarg_val = kwarg_val.to(cm_unit, equivalencies=u.spectral())
            elif 'spectral flux' in cm_unit.physical_type:
                # Flux/error column information
                kwarg_val = kwarg_val.to(cm_unit,
                                         equivalencies=u.spectral_density(
                                             1 * u.AA))
        elif tab_unit:
            # The user has provided no unit in the column mapping, so we
            # use the unit as defined in the table object.
            kwarg_val = u.Quantity(table[col_name], tab_unit)
        elif cm_unit is not None:
            # In this case, the user has defined a unit in the column mapping
            # but no unit has been defined in the table object.
            kwarg_val = u.Quantity(table[col_name], cm_unit)
        else:
            # Neither the column mapping nor the table contain unit information.
            # This may be desired e.g. for the mask or bit flag arrays.
            kwarg_val = table[col_name]

        spec_kwargs.setdefault(kwarg_name, kwarg_val)

    # Ensure that the uncertainties are a subclass of NDUncertainty
    if spec_kwargs.get('uncertainty') is not None:
        spec_kwargs['uncertainty'] = StdDevUncertainty(
            spec_kwargs.get('uncertainty'))

    return Spectrum1D(**spec_kwargs, wcs=wcs, meta=table.meta)
Beispiel #28
0
plt.xlabel(r'$\lambda$ ($\AA$)',fontsize=11)
plt.ylabel('ADU',fontsize=11)
plt.title('Spectrum of HD 32991',fontsize=11)
plt.show(block=False)

spectra=np.column_stack((w,sf))
np.savetxt('/home/ganesh/Desktop/spectpf.txt',spectra, newline='\n',delimiter=',')

#fitting spectrum
data2=pd.read_csv('/home/ganesh/Desktop/spectpf.txt',header=None)

y=u.Quantity(data2[1],u.dimensionless_unscaled)
x=u.Quantity(data2[0],u.angstrom)

spectrum = Spectrum1D(spectral_axis=x,flux=y)
g1_fit = fit_generic_continuum(spectrum)
y_continuum_fitted = g1_fit(x)

#continuum fitted spectrum
plt.figure()
plt.plot(x, y,'k')
plt.plot(x, y_continuum_fitted, 'orange')
plt.xlabel(r'$\lambda$ ($\AA$)',fontsize=11)
plt.show(block=False)

#Plot continuum normalised intensity
spec_normalized = spectrum / y_continuum_fitted

t=QTable([spec_normalized.spectral_axis,spec_normalized.flux])
t.write('/home/ganesh/Desktop/spectf.txt',format='ascii', delimiter=',')
Beispiel #29
0
def _fit_3D(initial_model, spectrum, window=None, n_cpu=None):
    """
    Fits an astropy CompoundModel to every spaxel in a cube
    using a multiprocessor pool running in parallel. Computes
    realizations of the models over each spaxel.

    Parameters
    ----------
    initial_model : :class: `astropy.modeling.CompoundModel`
        Initial guess for the model to be fitted.
    spectrum : :class:`specutils.spectra.Spectrum1D`
        The spectrum that stores the cube in its 'flux' attribute.
    window : `None` or :class:`specutils.spectra.SpectralRegion`
        See :func:`specutils.fitting.fitmodels.fit_lines`.
    n_cpu : `None` or int
        Number of cores to use for multiprocessing.
        Using all the cores at once is not recommended.
        If `None`, it will use max cores minus one.
        Set this to 1 for debugging.

    Returns
    -------
    :list: a list that stores 2D arrays. Each array contains one
        parameter from `astropy.modeling.CompoundModel` instances
        fitted to every spaxel in the input cube.
    :class:`specutils.spectra.Spectrum1D`
        The spectrum that stores the fitted model values in its 'flux'
        attribute.
    """
    if n_cpu is None:
        n_cpu = mp.cpu_count() - 1

    # Generate list of all spaxels to be fitted
    spaxels = _generate_spaxel_list(spectrum)

    fitted_models = []

    # Build cube with empty arrays, one per input spaxel. These
    # will store the flux values corresponding to the fitted
    # model realization over each spaxel.
    output_flux_cube = np.zeros(shape=spectrum.flux.shape)

    # Callback to collect results from workers into the cubes
    def collect_result(results):
        for i in range(len(results['x'])):
            x = results['x'][i]
            y = results['y'][i]
            model = results['fitted_model'][i]
            fitted_values = results['fitted_values'][i]

            # Store fitted model parameters
            fitted_models.append({"x": x, "y": y, "model": model})

            # Store fitted values
            output_flux_cube[x, y, :] = fitted_values

    # Run multiprocessor pool to fit each spaxel and
    # compute model values on that same spaxel.
    results = []
    pool = Pool(n_cpu)

    # The communicate overhead of spawning a process for each *individual*
    # parameter set is prohibitively high (it's actually faster to run things
    # sequentially). Instead, chunk the spaxel list based on the number of
    # available processors, and have each processor do the model fitting
    # on the entire subset of spaxel tuples, then return the set of results.
    for spx in np.array_split(spaxels, n_cpu):
        # Worker for the multiprocess pool.
        worker = SpaxelWorker(spectrum.flux,
                              spectrum.spectral_axis,
                              initial_model,
                              param_set=spx,
                              window=window)
        r = pool.apply_async(worker, callback=collect_result)
        results.append(r)
    for r in results:
        r.wait()

    pool.close()

    # Build output 3D spectrum
    funit = spectrum.flux.unit
    output_spectrum = Spectrum1D(spectral_axis=spectrum.spectral_axis,
                                 flux=output_flux_cube * funit)

    return fitted_models, output_spectrum
def _fit_3D(initial_model, spectrum):
    """
    Fits an astropy CompoundModel to every spaxel in a cube
    using a multiprocessor pool running in parallel. Computes
    realizations of the models over each spaxel.

    Parameters
    ----------
    spectrum : :class:`specutils.spectrum.Spectrum1D`
        The spectrum that stores the cube in its 'flux' attribute.

    Returns
    -------
    :list: a list that stores 2D arrays. Each array contains one
        parameter from `astropy.modeling.CompoundModel` instances
        fitted to every spaxel in the input cube.
    :class:`specutils.spectrum.Spectrum1D`
        The spectrum that stores the fitted model values in its 'flux'
        attribute.
    """

    # Generate list of all spaxels to be fitted
    spaxels = _generate_spaxel_list(spectrum)

    # Build cube with empty slabs, one per model parameter. These
    # will store only parameter values for now, so a cube suffices.
    parameters_cube = np.zeros(shape=(len(initial_model.parameters),
                                      spectrum.flux.shape[0],
                                      spectrum.flux.shape[1]))

    # Build cube with empty arrays, one per input spaxel. These
    # will store the flux values corresponding to the fitted
    # model realization over each spaxel.
    output_flux_cube = np.zeros(shape=spectrum.flux.shape)

    # Callback to collect results from workers into the cubes
    def collect_result(results):
        print('Enter callback')
        for i in range(len(results['x'])):
            x = results['x'][i]
            y = results['y'][i]
            model = results['fitted_model'][i]
            fitted_values = results['fitted_values'][i]
            # Store fitted model parameters
            for index, name in enumerate(model.param_names):
                param = getattr(model, name)
                parameters_cube[index, x, y] = param.value

            # Store fitted values
            output_flux_cube[x, y, :] = fitted_values

    # Run multiprocessor pool to fit each spaxel and
    # compute model values on that same spaxel.
    results = []
    pool = Pool(mp.cpu_count() - 1)

    # The communicate overhead of spawning a process for each *individual*
    # parameter set is prohibitively high (it's actually faster to run things
    # sequentially). Instead, chunk the spaxel list based on the number of
    # available processors, and have each processor do the model fitting
    # on the entire subset of spaxel tuples, then return the set of results.
    for spx in np.array_split(spaxels, mp.cpu_count() - 1):
        #for spx in np.array_split(spaxels, 1):
        # Worker for the multiprocess pool.
        worker = SpaxelWorker(spectrum.flux,
                              spectrum.spectral_axis,
                              initial_model,
                              param_set=spx)
        r = pool.apply_async(worker, callback=collect_result)
        results.append(r)
    for r in results:
        r.wait()

    pool.close()

    # Collect units from all parameters
    param_units = []
    for name in initial_model.param_names:
        param = getattr(initial_model, name)
        param_units.append(param.unit)

    # Re-format parameters cube to a dict of 2D Quantity arrays.
    fitted_parameters = _handle_parameter_units(initial_model, parameters_cube,
                                                param_units)

    # Build output 3D spectrum
    funit = spectrum.flux.unit
    output_spectrum = Spectrum1D(spectral_axis=spectrum.spectral_axis,
                                 flux=output_flux_cube * funit)

    return fitted_parameters, output_spectrum