Python astropy.units 模块，K 实例源码

def direct_ionization_cross_section(self, energy: u.erg):
        """
        Calculate direct ionization cross-section.

        The cross-sections are calculated one of two ways:

        - Using the method of [1]_ for hydrogenic and He-like ions
        - Using the scaled cross-sections of [2]_ for all other ions

        References
        ----------
        .. [1] Fontes, C. J., et al., 1999, Phys. Rev. A., `59 1329 <https://journals.aps.org/pra/abstract/10.1103/PhysRevA.59.1329>`_
        .. [2] Dere, K. P., 2007, A&A, `466, 771 <http://adsabs.harvard.edu/abs/2007A%26A...466..771D>`_
        """
        is_hydrogenic = (self.atomic_number - self.charge_state == 1) and (self.atomic_number >= 6)
        is_he_like = (self.atomic_number - self.charge_state == 2) and (self.atomic_number >= 10)

        if is_hydrogenic or is_he_like:
            return self._fontes_cross_section(energy)
        else:
            return self._dere_cross_section(energy)

def _dere_cross_section(self, energy: u.erg):
        """
        Calculate direct ionization cross-sections according to [1]_.

        References
        ----------
        .. [1] Dere, K. P., 2007, A&A, `466, 771 <http://adsabs.harvard.edu/abs/2007A%26A...466..771D>`_
        """
        # Cross-sections from diparams file
        cross_section_total = np.zeros(energy.shape)
        for ip,bt_c,bt_e,bt_cross_section in zip(self._diparams['ip'], self._diparams['bt_c'], self._diparams['bt_e'],
                                                 self._diparams['bt_cross_section']):
            U = energy/(ip.to(u.erg))
            scaled_energy = 1. - np.log(bt_c)/np.log(U - 1. + bt_c)
            f_interp = interp1d(bt_e.value, bt_cross_section.value, kind='cubic', fill_value='extrapolate')
            scaled_cross_section = f_interp(scaled_energy.value)*bt_cross_section.unit
            # Only nonzero at energies above the ionization potential
            scaled_cross_section *= (U.value > 1.)
            cross_section = scaled_cross_section*(np.log(U) + 1.)/U/(ip**2)
            if not hasattr(cross_section_total, 'unit'):
                cross_section_total = cross_section_total*cross_section.unit
            cross_section_total += cross_section

        return cross_section_total

def mean_spectra(region,line,file_extension,restFreq,spec_param):
    '''
    Sum spectra over entire mapped region
    Cubes are missing BUNIT header parameter. Fix. 
    '''
    filein = '{0}/0{}_{1}_{2}_trim.fits'.format(region,line,file_extension)
    #add_fits_units(filein,'K')
    cube = SpectralCube.read(filein)
    #trim_edge_cube(cube)
    slice_unmasked = cube.unmasked_data[:,:,:]
    if line == 'NH3_33':
        slice_unmasked[spec_param['mask33_chans'][0]:spec_param['mask33_chans'][1],:,:]=0.
    summed_spectrum = np.nanmean(slice_unmasked,axis=(1,2))
    cube2 = cube.with_spectral_unit(u.km/u.s,velocity_convention='radio',
                                    rest_value=restFreq*u.GHz)
    return summed_spectrum, cube2.spectral_axis

def __init__(self, data, header, temperature_bin_edges: u.K, **kwargs):
        self.temperature_bin_edges = temperature_bin_edges
        # sanitize header
        meta_base = header.copy()
        meta_base['temp_unit'] = self.temperature_bin_edges.unit.to_string()
        meta_base['bunit'] = data.unit.to_string()
        # build map list
        map_list = []
        for i in range(self.temperature_bin_edges.shape[0] - 1):
            tmp = GenericMap(data[:,:,i], meta_base)
            tmp.meta['temp_a'] = self.temperature_bin_edges[i].value
            tmp.meta['temp_b'] = self.temperature_bin_edges[i+1].value
            tmp.plot_settings.update(kwargs.get('plot_settings', {}))
            map_list.append(tmp)

        # call super method
        super().__init__(map_list)

def load_results(self, loop):
        """
        Load EBTEL output for a given loop object.

        Parameters
        ----------
        loop : `synthesizAR.Loop` object
        """
        # load text
        N_s = len(loop.field_aligned_coordinate)
        _tmp = np.loadtxt(loop.hydro_configuration['output_filename'])

        # reshape into a 1D loop structure with units
        time = _tmp[:, 0]*u.s
        electron_temperature = np.outer(_tmp[:, 1], np.ones(N_s))*u.K
        ion_temperature = np.outer(_tmp[:, 2], np.ones(N_s))*u.K
        density = np.outer(_tmp[:, 3], np.ones(N_s))*(u.cm**(-3))
        velocity = np.outer(_tmp[:, -2], np.ones(N_s))*u.cm/u.s
        # flip sign of velocity at apex
        i_mirror = np.where(np.diff(loop.coordinates.value[:, 2]) > 0)[0][-1] + 2
        velocity[:, i_mirror:] = -velocity[:, i_mirror:]

        return time, electron_temperature, ion_temperature, density, velocity

def non_equilibrium_ionization(self, time: u.s, temperature: u.K, density: u.cm**(-3), rate_matrix=None,
                                   initial_condition=None):
        """
        Compute the ionization fraction in non-equilibrium for a given temperature and density
        timeseries.
        """
        if rate_matrix is None:
            rate_matrix = self._rate_matrix()
        if initial_condition is None:
            initial_condition = self.equilibrium_ionization(rate_matrix=rate_matrix)

        interpolate_indices = [np.abs(self.temperature - t).argmin() for t in temperature]
        y = np.zeros(time.shape + (self.atomic_number+1,))
        y[0, :] = initial_condition[interpolate_indices[0], :]

        identity = u.Quantity(np.eye(self.atomic_number + 1))
        for i in range(1, time.shape[0]):
            dt = time[i] - time[i-1]
            term1 = identity - density[i]*dt/2.*rate_matrix[interpolate_indices[i], :, :]
            term2 = identity + density[i-1]*dt/2.*rate_matrix[interpolate_indices[i-1], :, :]
            y[i, :] = np.linalg.inv(term1) @ term2 @ y[i-1, :]
            y[i, :] = np.fabs(y[i, :])
            y[i, :] /= y[i, :].sum()

        return u.Quantity(y)

def __init__(self, element_name, temperature: u.K, hdf5_path=None, **kwargs):
        self.temperature = temperature
        if type(element_name) is str:
            element_name = element_name.capitalize()
        self.atomic_symbol = plasmapy.atomic.atomic_symbol(element_name)
        self.atomic_number = plasmapy.atomic.atomic_number(element_name)
        self.element_name = plasmapy.atomic.element_name(element_name)
        if hdf5_path is None:
            self.hdf5_dbase_root = fiasco.defaults['hdf5_dbase_root']
        else:
            self.hdf5_dbase_root = hdf5_path
        self._ion_kwargs = kwargs.get('ion_kwargs', {})
        self._ion_kwargs['hdf5_path'] = self.hdf5_dbase_root

def __init__(self, ion_name, temperature: u.K, *args, **kwargs):
        super().__init__(ion_name, *args, **kwargs)
        self.temperature = temperature
        # Get selected datasets
        # TODO: do not hardcode defaults, pull from rc file
        self._dset_names = {}
        self._dset_names['ioneq_filename'] = kwargs.get('ioneq_filename', 'chianti')
        self._dset_names['abundance_filename'] = kwargs.get('abundance_filename', 'sun_photospheric_1998_grevesse')
        self._dset_names['ip_filename'] = kwargs.get('ip_filename', 'chianti')

def radiative_recombination_rate(self):
        """
        Radiative recombination rate

        The recombination rate due to interaction with the ambient radiation field
        is calculated using a set of fit parameters using one of two methods:

        - Method of [1]_, (show expression)
        - Method of [2]_, (show expression)

        References
        ----------
        .. [1] Badnell, N. R., 2006, APJS, `167 334 <http://adsabs.harvard.edu/abs/2006ApJS..167..334B>`_
        .. [2] Shull, J. M., Van Steenberg, M., 1982, `48 95 <http://adsabs.harvard.edu/abs/1982ApJS...48...95S>`_
        """
        if self._rrparams['fit_type'][0] == 1 or self._rrparams['fit_type'][0] == 2:
            A = self._rrparams['A_fit']
            B = self._rrparams['B_fit']
            if self._rrparams['fit_type'] == 2:
                B = B + self._rrparams['C_fit'] * np.exp(-self._rrparams['T2_fit'] / self.temperature)
            T0 = self._rrparams['T0_fit']
            T1 = self._rrparams['T1_fit']

            return A/(np.sqrt(self.temperature/T0) * (1 + np.sqrt(self.temperature/T0))**(1. - B)
                      * (1. + np.sqrt(self.temperature/T1))**(1. + B))
        elif self._rrparams['fit_type'][0] == 3:
            return self._rrparams['A_fit'] * (self.temperature/(1e4*u.K))**(-self._rrparams['eta_fit'])
        else:
            raise ValueError('Unrecognized fit type {}'.format(self._rrparams['fit_type']))

def preprocessor(self, table, line, index):
        line = line.strip().split()
        if index == 0:
            filetype = int(line[0])
            table.append([filetype])
            if filetype == 1:
                self.dtypes = [int,float,float,float,float]
                self.units = [None,(u.cm**3)/u.s,None,u.K,u.K]
                self.headings = ['fit_type', 'A_fit', 'B_fit', 'T0_fit', 'T1_fit']
                self.descriptions = ['fit type','A fit parameter','B fit parameter',
                                     'T0 fit parameter','T1 fit parameter']
            elif filetype == 2:
                self.dtypes = [int,float,float,float,float,float,float]
                self.units = [None,(u.cm**3)/u.s,None,u.K,u.K,None,u.K]
                self.headings = ['fit_type', 'A_fit', 'B_fit', 'T0_fit', 'T1_fit', 'C_fit', 'T2_fit']
                self.descriptions = ['fit type','A fit parameter','B fit parameter',
                                     'T0 fit parameter','T1 fit parameter','C fit parameter',
                                     'T2 fit parameter']
            elif filetype == 3:
                self.dtypes = [int,float,float]
                self.units = [None,(u.cm**3)/u.s,None]
                self.headings = ['fit_type', 'A_fit', 'eta_fit']
                self.descriptions = ['fit type','A rad fit parameter','eta fit parameter']
            else:
                raise ValueError('Unrecognized .rrparams filetype {}'.format(filetype))
        else:
            if table[0][0] == 1 or table[0][0] == 2:
                table[0] += line[3:]
            else:
                table[0] += line[2:]

def preprocessor(self, table, line, index):
        line = line.strip().split()
        if index == 0:
            self._drparams_filetype = int(line[0])
            if self._drparams_filetype == 1:
                # Badnell type
                self.dtypes = [int,float,float]
                self.units = [None,u.K,(u.cm**3)/u.s*(u.K**(3/2))]
                self.headings = ['fit_type', 'E_fit', 'c_fit']
                self.descriptions = ['fit type', 'E fit parameter','c fit parameter']
            elif self._drparams_filetype == 2:
                # Shull type
                self.dtypes = [int,float,float,float,float]
                self.units = [None,(u.cm**3)/u.s*(u.K**(3/2)),u.dimensionless_unscaled,u.K,u.K]
                self.headings = ['fit_type', 'A_fit', 'B_fit', 'T0_fit', 'T1_fit']
                self.descriptions = ['fit type','A fit coefficient','B fit coefficient',
                                     'T0 fit coefficient','T1 fit coefficient']
            else:
                raise ValueError('Unrecognized drparams filetype {}'.format(self._drparams_filetype))
        else:
            if self._drparams_filetype == 1:
                tmp = np.array(line[2:],dtype=float)
                if index%2 == 0:
                    tmp_col = table[-1]
                    for i in range(tmp.shape[0]):
                        table.append([self._drparams_filetype,tmp_col[i],tmp[i]])
                    del table[0]
                else:
                    table.append(tmp)
            else:
                table.append([self._drparams_filetype]+line[2:])

def add_fits_units(filein,bunit):
    hdu = fits.open(filein)
    header = hdu[0].header
    header.set('BUNIT', 'K')
    hdu.writeto(filein,clobber=True)
    hdu.close()

def mask_spikes(cube,tmax):
    mask = cube > tmax * u.K
    cube2 = cube.with_mask(mask)
    return cube2

def make_slope_map(self, temperature_bounds=None, em_threshold=None, rsquared_tolerance=0.5):
        """
        Create map of emission measure slopes by fitting :math:`\mathrm{EM}\sim T^a` for a 
        given temperature range. Only those pixels for which the minimum :math:`\mathrm{EM}`
        across all temperature bins is above some threshold value.

        .. warning:: This method provides no measure of the goodness of the fit. Some slope values
                     may not provide an accurate fit to the data.
        """
        if temperature_bounds is None:
            temperature_bounds = u.Quantity((1e6, 4e6), u.K)
        if em_threshold is None:
            em_threshold = u.Quantity(1e25, u.cm**(-5))
        # cut on temperature
        temperature_bin_centers = (self.temperature_bin_edges[:-1] + self.temperature_bin_edges[1:])/2.
        index_temperature_bounds = np.where(np.logical_and(temperature_bin_centers >= temperature_bounds[0],
                                                           temperature_bin_centers <= temperature_bounds[1]))
        temperature_fit = temperature_bin_centers[index_temperature_bounds].value
        # unwrap to 2D and threshold
        data = self.as_array()*u.Unit(self[0].meta['bunit'])
        flat_data = data.reshape(np.prod(data.shape[:2]), temperature_bin_centers.shape[0])
        index_data_threshold = np.where(np.min(flat_data[:,index_temperature_bounds[0]], axis=1) >= em_threshold)
        flat_data_threshold = flat_data.value[index_data_threshold[0],:][:,index_temperature_bounds[0]]
        # very basic but vectorized fitting
        _, rss_flat, _, _, _ = np.polyfit(np.log10(temperature_fit), np.log10(flat_data_threshold.T), 0, full=True)
        coefficients, rss, _, _, _ = np.polyfit(np.log10(temperature_fit), np.log10(flat_data_threshold.T), 1, full=True)
        rsquared = 1. - rss/rss_flat
        slopes = np.where(rsquared >= rsquared_tolerance, coefficients[0], 0.)
        # rebuild into a map
        slopes_flat = np.zeros(flat_data.shape[0])
        slopes_flat[index_data_threshold[0]] = slopes
        slopes_2d = np.reshape(slopes_flat, data.shape[:2])
        base_meta = self[0].meta.copy()
        base_meta['temp_a'] = temperature_fit[0]
        base_meta['temp_b'] = temperature_fit[-1]
        base_meta['bunit'] = ''
        base_meta['detector'] = r'$\mathrm{EM}(T)$ slope'
        base_meta['comment'] = 'Linear fit to log-transformed LOS EM'
        plot_settings = self[0].plot_settings.copy()
        plot_settings['norm'] = None

        return GenericMap(slopes_2d, base_meta, plot_settings=plot_settings)

def make_emission_measure_map(time: u.s, field, instr, temperature_bin_edges=None, **kwargs):
    """
    Return a cube of maps showing the true emission meausure in each pixel
    as a function of electron temperature.
    """
    plot_settings = {'cmap': cm.get_cmap('magma'),
                     'norm': colors.SymLogNorm(1, vmin=1e25, vmax=1e29)}
    plot_settings.update(kwargs.get('plot_settings', {}))

    # read unbinned temperature and density
    with h5py.File(instr.counts_file, 'r') as hf:
        try:
            i_time = np.where(np.array(hf['time'])*u.Unit(hf['time'].attrs['units']) == time)[0][0]
        except IndexError:
            raise IndexError('{} is not a valid time in observing time for {}'.format(time, instr.name))
        unbinned_temperature = np.array(hf['electron_temperature'][i_time,:])
        temperature_unit = u.Unit(hf['electron_temperature'].attrs['units'])
        unbinned_density = np.array(hf['density'][i_time,:])
        density_unit = u.Unit(hf['density'].attrs['units'])

    # setup bin edges and weights
    if temperature_bin_edges is None:
        temperature_bin_edges = 10.**(np.arange(5.5, 7.5, 0.1))*u.K
    x_bin_edges = np.diff(instr.bin_range.x)/instr.bins.x*np.arange(instr.bins.x+1) + instr.bin_range.x[0]
    y_bin_edges = np.diff(instr.bin_range.y)/instr.bins.y*np.arange(instr.bins.y+1) + instr.bin_range.y[0]
    z_bin_edges = np.diff(instr.bin_range.z)/instr.bins.z*np.arange(instr.bins.z+1) + instr.bin_range.z[0]
    z_bin_indices = np.digitize(instr.total_coordinates.value[:,2], z_bin_edges, right=True)
    dh = np.diff(z_bin_edges)[z_bin_indices - 1]
    emission_measure_weights = (unbinned_density**2)*dh
    # bin in x,y,T space with emission measure weights
    xyT_coordinates = np.append(instr.total_coordinates.value[:,:2], 
                                unbinned_temperature[:,np.newaxis], axis=1)
    hist, _ = np.histogramdd(xyT_coordinates, 
                             bins=[x_bin_edges, y_bin_edges, temperature_bin_edges.value], 
                             weights=emission_measure_weights)

    meta_base = instr.make_fits_header(field, instr.channels[0])
    del meta_base['wavelnth']
    del meta_base['waveunit']
    meta_base['detector'] = r'$\mathrm{EM}(T)$'
    meta_base['comment'] = 'LOS Emission Measure distribution'
    data = np.transpose(hist, (1,0,2))*density_unit*density_unit*instr.total_coordinates.unit

    return EMCube(data, meta_base, temperature_bin_edges, plot_settings=plot_settings)