Python linelist.LineList類代碼示例

def test_subset():
    ism = LineList('ISM')
    subset = np.array([1215.6700, 1608.4511])*u.AA
    ism = ism.subset_lines(subset)
    assert len(ism._data) == 2
    np.testing.assert_allclose(ism['FeII 1608']['wrest'], 1608.4511*u.AA, rtol=1e-7)

    # Now with names
    ism = LineList('ISM')
    subset = ['HI 1215', 'HI 1025', 'CIV 1548']
    ism = ism.subset_lines(subset)
    np.testing.assert_allclose(ism['HI 1215']['wrest'], 1215.6700*u.AA, rtol=1e-7)

def set_llist(llist, in_dict=None, sort=True):
    """ Method to set a line list dict for the Widgets

    Parameters
    ----------
    sort : bool, optional
      Sort lines by rest wavelength
    """
    from linetools.lists.linelist import LineList
    from astropy.units.quantity import Quantity

    if in_dict is None:
        in_dict = dict(Lists=[])

    if isinstance(llist,basestring): # Set line list from a file
        in_dict['List'] = llist
        in_dict['Lists'].append(llist)
        if llist == 'None':
            in_dict['Plot'] = False
        else:
            in_dict['Plot'] = True
            # Load?
            if not (llist in in_dict):
                # Homebrew
                if llist == 'OVI':
                    gdlines = u.AA*[629.730, 702.332, 770.409, 780.324, 787.711, 832.927, 972.5367, 977.0201,
                        1025.7222, 1031.9261, 1037.6167, 1206.5, 1215.6700, 1260.4221]
                    llist_cls = LineList('Strong')
                    llist_cls = llist_cls.subset_lines(gdlines)

                    in_dict[llist] = llist_cls
                else:
                    llist_cls = LineList(llist)
                    # Sort
                    llist_cls._data.sort('wrest')
                    # Load
                    in_dict[llist] = llist_cls
    elif isinstance(llist, (Quantity, list)): # Set from a list of wrest
        in_dict['List'] = 'input.lst'
        in_dict['Lists'].append('input.lst')
        in_dict['Plot'] = True
        # Fill
        if sort:
            llist.sort()
        llist_cls = LineList('ISM')
        llist_cls = llist_cls.subset_lines(llist)
        in_dict['input.lst'] = llist_cls
    else:
        raise IOError('Not ready for this type of input')

    # Return
    return in_dict

    def fill_data(self, trans, linelist=None, closest=False, verbose=True):
        """ Fill atomic data and setup analy.

        Parameters
        ----------
        trans : Quantity or str
          Either a rest wavelength (e.g. 1215.6700*u.AA) or the name
          of a transition (e.g. 'CIV 1548'). For an unknown transition
          use string 'unknown'.
        linelist : LineList, optional
          Class of linelist or str setting LineList
        closest : bool, optional
          Take the closest line to input wavelength? [False]
        """

        # Deal with LineList
        if linelist is None:
            llist = LineList('ISM')
        elif isinstance(linelist,basestring):
            llist = LineList(linelist)
        elif isinstance(linelist,LineList):
            llist = linelist
        else:
            raise ValueError('Bad input for linelist')

        # Closest?
        llist.closest = closest

        # Data
        newline = llist[trans]
        try:
            self.data.update(newline)  # Expected to be a LineList dict object
        except TypeError:
            pdb.set_trace()


        # Update
        self.wrest = self.data['wrest']
        self.name = self.data['name']

        #
        self.analy.update( {
            'flg_eye': 0,
            'flg_limit': 0, # No limit
            'datafile': '', 
            'name': self.data['name']
            })

        # Additional fundamental attributes for Absorption Line
        self.attrib.update(abs_attrib.copy())

def set_llist(llist, in_dict=None, sort_by='wrest'):
    """ Method to set a line list dict for the Widgets

    Parameters
    ----------
    sort_by : str or list of str, optional
        Key(s)to sort the lines by. Default is 'wrest'.
        If sort_by='as_given', it preserves the order
        as given by llist.
    """
    from linetools.lists.linelist import LineList
    from astropy.units.quantity import Quantity

    if in_dict is None:
        in_dict = dict(Lists=[])

    if isinstance(llist,basestring): # Set line list from a file
        in_dict['List'] = llist
        in_dict['Lists'].append(llist)
        if llist == 'None':
            in_dict['Plot'] = False
        else:
            in_dict['Plot'] = True
            # Load?
            if not (llist in in_dict):
                # Homebrew
                if llist == 'OVI':
                    gdlines = u.AA*[629.730, 702.332, 770.409, 780.324, 787.711, 832.927, 972.5367, 977.0201,
                        1025.7222, 1031.9261, 1037.6167, 1206.5, 1215.6700, 1260.4221]
                    llist_cls = LineList('Strong', sort_by=sort_by)
                    llist_cls = llist_cls.subset_lines(gdlines, sort_by='as_given')

                    in_dict[llist] = llist_cls
                else:
                    llist_cls = LineList(llist, sort_by=sort_by)
                    # Load
                    in_dict[llist] = llist_cls
    elif isinstance(llist, (Quantity, list)): # Set from a list of wrest
        in_dict['List'] = 'input.lst'
        in_dict['Lists'].append('input.lst')
        in_dict['Plot'] = True
        # Fill
        llist_cls = LineList('ISM', sort_by=sort_by)
        llist_cls = llist_cls.subset_lines(llist, sort_by=sort_by)
        in_dict['input.lst'] = llist_cls
    else:
        raise IOError('Not ready for this type of input')

    # Return
    return in_dict

def test_all_transitions():
    error_msg = 'Something is wrong in all_transitions()'
    ism = LineList('ISM')
    #check simple case
    line = 'OVI'
    ovi_transitions = ism.all_transitions(line)
    assert len(ovi_transitions) == 2, error_msg
    #print(ovi_transitions['name'])
    #check unknown
    line = 'unknown'
    out = ism.all_transitions(line)
    assert type(out) == dict, error_msg
    #check case of single transition ion
    line = 'CIII'
    out = ism.all_transitions(line)
    assert type(out) == dict, error_msg
    #check case of transitions from excited levels
    line='FeII*'
    out = ism.all_transitions(line)
    assert len(out) == 27, "wrong line counts"
    print(out)
    # wrest
    out = ism.all_transitions(1215.6700*u.AA)
    assert len(out) == 30,"wrong line counts" # 30 Lyman series transitions
    #print('test_all_transitions() passed')
    h2 = LineList('H2')
    line = 'B19-0P(1)'
    out = h2.all_transitions(line)
    assert len(out) == 7

    def fill_data(self, trans, linelist=None, closest=False, verbose=True):
        """ Fill atomic data and setup analy.

        Parameters
        ----------
        trans : Quantity or str
          Either a rest wavelength (e.g. 1215.6700*u.AA) or the name
          of a transition (e.g. 'CIV 1548'). For an unknown transition
          use string 'unknown'.
        linelist : LineList, optional
          Class of linelist or str setting LineList
        closest : bool, optional
          Take the closest line to input wavelength? [False]
        """

        # Deal with LineList
        if linelist is None:
            if self.ltype == 'Abs':
                llist = LineList('ISM')
            elif self.ltype == 'Em':
                llist = LineList('Galaxy')
            else:
                raise ValueError("Not ready for ltype = {:s}".format(self.ltype))
        elif isinstance(linelist,basestring):
            llist = LineList(linelist)
        elif isinstance(linelist,LineList):
            llist = linelist
        else:
            raise ValueError('Bad input for linelist')

        # Closest?
        llist.closest = closest

        # Data
        newline = llist[trans]
        if newline is None:
            raise ValueError("Transition {} not found in LineList {:s}".format(trans, llist.list))
        try:
            self.data.update(newline)  # Expected to be a LineList dict object
        except TypeError:
            raise TypeError("Probably should not be here")


        # Update
        self.wrest = self.data['wrest']
        self.name = self.data['name']

        #
        self.update()  # is this used ?

    def fill_data(self, trans, linelist=None, closest=False):
        """ Fill atomic data and setup analy.

        Parameters
        ----------
        trans : Quantity or str
          Either a rest wavelength (e.g. 1215.6700*u.AA) or the name
          of a transition (e.g. 'CIV 1548'). For an unknown transition
          use string 'unknown'.
        linelist : LineList, optional
          Class of linelist or str setting LineList
        closest : bool, optional
          Take the closest line to input wavelength? [False]
        """

        # Deal with LineList
        if linelist is None:
            llist = LineList('ISM')
        elif isinstance(linelist,basestring):
            llist = LineList(linelist)
        elif isinstance(linelist,LineList):
            llist = linelist
        else:
            raise ValueError('Bad input for linelist')

        # Closest?
        llist.closest = closest

        # Data
        newline = llist[trans]
        self.data.update(newline)

        # Update
        self.wrest = self.data['wrest']
        self.trans = self.data['name']

        #
        self.analy.update( {
            'flg_eye': 0,
            'flg_limit': 0, # No limit
            'datafile': '', 
            'name': self.data['name']
            })

        # Additional attributes for Absorption Line
        self.attrib.update({'N': 0., 'Nsig': 0., 'flagN': 0, # Column
                       'b': 0.*u.km/u.s, 'bsig': 0.*u.km/u.s  # Doppler
                       } )

def test_all_transitions():
    error_msg = 'Something is wrong in all_transitions()'
    ism = LineList('ISM')
    #check simple case
    line = 'OVI'
    ovi_transitions = ism.all_transitions(line)
    assert len(ovi_transitions) == 2, error_msg
    #print(ovi_transitions['name'])
    #check unknown
    line = 'unknown'
    out = ism.all_transitions(line)
    assert type(out) == dict, error_msg
    #check case of single transition ion
    line = 'CIII'
    out = ism.all_transitions(line)
    assert type(out) == dict, error_msg
    # wrest
    out = ism.all_transitions(1215.6700*u.AA)
    assert len(out) == 30 # 30 Lyman series transitions

def test_strongest_transitions():
    error_msg = 'Something is wrong in strongest_transitions()'
    ism = LineList('ISM')
    wvlims = (1200,1800)*u.AA
    z = 0.5
    transitions = ism.strongest_transitions('HI',wvlims/(1+z),n_max=5)
    assert len(transitions) == 5,  error_msg
    assert transitions[0]['name'] == 'HI 1025' , error_msg
    assert isinstance(transitions,QTable), error_msg

    wvlims = (1500,1700)*u.AA
    z = 0.5
    transitions = ism.strongest_transitions('HI',wvlims/(1+z),n_max=5)
    assert isinstance(transitions,dict), error_msg #only Lyb should be available, so dict is expected
    assert transitions['name'] == 'HI 1025'

    wvlims = (1100,1200)*u.AA
    z = 0.0
    transitions = ism.strongest_transitions('HI',wvlims/(1+z),n_max=5)
    assert transitions is None, error_msg

def test_sortdata():
    error_msg = 'Something is wrong in sortdata()'
    ism = LineList('ISM', sort_by='name')
    assert ism.name[0] == 'AlII 1670', error_msg
    ism.sortdata('name', reverse=True)
    assert ism.name[0] == 'ZrIII 1798', error_msg
    ism.sortdata(['abundance', 'rel_strength'], reverse=True)
    assert ism.name[0] == 'HI 1215', error_msg
    ism.sortdata(['rel_strength'])
    assert ism.name[0] == 'CI** 1123b', error_msg

def test_available_transitions():
    error_msg = 'Something is wrong in available_transitions()'
    ism = LineList('ISM')
    wvlims = (900,1800)*u.AA
    z = 0.1

    transitions = ism.available_transitions(wvlims/(1+z),n_max_tuple=5)
    assert transitions[2]['name'] == 'HI 972' , error_msg
    assert isinstance(transitions,QTable), error_msg

    transitions = ism.available_transitions(wvlims/(1+z),n_max_tuple=2)
    assert transitions[2]['name'] == 'CIII 977' , error_msg

    wvlims = (1200,1800)*u.AA
    z = 0.5
    transitions = ism.available_transitions(wvlims/(1+z), n_max_tuple=2)
    assert transitions[0]['name'] == 'HI 1025', error_msg 
    assert 'OVI 1031' in transitions['name'], error_msg 
    assert 'CIII 977' in transitions['name'], error_msg

    wvlims = (1000,3000)*u.AA
    z = 1.5
    transitions = ism.available_transitions(wvlims/(1+z),n_max_tuple=2)
    assert 'NeVIII 770' in transitions['name'], error_msg
    assert 'MgX 609' in transitions['name'], error_msg
    assert 'HI 1215' not in transitions['name'], error_msg

    wvlims = (1215.6,1217)*u.AA
    z = 0
    transitions = ism.available_transitions(wvlims/(1+z),n_max_tuple=2)
    assert isinstance(transitions,dict), error_msg

 def init_lines(self):
     '''Fill up the component lines
     '''
     if self.linelist is None:
         self.linelist = LineList('Strong')
     # Get the lines
     #QtCore.pyqtRemoveInputHook()
     #xdb.set_trace()
     #QtCore.pyqtRestoreInputHook()
     all_trans = self.linelist.all_transitions(self.init_wrest)
     for trans in all_trans:
         self.lines.append(AbsLine(trans['wrest'],
             linelist=self.linelist))
     # Sync
     self.sync_lines()

class Component(object):
    def __init__(self, z, wrest, vlim=[-300.,300]*u.km/u.s,
        linelist=None):
        # Init
        self.init_wrest = wrest
        self.zcomp = z
        self.vlim = vlim
        self.attrib = {'N': 0., 'Nsig': 0., 'flagN': 0, # Column
                       'b': 0.*u.km/u.s, 'bsig': 0.*u.km/u.s,  # Doppler
                       'z': self.zcomp, 'zsig': 0.,
                       'Quality': 'None'}
        self.comment = 'None'
        #
        self.linelist = linelist
        self.lines = []
        self.init_lines()
        #
        self.name = 'z{:.5f}_{:s}'.format(
            self.zcomp,self.lines[0].data['name'].split(' ')[0])
        #
    def init_lines(self):
        '''Fill up the component lines
        '''
        if self.linelist is None:
            self.linelist = LineList('Strong')
        # Get the lines

        all_trans = self.linelist.all_transitions(self.init_wrest)
        #QtCore.pyqtRemoveInputHook()
        #xdb.set_trace()
        #QtCore.pyqtRestoreInputHook()
        if isinstance(all_trans,dict):
            all_trans = [all_trans]
        for trans in all_trans:
            self.lines.append(AbsLine(trans['wrest'],
                linelist=self.linelist))

        # Sync
        self.sync_lines()

    def sync_lines(self):
        '''Synchronize attributes of the lines
        '''
        for line in self.lines:
            line.attrib['N'] = self.attrib['N']
            line.attrib['b'] = self.attrib['b']
            line.attrib['z'] = self.attrib['z']

def test_set_extra_columns_to_datatable():
    ism = LineList('ISM')
    # bad calls
    try:
        ism.set_extra_columns_to_datatable(abundance_type='incorrect_one')
    except ValueError:
        pass
    try:
        ism.set_extra_columns_to_datatable(ion_correction='incorrect_one')
    except ValueError:
        pass
    # test expected strongest value
    ism.set_extra_columns_to_datatable(ion_correction='none', abundance_type='solar')
    np.testing.assert_allclose(ism['HI 1215']['rel_strength'], 14.704326420257642)
    tab = ism._data

    np.testing.assert_allclose(np.max(tab['rel_strength']), 14.704326420257642)

    def add_abslines_from_linelist(self, llist='ISM', wvlim=None, min_Wr=None, **kwargs):
        """
        It adds associated AbsLines satisfying some conditions (see parameters below).

        Parameters
        ----------
        llist : str
            Name of the linetools.lists.linelist.LineList
            object where to look for the transition names.
            Default is 'ISM', which means the function looks
            within `list = LineList('ISM')`.
        wvlims : Quantity array, optional
            Observed wavelength limits for AbsLines to be added.
            e.g. [1200, 2000]*u.AA.
        min_Wr : Quantity, optional
            Minimum rest-frame equivalent with for AbsLines to be added.
            This is calculated in the very low optical depth regime tau0<<1,
            where Wr is independent of Doppler parameter or gamma (see eq. 9.15 of
            Draine 2011). Still, a column density attribute for the AbsComponent
            is needed.

        Returns
        -------
        Adds AbsLine objects to the AbsComponent._abslines list.

        Notes
        -----
        **kwargs are passed to AbsLine.add_absline() method.

        """
        # get the transitions from LineList
        llist = LineList(llist)
        name = ions.ion_name(self.Zion, nspace=0)
        transitions = llist.all_transitions(name)
        # unify output to be always QTable
        if isinstance(transitions, dict):
            transitions = llist.from_dict_to_qtable(transitions)

        # check wvlims
        if wvlim is not None:
            cond = (transitions['wrest']*(1+self.zcomp) >= wvlim[0]) & \
                   (transitions['wrest']*(1+self.zcomp) <= wvlim[1])
            transitions = transitions[cond]

        # check outputs
        if len(transitions) == 0:
            warnings.warn("No transitions satisfying the criteria found. Doing nothing.")
            return

        # loop over the transitions when more than one found
        for transition in transitions:
            iline = AbsLine(transition['name'], z=self.zcomp)
            iline.limits.set(self.vlim)
            iline.attrib['coord'] = self.coord
            iline.attrib['logN'] = self.logN
            iline.attrib['sig_logN'] = self.sig_logN
            iline.attrib['flag_N'] = self.flag_N
            iline.attrib['N'] = 10**iline.attrib['logN'] / (u.cm * u.cm)
            iline.attrib['sig_N'] = 10**iline.attrib['sig_logN'] / (u.cm * u.cm)

            for key in self.attrib.keys():
                iline.attrib[key] = self.attrib[key]

            if min_Wr is not None:
                # check logN is defined
                logN = self.logN
                if logN == 0:
                    warnings.warn("AbsComponent does not have logN defined. Appending AbsLines "
                                 "regardless of min_Wr.")
                else:
                    N = 10**logN / (u.cm*u.cm)
                    Wr_iline = iline.get_Wr_from_N(N=N)  # valid for the tau0<<1 regime.
                    if Wr_iline < min_Wr:  # do not append
                        continue
            # add the absline
            self.add_absline(iline)