[Numpy-discussion] how can i optimize this function further more?

[Victor Noagbodji]

hello, i'm fairly new to numpy. i need help with a snow effect done
with pygame. the entire code is below. the performance drops in the
snowfall function. the original c code and a demo can be found here :
http://sol.gfxile.net/gp/ch04.html

as you can see. the original c code went the pixel by pixel way, but i
couldn't do that with pygame. i'm asking here because i think there
might be a way numpy could help. the idea of the snowfall function is
to plot white pixel one line down (while avoiding the green ground,
thus the test array1-array2 == white)

thanks a lot in advance for your help.

ps. sorry for the ugly code. this is the result of several
optimizations, plus help from the pygame mailing list.

import numpy
import cProfile
from math import sin, cos
from random import randint
from sys import exit, stdout

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

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

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

def init(screen_array, green=int(0x007f00)):
  lowest_level = 0
  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)
    lowest_level = max(p, lowest_level)
    for j in range(p, HEIGHT):
      screen_array[i, j] = green

  return lowest_level

def newsnow(screen_array, white=int(0xffffff), density=1):
  new_indices = numpy.array([randint(1, WIDTH-2) for i in xrange(density)])
  screen_array[new_indices, 0] = white

def snowfall(screen_array,
             white=int(0xffffff),
             fallrng=xrange(HEIGHT-2, -1, -1)):
  screen_array = numpy.transpose(screen_array)
  snow_in_next_layer = numpy.zeros(WIDTH, dtype=int)
  for j in fallrng:
    array1 = screen_array[j]
    array2 = screen_array[j+1]
    indices_where_snow_moved_down = numpy.where(array1-array2 == white)[0]
    snow_in_next_layer.fill(0)
    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()

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

  lowest_level = init(screen_array)
  display.update(screen_rect)

  fallrng = xrange(lowest_level-2, -1, -1)

  white = int(0xffffff)
  black = 0
  c = Clock()

  while True:
    newsnow(screen_array, density=8)
    snowfall(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

    c.tick()
    stdout.write('fps: ~%s\r' % round(c.get_fps()))
    stdout.flush()

if __name__ == '__main__':
  cProfile.runctx('main()', globals(), {'main': main}, 'last_gp3_stats')

-- 
paul victor noagbodji