playing with metaballs

isosurface.py

'''
Created on Aug 10, 2012

@author: spellrh
'''

import numpy as n
import pygame
from pygame.locals import *
from pygame import surfarray
import math, random

white = 255,255,255
black = 0,0,0

width,height = 300,300




def main():
    pygame.init()
#    screen = pygame.display.set_mode( (width,height) )
    pygame.display.set_caption('Metaballs')
    
    def makeSpot():
        '''
        this one is drawn with a surfArray
        Note that the SurfaceArray progresses down each column of the image(surface)
        and then across the image.  The PixelArray traverses across the screen first
        '''
        center = n.array( [width / 2, height / 2],dtype='int32')
        dmax = n.linalg.norm(center)
        m = n.ones((width,height,3),dtype='int32') * 255
        screen = pygame.display.set_mode( m.shape[:2],0,32)
        for i in range(height):
            for j in range(width):
                z = n.array([j,i])
                d = n.linalg.norm( center - z )
                dnorm = float(d/dmax)
                m[j][i] *= dnorm
        
        surfarray.blit_array(screen, m)
        pygame.display.flip()

    def makeCircle(radius):
        '''
        this one is drawn with a PixelArray
        '''
        width = 500
        height = 400
        rr = radius**2
        screen = pygame.display.set_mode((width,height),0,32)
#        screen.fill((255,0,0))
        pixArray = pygame.PixelArray(screen)
        
        cx = width / 2
        cy = height / 2
        
        
        for i in range(width):
            for j in range(height):
                xx = (i-cx)**2
                yy = (j-cy)**2
                
                if xx + yy == rr:
                    pixArray[i][j] = (255,200,0)
                else:
                    pixArray[i][j] = (0,0,0)
                
        del pixArray
        
        
        
    class MetaBall():
        def __init__(self, x,y,radius):
            self.x = x
            self.y = y
            self.r = float(radius)
            self.speed = random.uniform(-5,5)
            self.vx = random.random()*self.speed
            self.vy = random.random()*self.speed
        
        def compute(self,x,y):
            if self.x == x and self.y == y:
                return self.r
            return self.r/ math.sqrt((self.x-x)**2 + (self.y-y)**2)
        
        def update(self):
            self.x += self.vx *.9
            self.y += self.vy *.9
            
            if self.x < 0 or self.x > width:
                self.vx *= -1
            if self.y < 0 or self.y > height:
                self.vy *= -1
            
            pass
        
    metaballs = []
    count = 4
    for i in range(count):
        metaballs.append( MetaBall(random.random()*width, random.random()*height, 30) )
        
    screen = pygame.display.set_mode((width,height),0,32)
    max_t = 1.01
    min_t = .99
    
    def runMetaBalls():
        pixArray = pygame.PixelArray(screen)
        for y in range(height):
            for x in range(width):
                sum = 0.
                for m in metaballs:
                    sum += m.compute(x,y)
            
                if sum >= min_t and sum <= max_t :
                    pixArray[x][y] = 255,255,255
    
        del pixArray
    
    runMetaBalls()                        
    pygame.display.update()
    background = pygame.Surface(screen.get_size())
    background = background.convert() #makes the background the same format as screen; gives fastests results
    background.fill( (0,0,0) )
    
    
#    makeSpot()
#    makeCircle(50)
    
    while True:
        screen.blit(background, (0,0))
        for event in pygame.event.get():
            if event.type == QUIT:
                return

        for m in metaballs:
            m.update()
                            
        runMetaBalls()                        
        pygame.display.update()
        
if __name__ == '__main__':
    main()
    
    
    
    
    

'''

def main():
    pygame.init()
    screen = pygame.display.set_mode( (width,height) )
    pygame.display.set_caption('Metaballs')
    

    
    center = n.array( [width / 2, height / 2],dtype='int32')
    dmax = n.linalg.norm(center)
    m = n.ones((width,height,3),dtype='int32') * 255
    
    
    for i in range(height):
        for j in range(width):
            z = n.array([j,i])
            d = n.linalg.norm( center - z )
            dnorm = float(d/dmax)
#            m[j][i] *= dnorm
#            m[j][i][0] = int(m[j][i][0])
#            m[j][i][1] = int(m[j][i][1])
#            m[j][i][2] = int(m[j][i][2])
            
#            print "%s; %s; dist = %s ; m = %s" % (z, center, dnorm, m[j][i] )
            

    screen = pygame.display.set_mode( m.shape[:2],0,32)
    surfarray.blit_array(screen, m)
    pygame.display.flip()
#    pygame.display.set_caption("mandlebrot")
    while True:
        for event in pygame.event.get():
            if event.type == QUIT:
                return
    
    
    screen.blit(background, (0,0))
    screen.blit(canvas, (50,50))
    
    pygame.display.flip()
    
    clock = pygame.time.Clock()
    
    while True:
        clock.tick(60)
        canvas.fill(black)
        
        for event in pygame.event.get():
            if event.type == QUIT:
                return
            elif event.type == KEYDOWN and event.key == K_ESCAPE:
                return
            elif event.type == MOUSEBUTTONDOWN:
                pass

'''