Python util.parse_quantity函数代码示例

def parse_quantity_linspace(quantity_linspace_dictionary, add_one=True):
    """
    parse a dictionary of the following kind
    {'start': 5000 km/s,
     'stop': 10000 km/s,
     'num': 1000}

    Parameters
    ----------

    quantity_linspace_dictionary: ~dict

    add_one: boolean, default: True

    Returns
    -------

    ~np.array

    """

    start = parse_quantity(quantity_linspace_dictionary['start'])
    stop = parse_quantity(quantity_linspace_dictionary['stop'])

    try:
        stop = stop.to(start.unit)
    except u.UnitsError:
        raise ConfigurationError('"start" and "stop" keyword must be compatible quantities')

    num = quantity_linspace_dictionary['num']
    if add_one:
        num += 1

    return np.linspace(start.value, stop.value, num=num) * start.unit

    def test_spectrum_section(self):
        assert_almost_equal(self.config['spectrum']['start'],
                            parse_quantity(self.yaml_data['spectrum']['start']))
        assert_almost_equal(self.config['spectrum']['end'],
                            parse_quantity(self.yaml_data['spectrum']['stop']))

        assert self.config['spectrum']['bins'] == self.yaml_data['spectrum']['num']

    def parse_exponential(density_dict, v_inner, v_outer, time_explosion):
        time_0 = density_dict.pop('time_0', 19.9999584)
        if isinstance(time_0, basestring):
            time_0 = parse_quantity(time_0).to('s').value
        else:
            logger.debug('time_0 not supplied for density branch85 - using sensible default %g', time_0)
        try:
            rho_0 = density_dict.pop('rho_0')
            if isinstance(rho_0, basestring):
                rho_0 = parse_quantity(rho_0).to('g/cm^3').value
            else:
                raise KeyError
        except KeyError:
            rho_0 = 1e-2
            logger.warning('rho_o was not given in the config! Using %g', rho_0)
        try:
            v_0 = density_dict.pop('v_0')
            if isinstance(v_0, basestring):
                v_0 = parse_quantity(v_0).to('km/s').value
            
        except KeyError:
            v_0 = 1
            logger.warning('v_0 was not given in the config file! Using %f km/s', v_0)

        velocities = 0.5 * (v_inner + v_outer)
        densities = calc_exponential_density(velocities, v_0, rho_0)
        densities = u.Quantity(densities, 'g/cm^3')
        return densities

    def parse_power_law(density_dict, v_inner, v_outer, time_explosion):
        time_0 = density_dict.pop('time_0', 19.9999584)
        if isinstance(time_0, basestring):
            time_0 = parse_quantity(time_0).to('s')
        else:
            logger.debug('time_0 not supplied for density powerlaw - using sensible default %g', time_0)
        try:
            rho_0 = density_dict.pop('rho_0')
            if isinstance(rho_0, basestring):
                rho_0 = parse_quantity(rho_0)
            else:
                raise KeyError
        except KeyError:
            rho_0 = parse_quantity('1e-2 g/cm^3')
            logger.warning('rho_o was not given in the config! Using %g', rho_0)
        try:
            exponent = density_dict.pop('exponent')
        except KeyError:
            exponent = 2
            logger.warning('exponent was not given in the config file! Using %f', exponent)
        try:
            v_0 = density_dict.pop('v_0')
            if isinstance(v_0, basestring):
                v_0 = parse_quantity(v_0).to('cm/s')
            
        except KeyError:
            v_0 = parse_quantity('1 cm/s')
            logger.warning('v_0 was not given in the config file! Using %f km/s', v_0)

            
        velocities = 0.5 * (v_inner + v_outer)
        densities = calc_power_law_density(velocities, v_0, rho_0, exponent)
        densities = calculate_density_after_time(densities, time_0, time_explosion)
        return densities

def parse_spectral_bin(spectral_bin_boundary_1, spectral_bin_boundary_2):
    spectral_bin_boundary_1 = parse_quantity(spectral_bin_boundary_1).to('Angstrom', u.spectral())
    spectral_bin_boundary_2 = parse_quantity(spectral_bin_boundary_2).to('Angstrom', u.spectral())

    spectrum_start_wavelength = min(spectral_bin_boundary_1, spectral_bin_boundary_2)
    spectrum_end_wavelength = max(spectral_bin_boundary_1, spectral_bin_boundary_2)

    return spectrum_start_wavelength, spectrum_end_wavelength

    def parse_artis_density(density_file_dict, time_explosion):
        density_file = density_file_dict["name"]
        for i, line in enumerate(file(density_file)):
            if i == 0:
                no_of_shells = np.int64(line.strip())
            elif i == 1:
                time_of_model = u.Quantity(float(line.strip()), "day").to("s")
            elif i == 2:
                break

        velocities, mean_densities_0 = np.recfromtxt(density_file, skip_header=2, usecols=(1, 2), unpack=True)
        # converting densities from log(g/cm^3) to g/cm^3 and stretching it to the current ti
        velocities = u.Quantity(np.append([0], velocities), "km/s").to("cm/s")
        mean_densities_0 = u.Quantity(10 ** mean_densities_0, "g/cm^3")

        mean_densities = calculate_density_after_time(mean_densities_0, time_of_model, time_explosion)

        # Verifying information
        if len(mean_densities) == no_of_shells:
            logger.debug("Verified ARTIS file %s (no_of_shells=length of dataset)", density_file)
        else:
            raise ConfigurationError(
                "Error in ARTIS file %s - Number of shells not the same as dataset length" % density_file
            )

        min_shell = 1
        max_shell = no_of_shells

        v_inner = velocities[:-1]
        v_outer = velocities[1:]

        volumes = (4 * np.pi / 3) * (time_of_model ** 3) * (v_outer ** 3 - v_inner ** 3)
        masses = (volumes * mean_densities_0 / constants.M_sun).to(1)

        logger.info(
            "Read ARTIS configuration file %s - found %d zones with total mass %g Msun",
            density_file,
            no_of_shells,
            sum(masses.value),
        )

        if "v_lowest" in density_file_dict:
            v_lowest = parse_quantity(density_file_dict["v_lowest"]).to("cm/s").value
            min_shell = v_inner.value.searchsorted(v_lowest)
        else:
            min_shell = 1

        if "v_highest" in density_file_dict:
            v_highest = parse_quantity(density_file_dict["v_highest"]).to("cm/s").value
            max_shell = v_outer.value.searchsorted(v_highest)
        else:
            max_shell = no_of_shells

        v_inner = v_inner[min_shell:max_shell]
        v_outer = v_outer[min_shell:max_shell]
        mean_densities = mean_densities[min_shell:max_shell]

        return v_inner, v_outer, mean_densities, min_shell, max_shell

def parse_supernova_section(supernova_dict):
    """
    Parse the supernova section

    Parameters
    ----------

    supernova_dict: dict
        YAML parsed supernova dict

    Returns
    -------

    config_dict: dict

    """
    config_dict = {}

    # parse luminosity
    luminosity_value, luminosity_unit = supernova_dict["luminosity_requested"].strip().split()

    if luminosity_unit == "log_lsun":
        config_dict["luminosity_requested"] = (
            10 ** (float(luminosity_value) + np.log10(constants.L_sun.cgs.value)) * u.erg / u.s
        )
    else:
        config_dict["luminosity_requested"] = (float(luminosity_value) * u.Unit(luminosity_unit)).to("erg/s")

    config_dict["time_explosion"] = parse_quantity(supernova_dict["time_explosion"]).to("s")

    if "distance" in supernova_dict:
        config_dict["distance"] = parse_quantity(supernova_dict["distance"])
    else:
        config_dict["distance"] = None

    if "luminosity_wavelength_start" in supernova_dict:
        config_dict["luminosity_nu_end"] = parse_quantity(supernova_dict["luminosity_wavelength_start"]).to(
            "Hz", u.spectral()
        )
    else:
        config_dict["luminosity_nu_end"] = np.inf * u.Hz

    if "luminosity_wavelength_end" in supernova_dict:
        config_dict["luminosity_nu_start"] = parse_quantity(supernova_dict["luminosity_wavelength_end"]).to(
            "Hz", u.spectral()
        )
    else:
        config_dict["luminosity_nu_start"] = 0.0 * u.Hz

    return config_dict

def parse_supernova_section(supernova_dict):
    """
    Parse the supernova section

    Parameters
    ----------

    supernova_dict: dict
        YAML parsed supernova dict

    Returns
    -------

    config_dict: dict

    """
    config_dict = {}

    #parse luminosity
    luminosity_value, luminosity_unit = supernova_dict['luminosity_requested'].strip().split()

    if luminosity_unit == 'log_lsun':
        config_dict['luminosity_requested'] = 10 ** (
        float(luminosity_value) + np.log10(constants.L_sun.cgs.value)) * u.erg / u.s
    else:
        config_dict['luminosity_requested'] = (float(luminosity_value) * u.Unit(luminosity_unit)).to('erg/s')

    config_dict['time_explosion'] = parse_quantity(supernova_dict['time_explosion']).to('s')

    if 'distance' in supernova_dict:
        config_dict['distance'] = parse_quantity(supernova_dict['distance'])
    else:
        config_dict['distance'] = None

    if 'luminosity_wavelength_start' in supernova_dict:
        config_dict['luminosity_nu_end'] = parse_quantity(supernova_dict['luminosity_wavelength_start']). \
            to('Hz', u.spectral())
    else:
        config_dict['luminosity_nu_end'] = np.inf * u.Hz

    if 'luminosity_wavelength_end' in supernova_dict:
        config_dict['luminosity_nu_start'] = parse_quantity(supernova_dict['luminosity_wavelength_end']). \
            to('Hz', u.spectral())
    else:
        config_dict['luminosity_nu_start'] = 0.0 * u.Hz

    return config_dict

def test_parse_quantity():
    q1 = parse_quantity('5 km/s')
    assert q1.value == 5.
    assert q1.unit == u.Unit('km/s')

    with pytest.raises(MalformedQuantityError):
        parse_quantity(5)

    with pytest.raises(MalformedQuantityError):
        parse_quantity('abcd')

    with pytest.raises(MalformedQuantityError):
        parse_quantity('a abcd')

    with pytest.raises(MalformedQuantityError):
        parse_quantity('5 abcd')

    def parse_branch85(density_dict, v_inner, v_outer, time_explosion):

        time_0 = density_dict.pop('time_0', 19.9999584)
        if isinstance(time_0, basestring):
            time_0 = parse_quantity(time_0).to('s')
        else:
            time_0 *= u.s
            logger.debug('time_0 not supplied for density branch85 - using sensible default %g', time_0)

        density_coefficient = density_dict.pop('density_coefficient', None)
        if density_coefficient is None:
            density_coefficient = 3e29 * u.Unit('g/cm^3')
            logger.debug('density_coefficient not supplied for density type branch85 - using sensible default %g',
                         density_coefficient)
        else:
            density_coefficient = parse_quantity(density_coefficient)

        velocities = 0.5 * (v_inner + v_outer)
        densities = density_coefficient * (velocities.value * 1e-5) ** -7

        densities = calculate_density_after_time(densities, time_0, time_explosion)

        return densities

    def parse_exponential(density_dict, v_inner, v_outer, time_explosion):
        time_0 = density_dict.pop('time_0', 19.9999584)
        if isinstance(time_0, basestring):
            time_0 = parse_quantity(time_0).to('s').value
        else:
            logger.debug('time_0 not supplied for density branch85 - using sensible default %g', time_0)
        try:
            rho_0 = float(density_dict.pop('rho_0'))
        except KeyError:
            rho_0 = 1e-2
            logger.warning('rho_o was not given in the config! Using %g', rho_0)
        try:
            exponent = density_dict.pop('exponent')
        except KeyError:
            exponent = 2
            logger.warning('exponent was not given in the config file! Using %f', exponent)

        velocities = 0.5 * (v_inner + v_outer)
        densities = calculate_exponential_densities(velocities, v_inner[0], rho_0, exponent)

        return densities

def read_simple_ascii_density(fname):
    """
    Reading a density file of the following structure (example; lines starting with a hash will be ignored):
    The first density describes the mean density in the center of the model and is not used.
    5 s
    #index velocity [km/s] density [g/cm^3]
    0 1.1e4 1.6e8
    1 1.2e4 1.7e8

    Parameters
    ----------

    fname: str
        filename or path with filename


    Returns
    -------

    time_of_model: ~astropy.units.Quantity
        time at which the model is valid

    data: ~pandas.DataFrame
        data frame containing index, velocity (in km/s) and density
    """

    with open(fname) as fh:
        time_of_model_string = fh.readline().strip()
        time_of_model = parse_quantity(time_of_model_string)

    data = recfromtxt(fname, skip_header=1,
                      names=('index', 'velocity', 'density'),
                      dtype=(int, float, float))
    velocity = (data['velocity'] * u.km / u.s).to('cm/s')
    mean_density = (data['density'] * u.Unit('g/cm^3'))[1:]

    return time_of_model, velocity, mean_density

def test_quantity_parser_normal():
    q1 = parse_quantity('5 km/s')
    assert q1.value == 5.
    assert q1.unit == u.Unit('km/s')

def test_quantity_parser_malformed_quantity2():
    with pytest.raises(MalformedQuantityError):
        q1 = parse_quantity('5 abcd')

    def parse_artis_model_setup_files(model_file_section_dict, time_explosion):

        ###### Reading the structure part of the ARTIS file pair
        structure_fname = model_file_section_dict['structure_fname']

        for i, line in enumerate(file(structure_fname)):
            if i == 0:
                no_of_shells = np.int64(line.strip())
            elif i == 1:
                time_of_model = u.Quantity(float(line.strip()), 'day').to('s')
            elif i == 2:
                break

        artis_model_columns = ['velocities', 'mean_densities_0', 'ni56_fraction', 'co56_fraction', 'fe52_fraction',
                               'cr48_fraction']
        artis_model = np.recfromtxt(structure_fname, skip_header=2, usecols=(1, 2, 4, 5, 6, 7), unpack=True,
                                    dtype=[(item, np.float64) for item in artis_model_columns])
        #converting densities from log(g/cm^3) to g/cm^3 and stretching it to the current ti
        velocities = u.Quantity(np.append([0], artis_model['velocities']), 'km/s').to('cm/s')
        mean_densities_0 = u.Quantity(10 ** artis_model['mean_densities_0'], 'g/cm^3')

        mean_densities = calculate_density_after_time(mean_densities_0, time_of_model, time_explosion)


        #Verifying information
        if len(mean_densities) == no_of_shells:
            logger.debug('Verified ARTIS model structure file %s (no_of_shells=length of dataset)', structure_fname)
        else:
            raise ConfigurationError(
                'Error in ARTIS file %s - Number of shells not the same as dataset length' % structure_fname)

        v_inner = velocities[:-1]
        v_outer = velocities[1:]

        volumes = (4 * np.pi / 3) * (time_of_model ** 3) * ( v_outer ** 3 - v_inner ** 3)
        masses = (volumes * mean_densities_0 / constants.M_sun).to(1)

        logger.info('Read ARTIS configuration file %s - found %d zones with total mass %g Msun', structure_fname,
                    no_of_shells, sum(masses.value))

        if 'v_lowest' in model_file_section_dict:
            v_lowest = parse_quantity(model_file_section_dict['v_lowest']).to('cm/s').value
            min_shell = v_inner.value.searchsorted(v_lowest)
        else:
            min_shell = 1

        if 'v_highest' in model_file_section_dict:
            v_highest = parse_quantity(model_file_section_dict['v_highest']).to('cm/s').value
            max_shell = v_outer.value.searchsorted(v_highest)
        else:
            max_shell = no_of_shells
        artis_model = artis_model[min_shell:max_shell]
        v_inner = v_inner[min_shell:max_shell]
        v_outer = v_outer[min_shell:max_shell]
        mean_densities = mean_densities[min_shell:max_shell]

        ###### Reading the abundance part of the ARTIS file pair
        abundances_fname = model_file_section_dict['abundances_fname']
        abundances = pd.DataFrame(np.loadtxt(abundances_fname)[min_shell:max_shell, 1:].transpose(),
                                  index=np.arange(1, 31))

        ni_stable = abundances.ix[28] - artis_model['ni56_fraction']
        co_stable = abundances.ix[27] - artis_model['co56_fraction']
        fe_stable = abundances.ix[26] - artis_model['fe52_fraction']
        mn_stable = abundances.ix[25] - 0.0
        cr_stable = abundances.ix[24] - artis_model['cr48_fraction']
        v_stable = abundances.ix[23] - 0.0
        ti_stable = abundances.ix[22] - 0.0

        abundances.ix[28] = ni_stable
        abundances.ix[28] += artis_model['ni56_fraction'] * np.exp(
            -(time_explosion * inv_ni56_efolding_time).to(1).value)

        abundances.ix[27] = co_stable
        abundances.ix[27] += artis_model['co56_fraction'] * np.exp(
            -(time_explosion * inv_co56_efolding_time).to(1).value)
        abundances.ix[27] += (inv_ni56_efolding_time * artis_model['ni56_fraction'] /
                              (inv_ni56_efolding_time - inv_co56_efolding_time)) * \
                             (np.exp(-(inv_co56_efolding_time * time_explosion).to(1).value) - np.exp(
                                 -(inv_ni56_efolding_time * time_explosion).to(1).value))

        abundances.ix[26] = fe_stable
        abundances.ix[26] += artis_model['fe52_fraction'] * np.exp(
            -(time_explosion * inv_fe52_efolding_time).to(1).value)
        abundances.ix[26] += ((artis_model['co56_fraction'] * inv_ni56_efolding_time
                               - artis_model['co56_fraction'] * inv_co56_efolding_time
                               + artis_model['ni56_fraction'] * inv_ni56_efolding_time
                               - artis_model['ni56_fraction'] * inv_co56_efolding_time
                               - artis_model['co56_fraction'] * inv_ni56_efolding_time * np.exp(
            -(inv_co56_efolding_time * time_explosion).to(1).value)
                               + artis_model['co56_fraction'] * inv_co56_efolding_time * np.exp(
            -(inv_co56_efolding_time * time_explosion).to(1).value)
                               - artis_model['ni56_fraction'] * inv_ni56_efolding_time * np.exp(
            -(inv_co56_efolding_time * time_explosion).to(1).value)
                               + artis_model['ni56_fraction'] * inv_co56_efolding_time * np.exp(
            -(inv_ni56_efolding_time * time_explosion).to(1).value))
                              / (inv_ni56_efolding_time - inv_co56_efolding_time))

        abundances.ix[25] = mn_stable
        abundances.ix[25] += (inv_fe52_efolding_time * artis_model['fe52_fraction'] /
#.........这里部分代码省略.........