Permalink
01f939b Mar 27, 2015
@bollu @larsoner
executable file 206 lines (164 sloc) 5.91 KB
#!/usr/bin/env python
import numpy as np
import sys
from vispy.util.transforms import perspective
from vispy.util import transforms
from vispy import gloo
from vispy import app
from vispy import io
import galaxy_specrend
from galaxy_simulation import Galaxy
VERT_SHADER = """
#version 120
uniform mat4 u_model;
uniform mat4 u_view;
uniform mat4 u_projection;
//sampler that maps [0, n] -> color according to blackbody law
uniform sampler1D u_colormap;
//index to sample the colormap at
attribute float a_color_index;
//size of the star
attribute float a_size;
//type
//type 0 - stars
//type 1 - dust
//type 2 - h2a objects
//type 3 - h2a objects
attribute float a_type;
attribute vec2 a_position;
//brightness of the star
attribute float a_brightness;
varying vec3 v_color;
void main (void)
{
gl_Position = u_projection * u_view * u_model * vec4(a_position, 0.0,1.0);
//find base color according to physics from our sampler
vec3 base_color = texture1D(u_colormap, a_color_index).rgb;
//scale it down according to brightness
v_color = base_color * a_brightness;
if (a_size > 2.0)
{
gl_PointSize = a_size;
} else {
gl_PointSize = 0.0;
}
if (a_type == 2) {
v_color *= vec3(2, 1, 1);
}
else if (a_type == 3) {
v_color = vec3(.9);
}
}
"""
FRAG_SHADER = """
#version 120
//star texture
uniform sampler2D u_texture;
//predicted color from black body
varying vec3 v_color;
void main()
{
//amount of intensity from the grayscale star
float star_tex_intensity = texture2D(u_texture, gl_PointCoord).r;
gl_FragColor = vec4(v_color * star_tex_intensity, 0.8);
}
"""
galaxy = Galaxy(10000)
galaxy.reset(13000, 4000, 0.0004, 0.90, 0.90, 0.5, 200, 300)
# coldest and hottest temperatures of out galaxy
t0, t1 = 200.0, 10000.0
# total number of discrete colors between t0 and t1
n = 1000
dt = (t1 - t0) / n
# maps [0, n) -> colors
# generate a linear interpolation of temperatures
# then map the temperatures to colors using black body
# color predictions
colors = np.zeros(n, dtype=(np.float32, 3))
for i in range(n):
temperature = t0 + i * dt
x, y, z = galaxy_specrend.spectrum_to_xyz(galaxy_specrend.bb_spectrum,
temperature)
r, g, b = galaxy_specrend.xyz_to_rgb(galaxy_specrend.SMPTEsystem, x, y, z)
r = min((max(r, 0), 1))
g = min((max(g, 0), 1))
b = min((max(b, 0), 1))
colors[i] = galaxy_specrend.norm_rgb(r, g, b)
# load the PNG that we use to blend the star with
# to provide a circular look to each star.
def load_galaxy_star_image():
fname = io.load_data_file('galaxy/star-particle.png')
raw_image = io.read_png(fname)
return raw_image
class Canvas(app.Canvas):
def __init__(self):
# setup initial width, height
app.Canvas.__init__(self, keys='interactive', size=(800, 600))
# create a new shader program
self.program = gloo.Program(VERT_SHADER, FRAG_SHADER,
count=len(galaxy))
# load the star texture
self.texture = gloo.Texture2D(load_galaxy_star_image(),
interpolation='linear')
self.program['u_texture'] = self.texture
# construct the model, view and projection matrices
self.view = transforms.translate((0, 0, -5))
self.program['u_view'] = self.view
self.model = np.eye(4, dtype=np.float32)
self.program['u_model'] = self.model
self.program['u_colormap'] = colors
w, h = self.size
self.projection = perspective(45.0, w / float(h), 1.0, 1000.0)
self.program['u_projection'] = self.projection
# start the galaxy to some decent point in the future
galaxy.update(100000)
data = self.__create_galaxy_vertex_data()
# setup the VBO once the galaxy vertex data has been setup
# bind the VBO for the first time
self.data_vbo = gloo.VertexBuffer(data)
self.program.bind(self.data_vbo)
# setup blending
gloo.set_state(clear_color=(0.0, 0.0, 0.03, 1.0),
depth_test=False, blend=True,
blend_func=('src_alpha', 'one'))
self._timer = app.Timer('auto', connect=self.update, start=True)
def __create_galaxy_vertex_data(self):
data = np.zeros(len(galaxy),
dtype=[('a_size', np.float32, 1),
('a_position', np.float32, 2),
('a_color_index', np.float32, 1),
('a_brightness', np.float32, 1),
('a_type', np.float32, 1)])
# see shader for parameter explanations
pw, ph = self.physical_size
data['a_size'] = galaxy['size'] * max(pw / 800.0, ph / 800.0)
data['a_position'] = galaxy['position'] / 13000.0
data['a_color_index'] = (galaxy['temperature'] - t0) / (t1 - t0)
data['a_brightness'] = galaxy['brightness']
data['a_type'] = galaxy['type']
return data
def on_resize(self, event):
# setup the new viewport
gloo.set_viewport(0, 0, *event.physical_size)
# recompute the projection matrix
w, h = event.size
self.projection = perspective(45.0, w / float(h),
1.0, 1000.0)
self.program['u_projection'] = self.projection
def on_draw(self, event):
# update the galaxy
galaxy.update(50000) # in years !
# recreate the numpy array that will be sent as the VBO data
data = self.__create_galaxy_vertex_data()
# update the VBO
self.data_vbo.set_data(data)
# bind the VBO to the GL context
self.program.bind(self.data_vbo)
# clear the screen and render
gloo.clear(color=True, depth=True)
self.program.draw('points')
if __name__ == '__main__':
c = Canvas()
c.show()
if sys.flags.interactive == 0:
app.run()