Re: [pygame] is there a way to optimize this code?

[Ian Mallett <geometrian@xxxxxxxxx>]

Cool effect!

Well, short of doing anything drastic--like using the GPU to simulate it, let's try some things, in the order I thought of them.  

First, to establish a baseline, I added a framerate counter and got just over 6 fps.  


-Adding
import psyco
psyco.full()
increased 1 fps

-By far the slowest part of the loop in main() is the copied function from snowfall().  Best to restructure that into NumPy:
screen_array = numpy.transpose(screen_array)

for j in xrange(height-2,-1,-1):
    array1 = screen_array[j]

    array2 = screen_array[j+1]
    indices_where_snow_moved_down = numpy.where(array1-array2==white)[0]

    snow_in_next_layer = numpy.zeros(WIDTH)
    snow_in_next_layer[indices_where_snow_moved_down] = white

    screen_array[j  ] = array1 - snow_in_next_layer
    screen_array[j+1] = array2 + snow_in_next_layer

screen_array = numpy.transpose(screen_array)
Increased 29 fps.  

-shouldn't it be "random.randint(0,WIDTH-1)" instead of "random.randint(1, 638)" under newsnow(...)?

-Adding a density option to newsnow() makes sense too.  

-You could improve the second suggestion for another huge speed bump.  Last observation: it's 1:43AM here, and I'm tired; it's an exercise for the reader :D

HTH,
Ian
try:
    import psyco
    psyco.full()
except:
    print "Psyco will make this marginally faster"
import numpy
import random
import sys
from sys import exit
from math import sin, cos

import pygame
import pygame.event as event
import pygame.display as display

from pygame.locals import *
from pygame import Surface
from pygame.time import get_ticks
from pygame.surfarray import pixels2d, blit_array, make_surface

WIDTH = 640
HEIGHT = 480
RESOLUTION = (WIDTH, HEIGHT)

def init(width, height, screen_array, green=int(0x007f00)):
 for i in xrange(width):
   sins = (sin((i + 3247) * 0.02) * 0.3 +
           sin((i + 2347) * 0.04) * 0.1 +
           sin((i + 4378) * 0.01) * 0.6)
   p = int(sins * 100 + 380)
   for j in range(p, height):
     screen_array[i, j] = green

def newsnow(width, height, screen_array, white=int(0xffffff),density=1):
    for i in xrange(density):
        screen_array[random.randint(0,WIDTH-1),0] = white

def snowfall(width, height, screen_array, white=int(0xffffff), black=0, green=int(0x007f00)):
    #The basic idea is, starting from the bottom, move each row into the next row down.
    
    #array1 is the row.  array2 is the next row down.
    
    #indices_where_snow_moved_down is where there's black below the snow (i.e., the pixel
    #of the row minus the pixel of the next row down exactly equals white.
    
    #Then, a row of zeros, snow_in_next_layer, is created.  For the indices just calculated,
    #the particle in snow_in_next_layer is turned white.

    #This array of moved snow particles then is added to the next row, and subtracted from
    #the first.  Q.E.D.
    screen_array = numpy.transpose(screen_array)
    for j in xrange(height-2,-1,-1):
        array1 = screen_array[j]
        array2 = screen_array[j+1]
        indices_where_snow_moved_down = numpy.where(array1-array2==white)[0]
        snow_in_next_layer = numpy.zeros(WIDTH)
        snow_in_next_layer[indices_where_snow_moved_down] = white
        screen_array[j  ] = array1 - snow_in_next_layer
        screen_array[j+1] = array2 + snow_in_next_layer
    screen_array = numpy.transpose(screen_array)
def main():
    pygame.init()

    width = WIDTH
    height = HEIGHT

    screen_surf = display.set_mode(RESOLUTION)
    screen_rect = screen_surf.get_rect()
    screen_array = pixels2d(screen_surf)

    init(width, height, screen_array)

    white = int(0xffffff)
    black = 0

    display.update(screen_rect)

    Clock = pygame.time.Clock()
    while True:

        newsnow(width, height, screen_array, density=8)
        snowfall(width, height, screen_array)

        display.update(screen_rect)

        for e in event.get():
            type = e.type
            if type == QUIT:
                exit()
            elif type == KEYUP and e.key == K_ESCAPE:
                return

        Clock.tick()
        sys.stdout.write("Framerate: ~"+str(int(round(Clock.get_fps())))+" frames/second               \r")
        sys.stdout.flush()


if __name__ == '__main__': main()