pprocess

Annotated examples/PyGmy/ppygmy_points.py

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