""" (C) Steven Byrnes, 2014. This code is released under the MIT license http://opensource.org/licenses/MIT  This code should work in Python 2.7 or 3.3. It requires imagemagick to be installed; that's how it assembles images into animated GIFs. """  from __future__ import division, print_function  import pygame as pg from math import pi from cmath import exp  import subprocess, os directory_now = os.path.dirname(os.path.realpath(__file__))  frames_in_anim = 60 animation_loop_seconds = 2 #time in seconds for animation to loop one cycle  bgcolor = (255,255,255) #white potential_curve_color = (0,0,0) #black ecolor = (100,0,100) #electron is purple  linear_color = (0, 0, 150) shg_color = (0, 150, 0) const_color = (150, 0, 0)  eradius = 20  img_height = 500 img_width = 1000  top_arrow_y = 350 middle_arrow_y = 380 bottom_arrow_y = 410 arrow_width = 8  # Limits of the potential curve xmin = 100 xmax = 900 ymin = 40 ymax = 300  # pygame draws pixel-art, not smoothed. Therefore I am drawing it # bigger, then smoothly shrinking it down final_width = int(round(0.3 * img_width)) final_height = int(round(0.3 * img_height))  def potential_curve(x):     """     My potential curve y as a function of x     """     xscaled = (x-xmin) / (xmax - xmin)     if xscaled < 0.2:         yscaled = (0.2 - xscaled)**2 / (0.2**2)     else:         yscaled = (xscaled - 0.2)**2 / (0.8**2)     # flip it, because higher y-coordinate is lower in pygame drawing     yscaled = 1 - yscaled     return ymin + (ymax - ymin) * yscaled  curve_bottom_x = 0.79 * xmin + 0.21 * xmax curve_bottom_y = potential_curve(curve_bottom_x)  def electron_curve(x):     """     The path that the electron center travels along     """     # xscaled = (x-xmin) / (xmax - xmin)     y = min(potential_curve(x), potential_curve(x+eradius), potential_curve(x-eradius))     return y - eradius  # Constants and function for calculating electron motion linear_coef = 0.3 shg_coef = 0.07 displacement = 0.32  def e_x(phase):     """     x-position of electron as a function of phase (from 0 to 2pi)     """     xscaled = (linear_coef * exp(1j * phase) + shg_coef * exp(2j * phase)                + displacement).real     return xmin + xscaled * (xmax - xmin)  def draw_arrow(surf, tail_xy, head_xy, width=2, color=(0,0,0)):     """     draw a horizontal arrow     """     # tail_xy and head_xy are 2-tuples. Unpack them first     tail_x, tail_y = tail_xy     head_x, head_y = head_xy     assert head_y == tail_y     h = 16 # arrowhead height     b = 18 # arrowhead half-base     if tail_x < head_x:         # rightward arrow         triangle = [(head_x, head_y),                     (head_x - h, head_y - b),                     (head_x - h, head_y + b)]     else:         # leftward arrow         triangle = [(head_x, head_y),                     (head_x + h, head_y - b),                     (head_x + h, head_y + b)]     pg.draw.line(surf, color, (tail_x, tail_y), (head_x, head_y), width)     pg.draw.polygon(surf, color, triangle, 0)  def main():     """ function for creating the animated GIF """     # Make and save a drawing for each frame     filename_list = [os.path.join(directory_now, 'temp' + str(n) + '.png')                          for n in range(frames_in_anim)]      # Put the potential curve in the form of a list of points, to be drawn below     xs = range(xmin, xmax + 1,1)     ys = [potential_curve(x) for x in xs]     potential_curve_path = zip(xs, ys)          for frame in range(frames_in_anim):         phase = 2 * pi * frame / frames_in_anim         electron_x = e_x(phase)         electron_y = electron_curve(electron_x)                  # initialize surface         surf = pg.Surface((img_width,img_height))         surf.fill(bgcolor)                  # draw potential curve         pg.draw.lines(surf, potential_curve_color, False,                       potential_curve_path, 10)                  # draw vertical line to first arrow         pg.draw.line(surf, (0,0,0), (curve_bottom_x,curve_bottom_y),                      (curve_bottom_x, top_arrow_y), 3)                  # draw three arrows         linear_term = (linear_coef * exp(1j * phase)).real * (xmax - xmin)         shg_term = (shg_coef * exp(2j * phase)).real * (xmax - xmin)                  draw_arrow(surf,                    (curve_bottom_x, top_arrow_y),                    (curve_bottom_x + linear_term, top_arrow_y),                    width=arrow_width, color=linear_color)         draw_arrow(surf,                    (curve_bottom_x + linear_term, middle_arrow_y),                    (curve_bottom_x + linear_term + shg_term, middle_arrow_y),                    width=arrow_width, color=shg_color)         draw_arrow(surf,                    (curve_bottom_x + linear_term + shg_term, bottom_arrow_y),                    (electron_x, bottom_arrow_y),                    width=arrow_width, color=const_color)                  # draw electron         pg.draw.circle(surf, ecolor,                        ((int(round(electron_x)), int(round(electron_y)))),                        eradius, 0)          shrunk_surface = pg.transform.smoothscale(surf, (final_width, final_height))         pg.image.save(shrunk_surface, filename_list[frame])      seconds_per_frame = animation_loop_seconds / frames_in_anim     frame_delay = str(int(seconds_per_frame * 100))     command_list = ['convert', '-delay', frame_delay, '-loop', '0'] + filename_list + ['anim.gif']     # Use the "convert" command (part of ImageMagick) to build the animation     subprocess.call(command_list, cwd=directory_now)     # Earlier, we saved an image file for each frame of the animation. Now     # that the animation is assembled, we can (optionally) delete those files     if True:         for filename in filename_list:             os.remove(filename)     return  main()