[pygame] Just sharing

[Silver <rockachu2@xxxxxxxxx>]

I've been working on my little gravity program, except it really doesn't
apply to gravity anymore. Feel free to use it as you wish.

:
* press Esc to exit, q to almost stop all the particles, mouse to move
the central particle, any other key to add a new particle.
* 3d distances+ interactions
* 'spin'
* mouse attracts particles, scroll to control attraction speed
* uses multiprocessing (though not perfectly. Still working on the
comments I got from here, thanks!)

Note: If you aren't using python 3, uncomment the section in the file so
that the print statements don't raise a syntax error.

import pygame
from pygame import *
import random

import multiprocessing
import time
import math, sys

width = 1920
height = 1080
zheight = 1080
x = 0


## if running py2.7 or earlier, uncomment this:
"""
def print(n):
    print n



"""



def attract(itemlist):             ## primary functin
    mspd = 4
    gravitar = itemlist[1]
    itemlist = itemlist[0]
    exc = 0
    red = 0
    rep = 8
    for i in itemlist:
        if i.id != gravitar.id: #and i.spind == gravitar.spind:
            #print(i.spind)
            #print(gravitar.spind)
            n = True
            if i.spind == gravitar.spind:
                n = False
            
            
            
            if gravitar.mass > i.mass:
                greatermass = gravitar.mass
            else:
                greatermass = i.mass
            xdiff = gravitar.x - i.x  
            ydiff = gravitar.y - i.y
            zdiff = gravitar.z - i.z
            dist = ((xdiff**2) + (ydiff**2) + (zdiff**2))**.5
            if dist < greatermass:
                dist = greatermass
            
            if i.id == -1:
                dist **= .8
                
                
            #landmark
            if i.id == -1:
                acceleration = ((rep*(rep*gravitar.mass))/(dist**2)) / gravitar.mass
            else:
                acceleration = ((rep**3)/(dist**2))
            
            xc = xdiff/dist
            yc = ydiff/dist
            zc = zdiff/dist
            if n:
                if i.id != -1:
                    gravitar.dx += acceleration * xc
                    gravitar.dy += acceleration * yc
                    gravitar.dz += acceleration * zc
                else:
                    gravitar.dx -= acceleration * xc
                    gravitar.dy -= acceleration * yc
                    gravitar.dz -= acceleration * zc
                    
        else:
            exc += 1
    #if exc > 1:
        #print "Exc is" + str(exc)
    #    pass
        #raise ValueError, exc
    """if gravitar.dx > mspd:
        red += gravitar.dx - mspd
        gravitar.dx = mspd
    if gravitar.dx < -mspd:
        red += abs(gravitar.dx) - mspd
        gravitar.dx = -mspd   
    if gravitar.dy > mspd:
        red += gravitar.dy - mspd
        gravitar.dy = mspd
    if gravitar.dy < -mspd:
        red += abs(gravitar.dy) - mspd
        gravitar.dy = -mspd   
        
    if red > 0:
        gravitar.red = red
    else:
        gravitar.red = None"""

    if gravitar.x > width or gravitar.x < 0:
        gravitar.dx *= -.98
        gravitar.x += gravitar.dx*1.2
    else:
        gravitar.x += gravitar.dx
    if gravitar.y > height or gravitar.y < 0:
        gravitar.dy *= -.98
        gravitar.y += gravitar.dy*1.2
    else:
        gravitar.y += gravitar.dy
    if gravitar.z > zheight or gravitar.z < 0:
        gravitar.dz *= -.98
        gravitar.z += gravitar.dz*1.2
    else:
        gravitar.z += gravitar.dz
        
        
    gravitar.mass = ((abs(gravitar.dx) + abs(gravitar.dy) +abs(gravitar.dz))/3.0)
    #gravitar.x += gravitar.dx
    #gravitar.y += gravitar.dy
    
    #gravitar.dx *= .999
    #gravitar.dy *= .999
    gravitar.mass = 10/gravitar.mass
    return gravitar
    
    
    



class Gravitar(object):
    def __init__(self, id):
        self.id = id
        self.red = None
        self.x = random.random()*width
        self.y = random.random()*height
        self.z = random.random()*zheight
        self.dx = random.random()*3
        self.dy = random.random()*3
        self.dz = random.random()*3
        self.mass = 8
        self.spind = random.choice([.5,-.5])

class SGravitar(object):
    def __init__(self):
        self.red = None
        self.id = -1
        self.x = 0
        self.y = 0
        self.z = zheight/2
        self.dx = 0
        self.dy = 0
        self.dz = 0
        self.mass = 20
        self.spind = 1



def genlist():
    n = []
    global x
    for i in range( 2):
        n.append(Gravitar(i))
        x += 1
    
    return n

if __name__ == '__main__':
    
    pygame.init()
    screen = pygame.display.set_mode((width, height), pygame.FULLSCREEN)
    
    screeno = pygame.Surface((screen.get_rect().width, screen.get_rect().height))
    screeno.convert()
    screeno.fill((0,0,0))
    screeno.set_alpha(1, pygame.RLEACCEL)
    # Establish communication queues
    
    
    pool = multiprocessing.Pool(processes = multiprocessing.cpu_count())
    
    # Enqueue jobs
    
    
    # Add a poison pill for each consumer
    result = []
    for i in range(2):
        result.append(Gravitar(i))
        x += 1
    result.append(SGravitar())
    lastx = x
    r = 2.0
    # Start printing results
    n = 0
    while True:
        n += 1
        if not n % 5:
            screen.blit(screeno, (0,0))
        ## blit all here
        if x != lastx:
            print(x)
            lastx = x
        mousepos = pygame.mouse.get_pos()
        
        #for item in result:
            #pygame.draw.circle(screen, (155,155,155), ( int(item.x), int(item.y)), int(item.mass))
        #print result
        
        result = [[result,i] for i in result]
        res = pool.map_async(attract, result)
        result = []
        #print "Calculating frame"
        #print res
        result = res.get()
        
        for i in result:
            if i.id == -1:
                i.x, i.y = mousepos
                i.z = zheight/2
                i.mass = r        
        
        for item in result:
            u = item.z/zheight
            if u > 1:
                u = 1
            elif u < 0:
                u = 0
            pygame.draw.circle(screen, (255*u,255*u,255*u), ( int(item.x), int(item.y)), int(item.mass))
            if item.red:
                pygame.draw.circle(screen, (255,0,0), ( int(item.x), int(item.y)), int(item.red))
        pygame.display.update()
        for event in pygame.event.get():
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_q:
                    for i in result:
                        i.dx = 0.1
                        i.dy = 0.1
                        i.dz = 0.1
                elif event.key == pygame.K_ESCAPE:
                    raise SystemExit
                else:
                    x += 1 
                    result.append(Gravitar(x))
            if event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 4:
                    r -= .1
                if event.button == 5:
                    r += .1
        
        
        
        #print "Mapping..."
        

        
        #result = [[result,i] for i in result]
        #res = pool.map_async(attract, result)
        #result = []
        #print "Calculating frame"
        #print res
        #result = res.get()
        #print result