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

[Victor Noagbodji <noagbodjivictor@xxxxxxxxx>]

hi, i'm following SDL tutorial here, http://sol.gfxile.net/gp/ch04.html

this is the pygame equivalent. as you can see the code is not nice
looking. that's because i have tried all the optimization tricks i
know (local names, avoid dot operators, numpy array etc...) the result
is still awfully slow

help.

import numpy
import random
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)):
  for i in xrange(8):
    screen_array[random.randint(1, 638), 0] = white

def snowfall(width, height, screen_array, white=int(0xffffff), black=0):
  for j in xrange(height-2, -1, -1):
    for i in xrange(1, width-1):
      if screen_array[i, j] == white:
        if screen_array[i, j+1] == black:
          screen_array[i, j+1] = white
          screen_array[i, j] = black

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)

  while True:
    tick = get_ticks()

    newsnow(width, height, screen_array)
    #snowfall(width, height, screen_array)

    for j in xrange(height-2, -1, -1):
      for i in xrange(1, width-1):
        if screen_array[i, j] == white:
          if screen_array[i, j+1] == black:
            screen_array[i, j+1] = white
            screen_array[i, j] = black

    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

if __name__ == '__main__': main()

-- 
paul victor noagbodji