Blender/Python: Plot stars from the HYG data

Plots stars on the inside of a sphere using three different types of solids. Most stars are tetrahedrons, brighter stars are octahedrons and the brightest stars are plotted as dodecahedrons. ps. feel free to use this however you want but please don't republish the code as I don't want multiple copies all over the internet as I am likely to update this script every now and then... oh and if you make a million dollars out of it somehow, be sure to donate some of that money to the needy. If I ever find a file hosting service that I like, I'll post the completed .blend file there to save the trouble of running it seeing how the whole catalog can take six or more hours to run...

#!/usr/local/bin/python
#
# $Header: /home/zero/Blender/SOL-Centered-Starsphere-solid/RCS/starsphere-sol-solid-010.py,v 1.21 2015/01/31 03:35:33 zero Exp zero $
#
# PURPOSE:
#
# Script to build up a starfield using the data from the HYG catalogue
#
# DATA SOURCE:
#
# http://astronexus.com/node/34
#
# GET data file from http://astronexus.com/files/downloads/hygxyz.csv.gz
#
# DESCRIPTION:
#
# Stars are plotted on the "inside" of an imaginary sphere. Most stars
# are plotted as tetrahedrons , brighter stars as octahedrons and brightest
# stars as dodecahedrons. Stars are plotted according to their spectral type
# with each spectral type being a seperate object
#
# USAGE:
#
# I run this on my FreeBSD machine with 
#
# unix:~ $ blender -noaudio -P starsphere-sol-solid.py >& OUT < /dev/null &
#
# then once it has finished, (it can take many hours to run) I save the whole
# thing as a .blend file and then I *append* it into Blender via the
# "Shift-F1" feature (Menu: File->Append)
#
# I see no reason why this wouldn't run on other Unix variants, but if you are 
# going to use this on Windows then I've got no idea if it will work for you!
#
# NOTE that blender seems to take quite a while to exit after generating
# the field
import os  
import sys
from mathutils import *; from math import *
import re
import bpy
import traceback

#----------------------------------------------------------------------------
#
# Alter this to point to where you have put your downloaded copy of the Gzipped data
#
hygxyz_file = '/home/zero/Downloads/hygxyz.csv.gz'

#
# full path name of the appropriate unzipper to decompress the data file
#
dezipper = '/usr/bin/gzcat'

#----------------------------------------------------------------------------
# turn on debugging output
debug = 1

# anything fainter than this is skipped completely 
# this is used in the main loop and has no effect on the star size
# NOTE: any value above 20 will give you to *whole* 120,000+ stars
# from the HYG catalog and will take something like six hours to run.
# for testing it out, set the value to 5.0 or so and that will only plot
# the brightest stars
# skip_mag = 5.0 # Takes about twenty seconds
# skip_mag = 8.0 # Takes about thirty-five minutes
skip_mag = 8.0

# "base" minimum star brightness, this is used for size of star calculations.
# set to 20 to as brown dwarf stars tend to be this bright (absolute magnitude)
faintest = 20

# Plot all stars in a sphere so all stars are the same distance away
# from the camera, requires use of Visual magnitude
use_sphere = 1

# Stars will be plotted on the inside "surface" of an imarginary sphere
# with a radius this big
sphere_size = 2000

#----------------------------------------------------------------------------
# create mesh using visual magnitude
def create_star_solid(mesh,mag):
    # A tetrahedron
    # adjust for Sirius (the brightest star) it will be factor=0
    factor = mag + 1.44
    # now reverse the scale of size, Sirius will be of factor=2.2
    factor = (faintest - factor) / 10
    # now exaggerate the difference
    # TODO (should be the 2.5 times brighter for each magnitude)
    factor = factor * factor * factor

    # most stars will be a tetrahedron
    if   mag > 6: verts, faces = calculate_tetrahedron(factor)
    # brighter stars will be a octahedron
    elif mag > 3: verts, faces = calculate_octahedron(factor)
    # brightest stars will be a dodecahedron
    else:         verts, faces = calculate_dodecahedron(factor)

    mesh.from_pydata(verts, [], faces)
    mesh.update()

    return

# solid calculations are based on bits of code from blender itself. 
# under FreeBSD the file is located here...
# /usr/local/share/blender/scripts/addons/add_mesh_extra_objects/add_mesh_solid.py
def calculate_tetrahedron(factor):
    s = (sqrt(2)/3.0) * factor
    t = (-1/3) * factor
    u = (sqrt(6)/3) * factor
    v = 1 * factor
    verts = [(0,0,v),(2*s,0,t),(-s,u,t),(-s,-u,t)]
    faces = [[0,1,2],[0,2,3],[0,3,1],[1,3,2]]
    return verts, faces

def calculate_octahedron(factor):
    v = 1 * factor
    verts = [(v,0,0),(-v,0,0),(0,v,0),(0,-v,0),(0,0,v),(0,0,-v)]
    faces = [[4,0,2],[4,2,1],[4,1,3],[4,3,0],[5,2,0],[5,1,2],[5,3,1],[5,0,3]]
    return verts, faces

def calculate_dodecahedron(factor):
    s = 1/sqrt(3) * factor
    t = sqrt((3-sqrt(5))/6) * factor
    u = sqrt((3+sqrt(5))/6) * factor
    verts = [(s,s,s),(s,s,-s),(s,-s,s),(s,-s,-s),(-s,s,s),(-s,s,-s),(-s,-s,s),(-s,-s,-s),
         (t,u,0),(-t,u,0),(t,-u,0),(-t,-u,0),(u,0,t),(u,0,-t),(-u,0,t),(-u,0,-t),(0,t,u),
         (0,-t,u),(0,t,-u),(0,-t,-u)]
    faces = [[0,8,9,4,16],[0,12,13,1,8],[0,16,17,2,12],[8,1,18,5,9],[12,2,10,3,13],
             [16,4,14,6,17],[9,5,15,14,4],[6,11,10,2,17],[3,19,18,1,13],[7,15,5,18,19],
             [7,11,6,14,15],[7,19,3,10,11]]
    return verts, faces

#----------------------------------------------------------------------------
# extract the first thing that looks like a spectral type,
# ie. for stars that have multiple spectral types listed, the first is used
def get_spectral_type(spectrum):
    m = spectrum_regex.match(spectrum)
    if m:
        spectral_type = m.group(1)
    else:
        if debug > 1: print('unhandled spectral type "{}"'.format(spectrum))
        # if we can't make sense of the spectral type we are given, then
        # set the colour to be white instead of discarding it as unclassified
        spectral_type = 'A0'
    return spectral_type

#----------------------------------------------------------------------------
# create object under a parent
def create_star_object(parent,spectral_type,mag,x,y,z,RA,DEC):

    # create a mesh for this star
    star_mesh = bpy.data.meshes.new(spectral_type+'Mesh')
    # now create a new object for that mesh
    star_object = bpy.data.objects.new(spectral_type, star_mesh)

    # assign the parent
    star_object.parent = parent

    # set location
    star_object.location = (x, y, z)

    # link object to the scene
    bpy.context.scene.objects.link(star_object)

    # create less complex objects for dimmer stars
    create_star_solid(star_mesh, mag)

    # set the object origin, assumes the cursor is at 0,0,0
    bpy.ops.object.origin_set(type='ORIGIN_CURSOR')

    # point the Z axis at origin - this is done so that the tetrahedron
    # has the maximum area pointing at the origin so that the surface
    # area is the same no matter where the star was plotted. If this wasn't
    # done, half of the tetrahedron stars would be "side-on" and would appear
    # slightly dimmer than they should be
    star_object.matrix_world *= Matrix.Rotation(radians(RA*15), 4, "Z")
    star_object.matrix_world *= Matrix.Rotation(radians((DEC+90)*-1), 4, "Y")

    spectral_type_material = bpy.data.materials.get(spectral_type)

    # Set object material
    if spectral_type_material is None:
        spectral_type_material = bpy.data.materials.new(spectral_type)
        star_object.data.materials.append(spectral_type_material)

        # Extract the spectral Letter and Number
        m = spectral_type_regex.match(spectral_type)

        # spectral Letter
        s_l = m.group(1)
        # spectral Number
        s_n = m.group(2)

        if debug > 5: print('s_l={} s_n={}'.format(s_l, s_n))

        # bit of a dogdy algorithm to set the colour but it's good enough for now
        if   m_regex.search(s_l): diffuse = ( 1.00, 0.00 + (int(s_n) / 20.0), 0.00 ) 
        elif k_regex.search(s_l): diffuse = ( 1.00, 0.50 + (int(s_n) / 20.0), 0.00 )
        elif g_regex.search(s_l): diffuse = ( 1.00, 1.00, 0.00 + (int(s_n) / 20.0) )
        elif f_regex.search(s_l): diffuse = ( 1.00, 1.00, 0.50 + (int(s_n) / 20.0) )
        elif a_regex.search(s_l): diffuse = ( 1.00 - (int(s_n) / 40.0), 1.00 - (int(s_n) / 40.0), 1.00 )
        elif b_regex.search(s_l): diffuse = ( 0.75 - (int(s_n) / 40.0), 0.75 - (int(s_n) / 40.0), 1.00 )
        elif o_regex.search(s_l): diffuse = ( 0.50 - (int(s_n) / 80.0), 0.50 - (int(s_n) / 80.0), 1.00 )
        else: diffuse = ( 1.00, 1.00, 1.00 )

        # use the colour we established above
        spectral_type_material.diffuse_color = diffuse
        # specular colour isn't used but set it anyway
        spectral_type_material.specular_color = diffuse

        # turn on emit
        spectral_type_material.emit = 1

        # set various material properties that we aren't really interested
        # in ie. we don't have want to use the specular and we also
        # don't want these stars to be casting shadows, nor receiving them
        spectral_type_material.specular_intensity = 0
        spectral_type_material.ambient = 0
        spectral_type_material.use_shadows = 0
        #not sure about this one#spectral_type_material.use_raytrace = 0
        spectral_type_material.use_ray_shadow_bias = 0
        spectral_type_material.use_mist = 0
        spectral_type_material.use_cast_approximate = 0
        spectral_type_material.use_cast_buffer_shadows = 0
        spectral_type_material.use_cast_shadows = 0
    else:
        star_object.data.materials.append(spectral_type_material)

    # If this is a subsequent star in a spectral type, join it to the first.
    # All master star objects will be two characters long G0 B6 K4 etc etc
    # Subsequent star objects will be A0.001 B6.001 etc etc so only the first
    # star created in a spectral type will have a two character name
    if len(star_object.name) != 2:
        # deselect all
        bpy.ops.object.select_all(action='DESELECT')

        # Select the current star
        star_object.select = True

        # Select the master star object
        bpy.data.objects[spectral_type].select = True

        # Make the master star the active object
        bpy.context.scene.objects.active = bpy.data.objects[spectral_type]

        # now join them
        bpy.ops.object.join()

    return


#-------------------------------------------------------------------
# calculate X Y Z co-ordinates on the inside "surface" of a sphere
# algorithm from http://www.projectrho.com/public_html/starmaps/trigonometry.php
def plot_star_on_sphere(starfield_parent, spectral_type, mag, RA, DEC):
    PHI = radians(RA * 15)
    THETA = radians(DEC)
    RHO = sphere_size

    if debug > 4:
        print('RHO={} THETA={} PHI={} Mag={} Spectrum={} RA={} DEC={}'.format(
            RHO, THETA, PHI, mag, spectrum, RA, DEC))

    RVECT = RHO * cos(THETA)
    x = RVECT * cos(PHI)
    y = RVECT * sin(PHI)
    z = RHO * sin(THETA)

    # Add the star solid
    create_star_object(starfield_parent, spectral_type, mag, x, y, z, RA, DEC)

    return


#-------------------------------------------------------------------------------
#---------------main------------------------------------------------------------
#-------------------------------------------------------------------------------
#
# Compile up the regexes (avoids compiling the regexes for each star plotted)
#
m_regex = re.compile('^[mM]')
k_regex = re.compile('^[kK]')
g_regex = re.compile('^[gG]')
f_regex = re.compile('^[fF]')
a_regex = re.compile('^[aA]')
b_regex = re.compile('^[bB]')
o_regex = re.compile('^[oO]')
spectrum_regex = re.compile("([OBAFGKM]\d)")
spectral_type_regex = re.compile("([OBAFGKM])(\d)")

# Delete the default objects
# deselect all
bpy.ops.object.select_all(action='DESELECT')
# select all the default objects
bpy.data.objects['Cube'].select = True
bpy.data.objects['Lamp'].select = True
bpy.data.objects['Camera'].select = True
# and delete them
bpy.ops.object.delete()

# Make sky completely black
bpy.data.worlds['World'].horizon_color = ( 0, 0, 0 )

#-------------------------------------------------------------------------------
# Add our own camera
bpy.ops.object.camera_add(view_align=False, location=(0, 0, 0), rotation=(radians(78), 0, 0))
# get the camera object
starcam = bpy.context.object
# set the camera object name
starcam.name = "STARCAM"
# set the camera clip to be far, far away
bpy.data.cameras[-1].clip_end=9999
# set the size of the camera to be really wide
bpy.data.cameras[-1].lens=22

#-------------------------------------------------------------------------------
# Add the Star objects' parent
bpy.ops.object.add(type='EMPTY')
# set object mode
bpy.ops.object.mode_set(mode='OBJECT')
# get object
starfield_parent = bpy.context.object
# set object name
starfield_parent.name = "StarSphere"

#-MAIN-LOOP---------------------------------------------------------------------
# StarID,HIP,HD,HR,Gliese,BayerFlamsteed,ProperName,RA,DEC,Distance,PMRA,PMDEC,RV,Mag,AbsMag,Spectrum,ColorIndex,X,Y,Z,VX,VY,VZ
# read data file
for myline in os.popen(dezipper + ' ' + hygxyz_file):
    if debug > 8:
        sys.stdout.write(myline)

    # split the line into an array 'f' on the ',' field seperator
    f = myline.split(',')

    # apparently this is how you do things in python
    # this will attempt to parse the line
    # if for some reason it fails, it goes down to the bottom
    try:
        propername = str(f[6])
        # Skip the title record
        if propername == 'ProperName': continue
        # Skip the sun, it's supposed to be in the middle anyway
        if propername == 'Sol': continue
        bayerflamsteed = str(f[5])
        RA = float(f[7])
        DEC = float(f[8])

        mag = float(f[13])

        absmag = float(f[14])

        # skip out if the star isn't very bright (for debugging)
        #if debug > 5 and mag > 1.0: continue
        if mag > skip_mag: continue
        spectrum = f[15]

        spectral_type = get_spectral_type(spectrum)

        # skip out if we don't have a type - shouldn't happen, probably dud data
        if spectral_type == '': continue

        x = float(f[17])
        y = float(f[18])
        z = float(f[19])

        if use_sphere > 0: plot_star_on_sphere(starfield_parent, spectral_type, mag, RA, DEC)
        else:
            if debug > 1:
                print('Star rejected ProperName={} Mag={} AbsMag={} Spectrum={} x={} y={} z={}'.format(
                    propername, mag, absmag, spectrum, x, y, z))


        pass
    # this is called when the line is a dud
    except:
        if debug: print("something failed: input line was ", myline)
        traceback.print_exc()
        # Print some debugging info upon uncontrolled exit
        for i in bpy.data.meshes: print("MeshName=="+i.name)
        # exit completely if the debug setting is high
        if debug > 5: sys.exit()
        else: continue

# set object mode again
bpy.ops.object.mode_set(mode='OBJECT')