Skip to content

Instantly share code, notes, and snippets.

@Zulko
Last active June 23, 2020 04:23
Show Gist options
  • Save Zulko/54e5468759396c5cbbd2 to your computer and use it in GitHub Desktop.
Save Zulko/54e5468759396c5cbbd2 to your computer and use it in GitHub Desktop.
Animated atom (Vispy) turned into a video/GIF (MoviePy)
# -*- coding: utf-8 -*-
# vispy: gallery 30
# -----------------------------------------------------------------------------
# Copyright (c) 2014, Vispy Development Team. All Rights Reserved.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.
# -----------------------------------------------------------------------------
# Author: Nicolas P .Rougier
# Date: 06/03/2014
# Abstract: Fake electrons orbiting
# Keywords: Sprites, atom, particles
#
# Modified by Zulko for animation with MoviePy
# See result here : http://i.imgur.com/6PNEYB9.gif 250 %}
# -----------------------------------------------------------------------------
import numpy as np
from vispy import gloo
from vispy import app
from vispy.util.transforms import perspective, translate, rotate
from vispy.gloo.util import _screenshot
# Create vertices
n, p = 100, 150
data = np.zeros(p * n, [('a_position', np.float32, 2),
('a_color', np.float32, 4),
('a_rotation', np.float32, 4)])
trail = .5 * np.pi
data['a_position'][:, 0] = np.resize(np.linspace(0, trail, n), p * n)
data['a_position'][:, 0] += np.repeat(np.random.uniform(0, 2 * np.pi, p), n)
data['a_position'][:, 1] = np.repeat(np.linspace(0, 2 * np.pi, p), n)
data['a_color'] = 1, 1, 1, 1
data['a_color'] = np.repeat(
np.random.uniform(0.75, 1.00, (p, 4)).astype(np.float32), n, axis=0)
data['a_color'][:, 3] = np.resize(np.linspace(0, 1, n), p * n)
data['a_rotation'] = np.repeat(
np.random.uniform(0, 2 * np.pi, (p, 4)).astype(np.float32), n, axis=0)
vert = """
#version 120
uniform mat4 u_model;
uniform mat4 u_view;
uniform mat4 u_projection;
uniform float u_size;
uniform float u_clock;
attribute vec2 a_position;
attribute vec4 a_color;
attribute vec4 a_rotation;
varying vec4 v_color;
mat4 build_rotation(vec3 axis, float angle)
{
axis = normalize(axis);
float s = sin(angle);
float c = cos(angle);
float oc = 1.0 - c;
return mat4(oc * axis.x * axis.x + c,
oc * axis.x * axis.y - axis.z * s,
oc * axis.z * axis.x + axis.y * s,
0.0,
oc * axis.x * axis.y + axis.z * s,
oc * axis.y * axis.y + c,
oc * axis.y * axis.z - axis.x * s,
0.0,
oc * axis.z * axis.x - axis.y * s,
oc * axis.y * axis.z + axis.x * s,
oc * axis.z * axis.z + c,
0.0,
0.0, 0.0, 0.0, 1.0);
}
void main (void) {
v_color = a_color;
float x0 = 1.5;
float z0 = 0.0;
float theta = a_position.x + u_clock;
float x1 = x0*cos(theta) + z0*sin(theta);
float y1 = 0.0;
float z1 = (z0*cos(theta) - x0*sin(theta))/2.0;
mat4 R = build_rotation(a_rotation.xyz, a_rotation.w);
gl_Position = u_projection * u_view * u_model * R * vec4(x1,y1,z1,1);
gl_PointSize = 8.0 * u_size * sqrt(v_color.a);
}
"""
frag = """
#version 120
varying vec4 v_color;
varying float v_size;
void main()
{
float d = 2*(length(gl_PointCoord.xy - vec2(0.5,0.5)));
gl_FragColor = vec4(v_color.rgb, v_color.a*(1-d));
}
"""
class Canvas(app.Canvas):
def __init__(self):
app.Canvas.__init__(self, keys='interactive')
self.size = 800, 800
self.program = gloo.Program(vert, frag)
self.view = np.eye(4, dtype=np.float32)
self.model = np.eye(4, dtype=np.float32)
self.projection = np.eye(4, dtype=np.float32)
self.translate = 4.5
translate(self.view, 0, 0, -self.translate)
self.program.bind(gloo.VertexBuffer(data))
self.program['u_model'] = self.model
self.program['u_view'] = self.view
self.program['u_size'] = 5 / self.translate
gloo.set_state('translucent', depth_test=False)
self.program['u_clock'] = 0.0
def on_resize(self, event):
width, height = event.size
gloo.set_viewport(0, 0, width, height)
self.projection = perspective(45.0, width / float(height), 1.0, 1000.0)
self.program['u_projection'] = self.projection
def animation(self, t):
""" Added for animation with MoviePy """
self.program['u_clock'] = 2*t
gloo.clear('black')
self.program.draw('points')
return _screenshot((0, 0, self.size[0], self.size[1]))[:,:,:3]
if __name__ == '__main__':
from moviepy.editor import VideoClip
canvas = Canvas()
canvas.show()
clip = VideoClip(canvas.animation, duration=np.pi).resize(0.3)
clip.write_videofile('atom3.mp4', fps=20)
#clip.write_gif('atom3.gif', fps=20, opt='OptimizePlus')
@intoro
Copy link

intoro commented Feb 19, 2018

I'm getting this:

File "draw.py", line 136, in
canvas = Canvas()
File "draw.py", line 108, in init
translate(self.view, 0, 0, -self.translate)
TypeError: translate() takes from 1 to 2 positional arguments but 4 were given

I'm using Ubuntu and virtualenv, any suggestions on what to change?

@arpanisi
Copy link

Check this: http://vispy.org/util.html
translate takes the argument offset in the form [x,y,z]

@xinii
Copy link

xinii commented Jan 29, 2019

I'm getting this:

File "draw.py", line 136, in
canvas = Canvas()
File "draw.py", line 108, in init
translate(self.view, 0, 0, -self.translate)
TypeError: translate() takes from 1 to 2 positional arguments but 4 were given

I'm using Ubuntu and virtualenv, any suggestions on what to change?

self.view = translate([0, 0, -self.translate])

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment