Write spectral types across the top of a color-magnitude diagram

gistfile1.txt

import numpy as np
import matplotlib.pyplot as plt
import astropy.io.ascii as at
from scipy.interpolate import interp1d

# Read in Kraus & Hillenbrand for later use
kh = at.read('kraushillenbrand5.dat')
ksub = np.array([-23,-19,-15,-12,-11,-9,-7,-6,-5,-4,-3,-2,-1])
kh_rpmK0 = (kh['Mr'] - 0.035*(kh['Mr']-kh['Mi']) - 0.007 - kh['MK'])
kh_rpmK = kh_rpmK0[ksub]
kh_rpmK[0] += 0.1
kh_spt = kh['SpT'][ksub]

def add_spt(ax, texty=None):
    """Add spectral types across the top of the axis, assuming (r'-K)."""

    kh_color = kh_rpmK

    # Only use the relevant colors given the axis limits
    xlims = ax.get_xlim()
    good = np.where((kh_color>xlims[0]) & (kh_color<xlims[1]))[0]
    use_colors = kh_color[good]
    use_spts = kh_spt[good]

    # Determine the proper text position
    if texty is None:
        ylims = ax.get_ylim()
        texty = ylims[1]*1.05

    # Add the spectral type indicators
    for this_color, this_spt in itertools.izip(use_colors,use_spts):
        ax.text(this_color,texty,this_spt,fontsize="large")

    # Remove the tick marks to reduce confusion
    ax.tick_params(which="both",top=False)

kraushillenbrand5.dat

SpT	Mu	Mg	Mr	Mi	Mz	MJ	MH	MK	Mbol	Teff	Mass
B8	0.32	-0.39	-0.04	0.34	0.62	0.01	0.10	0.11	-1.00	11900	3.8
A0	1.58	0.47	0.72	1.04	1.28	0.54	0.58	0.56	0.30	9520	2.9
A2	2.41	1.22	1.39	1.65	1.87	1.12	1.15	1.12	1.10	8970	2.4
A5	3.14	1.88	1.95	2.15	2.32	1.53	1.52	1.48	1.75	8200	2.0
A7	3.47	2.21	2.23	2.40	2.55	1.75	1.71	1.66	2.08	7580	1.8
F0	3.94	2.77	2.68	2.79	2.90	2.10	2.01	1.96	2.61	7200	1.6
F2	4.23	3.10	2.96	3.04	3.13	2.32	2.20	2.14	2.89	6890	1.5
F5	5.01	3.90	3.68	3.69	3.74	2.85	2.67	2.61	3.61	6440	1.25
F8	5.76	4.60	4.29	4.26	4.28	3.31	3.08	3.01	4.24	6200	1.17
G0	6.09	4.89	4.52	4.44	4.44	3.53	3.27	3.20	4.47	6030	1.11
G2	6.35	5.07	4.65	4.54	4.51	3.64	3.38	3.30	4.60	5860	1.06
G5	6.78	5.40	4.92	4.79	4.74	3.86	3.56	3.48	4.89	5770	1.04
G8	7.55	6.03	5.50	5.32	5.25	4.31	3.95	3.86	5.30	5570	0.98
K0	8.08	6.38	5.77	5.55	5.45	4.49	4.10	4.00	5.69	5250	0.90
K2	8.89	6.94	6.23	5.94	5.80	4.80	4.35	4.24	6.08	4900	0.82
K4	9.90	7.62	6.77	6.40	6.20	5.08	4.56	4.43	6.55	4590	0.75
K5	10.36	7.98	7.03	6.59	6.35	5.20	4.64	4.51	6.68	4350	0.70
K7	11.27	8.59	7.45	6.90	6.58	5.46	4.85	4.70	6.89	4060	0.63
M0	12.46	9.90	8.50	7.83	7.46	6.04	5.37	5.18	7.60	3850	0.59
M1	13.00	10.47	9.00	8.12	7.64	6.33	5.68	5.47	7.97	3680	0.54
M2	13.66	11.36	9.76	8.73	8.15	6.73	6.09	5.86	8.44	3510	0.42
M3	14.55	12.37	10.77	9.44	8.74	7.31	6.68	6.44	9.09	3350	0.29
M4	15.83	13.55	11.99	10.48	9.64	8.10	7.49	7.22	9.92	3180	0.20
M5	17.38	15.22	13.67	11.76	10.71	9.08	8.47	8.16	11.01	3010	0.15
M6	18.71	16.56	14.99	12.98	11.88	10.15	9.50	9.16	12.06	2840	0.12
M7	19.74	17.82	16.21	13.94	12.68	10.76	10.08	9.69	12.70	2720	0.11
M8	21.05	19.40	17.60	14.83	13.21	11.19	10.46	10.03	13.13	2600	0.102
M9	21.72	19.93	18.19	15.38	13.69	11.49	10.73	10.26	13.43	2400	0.088
L0	22.33	20.98	18.48	15.85	14.01	11.76	10.96	10.44	13.69	2200	0.078