line survey后续二：标分子名

为了把标完分子后的图不至于很杂乱，尽量做到清晰简洁。
思路：

1. 窗口很宽，先把窗口拆分为二；
2. 标分子是个大工程。我们可以先把跃迁线较少的分子标上；
3. 对于很多跃迁线的分子，如果每一个跃迁位置都标分子就会显得很杂乱，我的做法是把靠的比较近的跃迁统一只标一个分子名；
4. 调整标记的上下左右位置。
   代码如下：

import matplotlib.pyplot as plt
import numpy as np
plt.ion()

spw = 'spw35'
data_file='I17233-3606_%s_17h26m42.432s_-36d9m18.075s.txt' %(spw)
model_file='%s_spectrum_output.dat' %(spw)
tran_file='tag_transitions_%s.txt' %(spw)

#Pls give the molecules with a lot transitions here.
molecules = ['H2CCN;v=0;hyp','CH3COCH3;v=0','CH3OCH3;v=0','CH3NC;v=0;#1','C2H5CN;v=0',
             'CH3OCHO;v18=1','CH3OCHO;v=0','H2CCO;v=0','C2H5CN;v12=1','C2H5CN;v20=1',
             'C2H5C-13-N;v=0','CH3C-13-H2CN;v=0','C-13-H3CH2CN;v=0','S-33-O2;v=0',
             'HCCCCCN;v11=2','HCCCCCN;v11=3','HCCCN;v7=2']
             
y_up = 174
y_down = -7


#define the font1 styles
font1 = {'weight' : 'normal',
         'size'   : 15,
        }
#define the font2 styles
font2 = {'weight' : 'normal',
         'size'   : 30,
        }

#define the font3 styles
font3 = {'weight' : 'normal',
         'size'   : 15,
        }


def looptoget(tag,tags,tag_x,n):
    a = 0
    flag = True
    end_x = 0
    try:
        new_tags = tags[n+a+1:][tag_x[n+a+1:]<tag_x[n+a+1]+50]
    except IndexError:
        new_tags = []
        a = 0
        flag=False
    while flag:
        if (tag in new_tags):
            try:
                new_tags = tags[n+a+1:][tag_x[n+a+1:]<tag_x[n+a+1]+50]
            except IndexError:
                new_tags = []
            a = a+1
        else:
            flag=False
            a = a+1
    if (a!=0):
        try:
            end_x=tag_x[n+a]
        except IndexError:
            end_x = tag_x[-1]
            a = a-2
    return (a,end_x)

def lastpos(lists,mystr):
    p = 0
    for i in range(len(lists)):
        if mystr==lists[i]:
            p = i
    return p

def getlineyup(x_,tag_x_,y__,n):
    x_1 = x_[tag_x_[n]<x_]
    y_0 = y__[tag_x_[n]<x_]
    y_1 = y_0[x_1<tag_x_[n]+10]
    line_y_upper = np.nanmax(y_1)
    return line_y_upper
    
def getyupper(text_position,line_x,line_y_upper,y_up):
    x,y = np.array(text_position).T[0],np.array(text_position).T[1]
    x_ = np.array(text_position).T[0]
    x_.sort()
    x_ = x_[x_<line_x[1]]
    if len(x_)==0:
        y_ = 0
        x_ = [0,0]
    else:
        p = x.tolist().index(x_[-1])
        y_ = y[p]
    if ((abs(x_[-1]-line_x[1])>10) or (line_y_upper-y_>y_up/4)):
        line_y_upper = line_y_upper
    else:
        line_y_upper = y_+y_up/4
    return line_y_upper

def loaddata(data_file,model_file,tran_file):
    #load observed line data
    exp_x,exp_y = np.loadtxt(data_file, usecols=(0,1), unpack=True)
    x_ = exp_x[:int(len(exp_x)/2)]
    y_ = exp_y[:int(len(exp_x)/2)]
    _x = exp_x[int(len(exp_x)/2):]
    _y = exp_y[int(len(exp_x)/2):]
    #load model line data
    model_x,model_y = np.loadtxt(model_file, usecols=(0,1), unpack=True)
    x__ = model_x[:int(len(model_x)/2)]
    y__ = model_y[:int(len(model_x)/2)]
    __x = model_x[int(len(model_x)/2):]
    __y = model_y[int(len(model_x)/2):]
    #load transition information file
    tags,note=np.loadtxt(tran_file, dtype=str, usecols=(0,3), unpack=True)
    tag_x, tag_y=np.loadtxt(tran_file, usecols=(1,2), unpack=True)
    tags_ = tags[tag_x<np.max(x_)]
    note_ = note[tag_x<np.max(x_)]
    tag_x_ = tag_x[tag_x<np.max(x_)]
    tag_y_ = tag_y[tag_x<np.max(x_)]
    _tags = tags[tag_x>np.max(x_)]
    _note = note[tag_x>np.max(x_)]
    _tag_x = tag_x[tag_x>np.max(x_)]
    _tag_y = tag_y[tag_x>np.max(x_)]
    return(x_,y_,x__,y__,tags_,note_,tag_x_,tag_y_,_x,_y,__x,__y,_tags,_note,_tag_x,_tag_y)

def plotline(x_, y_,x__, y__, y_down, y_up):
    plt.plot(x_, y_, drawstyle='steps', color='black', linewidth = 2,alpha=1.0)
    plt.axis([x_.min(),x_.max(), y_down, y_up])
    plt.xlabel('Frequency (MHz)', font1)
    plt.ylabel('T$_b$ (K)', font1)
    plt.tick_params(direction='in')
    plt.xticks(size = 20)
    plt.yticks(size = 20)
    plt.plot(x__, y__, drawstyle='steps', color='red', linewidth = 1.5, alpha=0.7)

#plot the molecule transitions dose not in the list of molecules
def plotsimpletransition(molecules,tag_x_,x_,y__,tag_y_,tags_,note_,y_up,n):
    if (note_[n] in molecules)==False:
        if n==0:
            if tag_x_[0]-x_[0]<10:
                line_x = [x_[0],x_[0]+9]
            else:
                line_x = [tag_x_[0],tag_x_[0]]
            line_y_upper = getlineyup(x_,tag_x_,y__,n)+y_up/10
        else:
            line_x = [tag_x_[n],tag_x_[n]]
            line_y_upper = getlineyup(x_,tag_x_,y__,n)+y_up/10
            if ((tag_x_[n]-tag_x_[n-1]<10)==True) and ((note_[n-1] in molecules)==False):
                line_y_upper = line_y_upper+y_up/5
                if n>4:
                    a = 2
                    while a<16:
                        if tag_x_[n]-tag_x_[n-a]<10:
                            line_y_upper = line_y_upper+y_up/5
                            a=a+1
                        else:
                            a = 16
        if line_y_upper>y_up-30:
            line_y_upper=y_up-30
            
        if tag_x_[n] in [98675.3573,98619.8544,99691.1276,100076.5186]:
            line_y_upper=line_y_upper+20
        if tag_x_[n] in [99672.1388]:
            line_y_upper=line_y_upper-25
        if tag_x_[n] in [101341.716]:
            line_y_upper=line_y_upper-15
        line_y = [tag_y_[n],line_y_upper]
        plt.plot([line_x[0],(line_x[0]+line_x[1])/2],line_y,color='blue',linewidth=1,linestyle='dashed')
        plt.text(line_x[1],line_y[1],tags_[n],fontdict=font3,rotation=angle,horizontalalignment='center', verticalalignment='bottom')
        text_position.append([line_x[1],line_y[1]])
    return(text_position)

#plot the molecule transitions in the list of molecules
def plotmultitrans(tags_,mol,tag_x_,x_,y__,note_,text_position,y_up,tag_y_,alpha):
    n = 0
    while n<(len(tags_)):
        if note_[n]==mol:
            a,end_x = looptoget(tags_[n],tags_,tag_x_,n)
            new_tags = tags_[n:n+a+1]
            if tag_x_[n]==100311.14:
                print(a)
            if a>1:
                line_y_upper = np.nanmax(tag_y_[n:n+a])+y_up/8
                value = []
                for i in range(a):
                    line_x = [tag_x_[n+i],tag_x_[n+i]]
                    line_y_upper = getyupper(text_position,line_x,line_y_upper,y_up)
                    value.append(line_y_upper)
                line_y_upper = np.nanmax(line_y_upper)+alpha
                #Coution: The line_y_upper of all transitions in a range should be the same.
                if tag_x_[n] in [98278.871,98281.105,98800.3974,98800.8904,98829.4957]:
                    line_y_upper = line_y_upper +70
                if tag_x_[n] in [99134.515,99135.1714,99487.1611,99502.2385]:
                    line_y_upper = line_y_upper +30
                for i in range(a+1):
                    if note_[n+i]==mol:
                        line_x = [tag_x_[n+i],tag_x_[n+i]]
                        if line_y_upper>y_up-40:
                            line_y_upper=y_up-40-alpha
                        if tag_x_[n] in [100078.608,100080.542]:
                            line_y_upper = line_y_upper -10
                        line_y = [tag_y_[n+i],line_y_upper]
                        plt.plot(line_x,line_y,color='blue',linewidth=1,linestyle='dashed')
            else:
                line_y_upper = getlineyup(x_,tag_x_,y__,n)+y_up/8
                line_x = [tag_x_[n],tag_x_[n]]
                line_y_upper = getyupper(text_position,line_x,line_y_upper,y_up)
                if tag_x_[n] in [98900.1552,99172.5214,99900.0915,99961.2448]:
                    line_y_upper = line_y_upper+35
                if tag_x_[n] in [97672.0182]:
                    line_y_upper = line_y_upper+80
                if tag_x_[n] in [100673.189]:
                    line_y_upper = line_y_upper-20
                if tag_x_[n] in [98176.293]:
                    line_x = [tag_x_[n],tag_x_[n]+5]
                    
                if line_y_upper>y_up-40:
                    line_y_upper=y_up-40
                if tag_x_[n] in [98600.7203]:
                    line_y_upper = line_y_upper-60
                line_y = [tag_y_[n],line_y_upper]
                plt.plot(line_x,line_y,color='blue',linewidth=1,linestyle='dashed')
                
            p = lastpos(new_tags,tags_[n])
            line_x = [tag_x_[n],tag_x_[n+p]]
            if tag_x_[n] in [98176.293]:
                line_x = [tag_x_[n]+12,tag_x_[n+p]+12]
            if line_y_upper>y_up-40:
                line_y_upper=y_up-40
            line_y = [line_y_upper,line_y_upper]
            plt.plot(line_x,line_y,color='blue',linewidth=1)
            if line_x==[98606.856,98611.163]:
                line_x = [98611.163,98611.163]
            plt.text(np.mean(line_x),line_y[1],tags_[n],fontdict=font3,rotation=angle,horizontalalignment='center', verticalalignment='bottom')
            text_position.append([line_x[1],line_y[1]])
            if p!=0:
                n = n+p+1
            else:
                n = n+1
        else:
            n = n+1
    return(text_position)


angle=90
text_position = []
x_,y_,x__,y__,tags_,note_,tag_x_,tag_y_,_x,_y,__x,__y,_tags,_note,_tag_x,_tag_y = loaddata(data_file,model_file,tran_file)
plt.figure(figsize=(20,10))
plotline(x_, y_,x__, y__, y_down, y_up)
for n in range(len(tags_)):
    text_position = plotsimpletransition(molecules,tag_x_,x_,y__,tag_y_,tags_,note_,y_up,n)

for i in range(len(molecules)):
    mol = molecules[i]
    text_position = plotmultitrans(tags_,mol,tag_x_,x_,y__,note_,text_position,y_up,tag_y_,i*3)
    
plt.savefig('spec_tag_%s_1.pdf' %(spw))
plt.close('all')


plt.figure(figsize=(20,10))
plotline(_x, _y,__x, __y, y_down, y_up)
text_position = []
for n in range(len(_tags)):
    text_position = plotsimpletransition(molecules,_tag_x,_x,__y,_tag_y,_tags,_note,y_up,n)

for i in range(len(molecules)):
    mol = molecules[i]
    text_position = plotmultitrans(_tags,mol,_tag_x,_x,__y,_note,text_position,y_up,_tag_y,i*3)    
plt.savefig('spec_tag_%s_2.pdf' %(spw))

得到下图