pprocess

Annotated examples/PyGmy/ppygmy_queue.py

175:eaaedc7eaf62
2016-12-19 Paul Boddie Added support for counting CPU cores on Mac OS X.
paulb@101 1
#!/usr/bin/env python
paulb@101 2
paulb@101 3
"""
paulb@101 4
An adaptation of pygmy.py ("a rubbish raytracer") employing pprocess
paulb@101 5
functionality in order to take advantage of multiprocessing environments.
paulb@101 6
paulb@101 7
--------
paulb@101 8
paulb@101 9
Copyright (C) 2005 Dave Griffiths
paulb@101 10
Copyright (C) 2006, 2007 Paul Boddie <paul@boddie.org.uk>
paulb@101 11
paulb@101 12
This program is free software; you can redistribute it and/or
paulb@101 13
modify it under the terms of the GNU General Public License
paulb@101 14
as published by the Free Software Foundation; either version 2
paulb@101 15
of the License, or (at your option) any later version.
paulb@101 16
paulb@101 17
This program is distributed in the hope that it will be useful,
paulb@101 18
but WITHOUT ANY WARRANTY; without even the implied warranty of
paulb@101 19
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
paulb@101 20
GNU General Public License for more details.
paulb@101 21
paulb@101 22
You should have received a copy of the GNU General Public License
paulb@101 23
along with this program; if not, write to the Free Software
paulb@101 24
Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
paulb@101 25
"""
paulb@101 26
paulb@101 27
import Image, ImageDraw, random, copy
paulb@101 28
from math import *
paulb@101 29
import pprocess
paulb@101 30
import sys
paulb@101 31
paulb@101 32
def sq(a):
paulb@101 33
    return a*a
paulb@101 34
    
paulb@101 35
class vec:
paulb@101 36
    def __init__(self, x, y, z):
paulb@101 37
        self.x=float(x)
paulb@101 38
        self.y=float(y)
paulb@101 39
        self.z=float(z)
paulb@101 40
paulb@101 41
    def __add__(self,other):
paulb@101 42
        return vec(self.x+other.x,self.y+other.y,self.z+other.z)
paulb@101 43
paulb@101 44
    def __sub__(self,other):
paulb@101 45
        return vec(self.x-other.x,self.y-other.y,self.z-other.z)
paulb@101 46
paulb@101 47
    def __mul__(self,amount):
paulb@101 48
        return vec(self.x*amount,self.y*amount,self.z*amount)
paulb@101 49
paulb@101 50
    def __div__(self,amount):
paulb@101 51
        return vec(self.x/amount,self.y/amount,self.z/amount)
paulb@101 52
paulb@101 53
    def __neg__(self):
paulb@101 54
        return vec(-self.x,-self.y,-self.z)
paulb@101 55
paulb@101 56
    def dot(self,other):
paulb@101 57
        return (self.x*other.x)+(self.y*other.y)+(self.z*other.z)
paulb@101 58
        
paulb@101 59
    def cross(self,other):
paulb@101 60
        return vec(self.y*other.z - self.z*other.y,
paulb@101 61
                    self.z*other.x - self.x*other.z,
paulb@101 62
                    self.x*other.y - self.y*other.x)
paulb@101 63
paulb@101 64
    def dist(self,other):
paulb@101 65
        return sqrt((other.x-self.x)*(other.x-self.x)+
paulb@101 66
                    (other.y-self.y)*(other.y-self.y)+
paulb@101 67
                    (other.z-self.z)*(other.z-self.z))
paulb@101 68
paulb@101 69
    def sq(self):
paulb@101 70
        return sq(self.x)+sq(self.y)+sq(self.z)
paulb@101 71
paulb@101 72
    def mag(self):
paulb@101 73
        return self.dist(vec(0,0,0))
paulb@101 74
paulb@101 75
    def norm(self):
paulb@101 76
        mag=self.mag()
paulb@101 77
        if mag!=0:
paulb@101 78
            self.x=self.x/mag
paulb@101 79
            self.y=self.y/mag
paulb@101 80
            self.z=self.z/mag
paulb@101 81
            
paulb@101 82
    def reflect(self,normal):
paulb@101 83
        vdn=self.dot(normal)*2
paulb@101 84
        return self-normal*vdn
paulb@101 85
paulb@101 86
class line: 
paulb@101 87
    def __init__(self, start, end):
paulb@101 88
        self.start=start
paulb@101 89
        self.end=end
paulb@101 90
            
paulb@101 91
    def vec(self):
paulb@101 92
        return self.end-self.start
paulb@101 93
    
paulb@101 94
    def closestpoint(self, point):
paulb@101 95
        l=self.end-self.start
paulb@101 96
        l2=point-self.start        
paulb@101 97
        t=l.dot(l2)
paulb@101 98
        if t<=0: return self.start
paulb@101 99
        if t>l.mag(): return self.end
paulb@101 100
        return self.start+l*t
paulb@101 101
    
paulb@101 102
class renderobject:
paulb@101 103
    def __init__(self, shader):
paulb@101 104
        self.shader=shader        
paulb@101 105
    
paulb@101 106
    def intersect(self,l):
paulb@101 107
        return "none",vec(0,0,0),vec(0,0,0) # type, position, normal
paulb@101 108
paulb@101 109
class plane(renderobject):
paulb@101 110
    def __init__(self,plane,dist,shader):
paulb@101 111
        renderobject.__init__(self,shader)
paulb@101 112
        self.plane=plane
paulb@101 113
        self.dist=dist
paulb@101 114
        
paulb@101 115
    def intersect(self,l):        
paulb@101 116
        vd=self.plane.dot(l.vec()) 
paulb@101 117
        if vd==0: return "none",vec(0,0,0),vec(0,0,0)
paulb@101 118
        v0 = -(self.plane.dot(l.start)+self.dist)
paulb@101 119
        t = v0/vd
paulb@101 120
        if t<0 or t>1: return "none",vec(0,0,0),vec(0,0,0)
paulb@101 121
        return "one",l.start+(l.vec()*t),self.plane
paulb@101 122
        
paulb@101 123
    
paulb@101 124
class sphere(renderobject):
paulb@101 125
    def __init__(self, pos, radius, shader):
paulb@101 126
        renderobject.__init__(self,shader)
paulb@101 127
        self.pos=pos
paulb@101 128
        self.radius=radius
paulb@101 129
paulb@101 130
    def disttoline(self,l):
paulb@101 131
        return self.pos.dist(l.closestpoint(self.pos))
paulb@101 132
    
paulb@101 133
    def intersect(self,l):
paulb@101 134
        lvec=l.vec()
paulb@101 135
        a = sq(lvec.x)+sq(lvec.y)+sq(lvec.z)
paulb@101 136
          
paulb@101 137
        b = 2*(lvec.x*(l.start.x-self.pos.x)+ \
paulb@101 138
               lvec.y*(l.start.y-self.pos.y)+ \
paulb@101 139
               lvec.z*(l.start.z-self.pos.z))
paulb@101 140
               
paulb@101 141
        c = self.pos.sq()+l.start.sq() - \
paulb@101 142
             2*(self.pos.x*l.start.x+self.pos.y*l.start.y+self.pos.z*l.start.z)-sq(self.radius)
paulb@101 143
    
paulb@101 144
        i = b*b-4*a*c 
paulb@101 145
paulb@101 146
        intersectiontype="none"
paulb@101 147
        pos=vec(0,0,0)
paulb@101 148
        norm=vec(0,0,0)
paulb@101 149
        t=0
paulb@101 150
        
paulb@101 151
        if i>0 :    
paulb@101 152
            if i==0:
paulb@101 153
                intersectiontype="one"
paulb@101 154
                t = -b/(2*a);
paulb@101 155
            else:  
paulb@101 156
                intersectiontype="two"
paulb@101 157
                t = (-b - sqrt( b*b - 4*a*c )) / (2*a)
paulb@101 158
                # just bother with one for the moment
paulb@101 159
                # t2= (-b + sqrt( b*b - 4*a*c )) / (2*a) 
paulb@101 160
paulb@101 161
            if t>0 and t<1: 
paulb@101 162
                pos = l.start+lvec*t
paulb@101 163
                norm=pos-self.pos 
paulb@101 164
                norm.norm()
paulb@101 165
            else: 
paulb@101 166
                intersectiontype="none"
paulb@101 167
                
paulb@101 168
        return intersectiontype,pos,norm
paulb@101 169
        
paulb@101 170
    def intersects(self,l):
paulb@101 171
        return self.disttoline(l)<self.radius
paulb@101 172
paulb@101 173
class light:
paulb@101 174
    def __init__(self):
paulb@101 175
        pass
paulb@101 176
        
paulb@101 177
    def checkshadow(self, obj, objects,l):
paulb@101 178
        # shadowing built into the lights (is this right?)
paulb@101 179
        for ob in objects:
paulb@101 180
            if ob is not obj:
paulb@101 181
                intersects,pos,norm = ob.intersect(l)
paulb@101 182
                if intersects is not "none":
paulb@101 183
                    return 1
paulb@101 184
        return 0
paulb@101 185
        
paulb@101 186
    def light(self, obj, objects, pos, normal):
paulb@101 187
        pass
paulb@101 188
paulb@101 189
class parallellight(light):
paulb@101 190
    def __init__(self, direction, col):
paulb@101 191
        direction.norm()
paulb@101 192
        self.direction=direction
paulb@101 193
        self.col=col
paulb@101 194
paulb@101 195
    def inshadow(self, obj, objects, pos):        
paulb@101 196
        # create a longish line towards the light
paulb@101 197
        l = line(pos,pos+self.direction*1000) 
paulb@101 198
        return self.checkshadow(obj,objects,l)
paulb@101 199
paulb@101 200
    def light(self, shaderinfo):
paulb@101 201
        if self.inshadow(shaderinfo["thisobj"],shaderinfo["objects"],shaderinfo["position"]): return vec(0,0,0)
paulb@101 202
        return self.col*self.direction.dot(shaderinfo["normal"])
paulb@101 203
paulb@101 204
class pointlight(light):
paulb@101 205
    def __init__(self, position, col):
paulb@101 206
        self.position=position
paulb@101 207
        self.col=col
paulb@101 208
    
paulb@101 209
    def inshadow(self, obj, objects, pos):
paulb@101 210
        l = line(pos,self.position)
paulb@101 211
        return self.checkshadow(obj,objects,l)
paulb@101 212
            
paulb@101 213
    def light(self, shaderinfo):
paulb@101 214
        if self.inshadow(shaderinfo["thisobj"],shaderinfo["objects"],shaderinfo["position"]): return vec(0,0,0)
paulb@101 215
        direction = shaderinfo["position"]-self.position;
paulb@101 216
        direction.norm()
paulb@101 217
        direction=-direction
paulb@101 218
        return self.col*direction.dot(shaderinfo["normal"])
paulb@101 219
paulb@101 220
class shader:
paulb@101 221
    def __init__(self):
paulb@101 222
        pass
paulb@101 223
paulb@101 224
    # a load of helper functions for shaders, need much improvement
paulb@101 225
paulb@101 226
    def getreflected(self,shaderinfo):
paulb@101 227
        depth=shaderinfo["depth"]
paulb@101 228
        col=vec(0,0,0)
paulb@101 229
        if depth>0:
paulb@101 230
            lray=copy.copy(shaderinfo["ray"])
paulb@101 231
            ray=lray.vec()
paulb@101 232
            normal=copy.copy(shaderinfo["normal"])
paulb@101 233
            ray=ray.reflect(normal)
paulb@101 234
            reflected=line(shaderinfo["position"],shaderinfo["position"]+ray)
paulb@101 235
            obj=shaderinfo["thisobj"]
paulb@101 236
            objects=shaderinfo["objects"]
paulb@101 237
            newshaderinfo = copy.copy(shaderinfo)
paulb@101 238
            newshaderinfo["ray"]=reflected
paulb@101 239
            newshaderinfo["depth"]=depth-1
paulb@101 240
            # todo - depth test
paulb@101 241
            for ob in objects:
paulb@101 242
                if ob is not obj:
paulb@101 243
                    intersects,position,normal = ob.intersect(reflected)
paulb@101 244
                    if intersects is not "none":
paulb@101 245
                        newshaderinfo["thisobj"]=ob
paulb@101 246
                        newshaderinfo["position"]=position
paulb@101 247
                        newshaderinfo["normal"]=normal
paulb@101 248
                        col=col+ob.shader.shade(newshaderinfo)
paulb@101 249
        return col
paulb@101 250
paulb@101 251
    def isoccluded(self,ray,shaderinfo):
paulb@101 252
        dist=ray.mag()
paulb@101 253
        test=line(shaderinfo["position"],shaderinfo["position"]+ray)
paulb@101 254
        obj=shaderinfo["thisobj"]
paulb@101 255
        objects=shaderinfo["objects"]
paulb@101 256
        # todo - depth test
paulb@101 257
        for ob in objects:
paulb@101 258
            if ob is not obj:
paulb@101 259
                intersects,position,normal = ob.intersect(test)
paulb@101 260
                if intersects is not "none":
paulb@101 261
                    return 1
paulb@101 262
        return 0
paulb@101 263
paulb@101 264
    def doocclusion(self,samples,shaderinfo):
paulb@101 265
        # not really very scientific, or good in any way...
paulb@101 266
        oc=0.0
paulb@101 267
        for i in range(0,samples):
paulb@101 268
            ray=vec(random.randrange(-100,100),random.randrange(-100,100),random.randrange(-100,100))
paulb@101 269
            ray.norm()
paulb@101 270
            ray=ray*2.5
paulb@101 271
            if self.isoccluded(ray,shaderinfo):
paulb@101 272
                oc=oc+1
paulb@101 273
        oc=oc/float(samples)
paulb@101 274
        return 1-oc
paulb@101 275
paulb@101 276
    def getcolour(self,ray,shaderinfo):
paulb@101 277
        depth=shaderinfo["depth"]
paulb@101 278
        col=vec(0,0,0)
paulb@101 279
        if depth>0:
paulb@101 280
            test=line(shaderinfo["position"],shaderinfo["position"]+ray)
paulb@101 281
            obj=shaderinfo["thisobj"]
paulb@101 282
            objects=shaderinfo["objects"]
paulb@101 283
            newshaderinfo = copy.copy(shaderinfo)
paulb@101 284
            newshaderinfo["ray"]=test
paulb@101 285
            newshaderinfo["depth"]=depth-1
paulb@101 286
            # todo - depth test
paulb@101 287
            for ob in objects:
paulb@101 288
                if ob is not obj:
paulb@101 289
                    intersects,position,normal = ob.intersect(test)
paulb@101 290
                    if intersects is not "none":
paulb@101 291
                        newshaderinfo["thisobj"]=ob
paulb@101 292
                        newshaderinfo["position"]=position
paulb@101 293
                        newshaderinfo["normal"]=normal
paulb@101 294
                        col=col+ob.shader.shade(newshaderinfo)
paulb@101 295
        return col
paulb@101 296
paulb@101 297
    def docolourbleed(self,samples,shaderinfo):
paulb@101 298
        # not really very scientific, or good in any way...
paulb@101 299
        col=vec(0,0,0)
paulb@101 300
        for i in range(0,samples):
paulb@101 301
            ray=vec(random.randrange(-100,100),random.randrange(-100,100),random.randrange(-100,100))
paulb@101 302
            ray.norm()
paulb@101 303
            ray=ray*5
paulb@101 304
            col=col+self.getcolour(ray,shaderinfo)
paulb@101 305
        col=col/float(samples)
paulb@101 306
        return col
paulb@101 307
paulb@101 308
    def shade(self,shaderinfo):
paulb@101 309
        col=vec(0,0,0)
paulb@101 310
        for lite in shaderinfo["lights"]:
paulb@101 311
            col=col+lite.light(shaderinfo)
paulb@101 312
        return col
paulb@101 313
paulb@101 314
class world:
paulb@101 315
    def __init__(self,width,height):
paulb@101 316
        self.lights=[]
paulb@101 317
        self.objects=[]
paulb@101 318
        self.cameratype="persp"
paulb@101 319
        self.width=width
paulb@101 320
        self.height=height
paulb@101 321
        self.backplane=2000.0
paulb@101 322
        self.imageplane=5.0
paulb@101 323
        self.aspect=self.width/float(self.height)
paulb@101 324
paulb@101 325
    def render_row(self, channel, sy):
paulb@101 326
paulb@101 327
        """
paulb@101 328
        Render the given row, using the 'channel' provided to communicate
paulb@101 329
        result data back to the coordinating process, and using 'sy' as the row
paulb@101 330
        position. A tuple containing 'sy' and a list of result numbers is
paulb@101 331
        returned by this function via the given 'channel'.
paulb@101 332
        """
paulb@101 333
paulb@101 334
        row = []
paulb@101 335
        for sx in range(0,self.width):
paulb@101 336
            x=2*(0.5-sx/float(self.width))*self.aspect
paulb@101 337
            y=2*(0.5-sy/float(self.height))
paulb@101 338
            if self.cameratype=="ortho":
paulb@101 339
                ray = line(vec(x,y,0),vec(x,y,self.backplane))
paulb@101 340
            else:
paulb@101 341
                ray = line(vec(0,0,0),vec(x,y,self.imageplane))
paulb@101 342
                ray.end=ray.end*self.backplane
paulb@101 343
paulb@101 344
            col=vec(0,0,0)
paulb@101 345
            depth=self.backplane
paulb@101 346
            shaderinfo={"ray":ray,"lights":self.lights,"objects":self.objects,"depth":2}
paulb@101 347
paulb@101 348
            for obj in self.objects:                            
paulb@101 349
                intersects,position,normal = obj.intersect(ray)
paulb@101 350
                if intersects is not "none":
paulb@101 351
                    if position.z<depth and position.z>0:
paulb@101 352
                        depth=position.z
paulb@101 353
                        shaderinfo["thisobj"]=obj
paulb@101 354
                        shaderinfo["position"]=position
paulb@101 355
                        shaderinfo["normal"]=normal
paulb@101 356
                        col=obj.shader.shade(shaderinfo)
paulb@101 357
            row.append(col)
paulb@101 358
paulb@101 359
        channel.send((sy, row))
paulb@101 360
paulb@101 361
    def render(self, filename, limit):
paulb@101 362
paulb@101 363
        """
paulb@101 364
        Render the image with many processes, saving it to 'filename', using the
paulb@101 365
        given process 'limit' to constrain the number of processes used.
paulb@101 366
        """
paulb@101 367
paulb@101 368
        image = Image.new("RGB", (self.width, self.height))
paulb@101 369
        draw = ImageDraw.Draw(image)
paulb@101 370
        total = self.width * self.height
paulb@101 371
        count = 0
paulb@101 372
paulb@101 373
        queue = pprocess.Queue(limit=limit)
paulb@101 374
        render_row = queue.manage(self.render_row)
paulb@101 375
paulb@101 376
        for y in range(0, self.height):
paulb@101 377
            render_row(y)
paulb@101 378
paulb@101 379
        for sy, row in queue:
paulb@101 380
            for sx, col in enumerate(row):
paulb@101 381
                draw.point((sx,sy),fill=(col.x*255,col.y*255,col.z*255))
paulb@101 382
                count = count + 1
paulb@101 383
paulb@101 384
            percent = int((count/float(total))*100)
paulb@101 385
            sys.stdout.write(("\010" * 9) + "%3d%% %3d" % (percent, sy))
paulb@101 386
            sys.stdout.flush()
paulb@101 387
paulb@101 388
        image.save(filename)
paulb@101 389
paulb@101 390
# vim: tabstop=4 expandtab shiftwidth=4