"""
  File:     HRDiagram.py
  Author:   Thomas Mueller, HdA
  
  Description:
  ------------
    Load Hipparcos catalog (binary fits file) and draw a Hertzsprung-Russell
    diagram of absolute versus B-V magnitude.
     
    The catalog data can be downloaded from VizieR 
         <http://vizier.u-strasbg.fr/viz-bin/VizieR>
         
    * Search for "hip" yields "I/234", "I/311", ...
    * Select "I/311/hip2"
    * Set Preferences to 
        - max = unlimited
        - FITS (binary) Table
        - Position in Decimal
    * Check that the following columns are selected:
        "HIP", "Plx", "Hpmag", "B-V"        
    * Press "Submit" to download the data.
    * Rename the file from "asu.fit" to e.g. "I_311_Hipparcos.fit".
"""
     
from astropy.io import fits
import numpy as np
import matplotlib.pyplot as plt

# Hipparcos file
filename = 'HIP-Sterne bis 6m.fit'

# Read fits file
hdulist = fits.open(filename)
hipdata = hdulist[1].data
hdulist.close()

# extract data from Hipparcos catalogue
ra = hipdata.field('RArad')
de = hipdata.field('DErad')
plx = hipdata.field('Plx')
BmV = hipdata.field('B-V')
mag = hipdata.field('Hpmag')

# select only those entries, where the parallax is positive
val = np.where(plx > 0)

# calculate absolute magnitude
AMag = mag[val] + 5 * (1 + np.log10(plx[val] * 0.001))

# draw 
fig = plt.figure(figsize = (9, 12))
ax = fig.add_subplot(111)
ax.set_title("H-R diagram of Hipparcos catalog")
ax.title.set_fontsize(18)

# stars are plotted using scatter-plot
ax.scatter(BmV[val], AMag, s = 25, marker = 'o')

# set x-axis
ax.set_xlim([-0.25, 1.5])
ax.set_xlabel('B-V [mag]')
ax.xaxis.label.set_fontsize(16)

# set y-axis
ax.set_ylim([4.5, -0.5])
ax.set_ylabel('Absolute magnitude M [mag]')
ax.yaxis.label.set_fontsize(16)

# gaphic output
ax.grid()
plt.savefig("HRdiagram.png", dpi=100, bbox_inches='tight', pad_inches = 0)
plt.show()