scripted/simpleEmitter.py

#-
# ==========================================================================
# Copyright (C) 1995 - 2006 Autodesk, Inc. and/or its licensors.  All 
# rights reserved.
#
# The coded instructions, statements, computer programs, and/or related 
# material (collectively the "Data") in these files contain unpublished 
# information proprietary to Autodesk, Inc. ("Autodesk") and/or its 
# licensors, which is protected by U.S. and Canadian federal copyright 
# law and by international treaties.
#
# The Data is provided for use exclusively by You. You have the right 
# to use, modify, and incorporate this Data into other products for 
# purposes authorized by the Autodesk software license agreement, 
# without fee.
#
# The copyright notices in the Software and this entire statement, 
# including the above license grant, this restriction and the 
# following disclaimer, must be included in all copies of the 
# Software, in whole or in part, and all derivative works of 
# the Software, unless such copies or derivative works are solely 
# in the form of machine-executable object code generated by a 
# source language processor.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND. 
# AUTODESK DOES NOT MAKE AND HEREBY DISCLAIMS ANY EXPRESS OR IMPLIED 
# WARRANTIES INCLUDING, BUT NOT LIMITED TO, THE WARRANTIES OF 
# NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR 
# PURPOSE, OR ARISING FROM A COURSE OF DEALING, USAGE, OR 
# TRADE PRACTICE. IN NO EVENT WILL AUTODESK AND/OR ITS LICENSORS 
# BE LIABLE FOR ANY LOST REVENUES, DATA, OR PROFITS, OR SPECIAL, 
# DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES, EVEN IF AUTODESK 
# AND/OR ITS LICENSORS HAS BEEN ADVISED OF THE POSSIBILITY 
# OR PROBABILITY OF SUCH DAMAGES.
#
# ==========================================================================
#+


import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx
import maya.OpenMayaRender as OpenMayaRender
import math
import random
import sys

kSimpleEmitterNodeName = "spSimpleEmitter"
kSimpleEmitterNodeID = OpenMaya.MTypeId(0x8700F)

SEGMENTS = 20
EDGES = 30
TORUS_PI = 3.14159265
TORUS_2PI = 2.0 * TORUS_PI

glRenderer = OpenMayaRender.MHardwareRenderer.theRenderer()
glFT = glRenderer.glFunctionTable()

class simpleEmitter(OpenMayaMPx.MPxEmitterNode):
    def __init__(self):
        OpenMayaMPx.MPxEmitterNode.__init__(self)

    def isFullValue( self, plugIndex, block ):
        value = 1
        mIsFull = OpenMayaMPx.cvar.MPxEmitterNode_mIsFull

        try:
            mhValue = block.inputArrayValue( mIsFull )
            mhValue.jumpToElement( plugIndex )
            hValue = mhValue.inputValue( )
            value = hValue.asBool()
        except:
            sys.stderr.write("Error getting the input array value\n")
            raise
        return value
                    
    
    def getWorldPosition( self, point ):
        try:
            thisNode = simpleEmitter.thisMObject( self )
            fnThisNode = OpenMaya.MFnDependencyNode(thisNode)

            worldMatrixAttr = fnThisNode.attribute( "worldMatrix" )

            matrixPlug = OpenMaya.MPlug( thisNode, worldMatrixAttr )
            matrixPlug = matrixPlug.elementByLogicalIndex( 0 )

            matrixObject = OpenMaya.MObject()
            matrixObject = matrixPlug.asMObject(  )

            worldMatrixData = OpenMaya.MFnMatrixData( matrixObject )
            worldMatrix = worldMatrixData.matrix( )

            point.x = worldMatrix( 3, 0 )
            point.y = worldMatrix( 3, 1 )
            point.z = worldMatrix( 3, 2 )
        except:
            sys.stderr.write( "Error in getWorldPosition\n" )
            raise

    def currentTimeValue( self, block ):
        try:
            mCurrentTime = OpenMayaMPx.cvar.MPxEmitterNode_mCurrentTime
            hValue = block.inputValue( mCurrentTime )
            value = hValue.asTime()
        except:
            sys.stderr.write("Error getting current time value, returning 0")
            value = OpenMaya.MTime(0.0)
        return value

    def startTimeValue( self, plugIndex, block ):
        try:
            mStartTime = OpenMayaMPx.cvar.MPxEmitterNode_mStartTime
            mhValue = block.inputArrayValue( mStartTime )
            mhValue.jumpToElement( plugIndex )
            hValue = mhValue.inputValue( )
            value = hValue.asTime( )        
        except:
            sys.stderr.write("Error getting start time value, setting to 0")
            value = OpenMaya.MTime(0.0)
        return value

    def deltaTimeValue( self, plugIndex, block ):
        try:
            mDeltaTime = OpenMayaMPx.cvar.MPxEmitterNode_mDeltaTime

            mhValue = block.inputArrayValue( mDeltaTime )
            mhValue.jumpToElement( plugIndex )

            hValue = mhValue.inputValue()
            value = hValue.asTime()
        except:
            sys.stderr.write("Error getting delta time value, setting to 0\n")
            value = OpenMaya.MTime(0.0)
        return value

    def rateValue( self, block ):
        try:
            mRate = OpenMayaMPx.cvar.MPxEmitterNode_mRate
            hValue = block.inputValue( mRate )
            value = hValue.asDouble()       
        except:
            sys.stderr.write("Error getting rate value, setting to 0\n")
            value = 0.0
        return value

    def directionValue( self, block ):
        try:
            mDirection = OpenMayaMPx.cvar.MPxEmitterNode_mDirection
            hValue = block.inputValue( mDirection )
            value = hValue.asDouble3()
            dirV = OpenMaya.MVector(value[0], value[1], value[2])
        except:
            sys.stderr.write("Error getting direction value, setting to 0,0,0\n")
            dirV = OpenMaya.MVector(0.0, 0.0, 0.0)
        return dirV

    def speedValue( self, block ):
        try:
            mSpeed = OpenMayaMPx.cvar.MPxEmitterNode_mSpeed
            hValue = block.inputValue( mSpeed )
            value = hValue.asDouble()       
        except:
            sys.stderr.write("Error getting speed value, setting to 0\n")
            value = 0.0
        return value

    def inheritFactorValue( self, plugIndex, block ):
        try:
            mInheritFactor = OpenMayaMPx.cvar.MPxEmitterNode_mInheritFactor
            mhValue = block.inputArrayValue( mInheritFactor )
            mhValue.jumpToElement( plugIndex )
            hValue = mhValue.inputValue( )
            value = hValue.asDouble()       
        except:
            sys.stderr.write("Error getting inherit factor value, setting to 0\n")
            value = 0.0
        return value

    def useRotation( self, direction ):
        try:
            thisNode = simpleEmitter.thisMObject(self)
            fnThisNode = OpenMaya.MFnDependencyNode( thisNode )

            worldMatrixAttr = fnThisNode.attribute( "worldMatrix" )

            matrixPlug = OpenMaya.MPlug( thisNode, worldMatrixAttr )
            matrixPlug = matrixPlug.elementByLogicalIndex( 0 )

            matrixObject = OpenMaya.MObject()
            matrixObject = matrixPlug.asMObject( )

            worldMatrixData = OpenMaya.MFnMatrixData( matrixObject )
            worldMatrix = worldMatrixData.matrix( )

            rotatedVector = OpenMaya.MVector()
            rotatedVector = direction * worldMatrix
        except:
            sys.stderr.write("Error getting rotation value, setting to 0,0,0\n")
            rotatedVector = OpenMaya.MVector(0.0, 0.0, 0.0)
        return rotatedVector


    def compute(self, plug, block):
        mOutput = OpenMayaMPx.cvar.MPxEmitterNode_mOutput

        # Determine if we are requesting the output plug for this emitter node.
        #
        if plug == mOutput:
            # Get the logical index of the element this plug refers to,
            # because the node can be emitting particles into more than
            # one particle shape.
            #
            try:
                multiIndex = plug.logicalIndex( )

                # Get output data arrays (position, velocity, or parentId)
                # that the particle shape is holding from the previous frame.
                #
                hOutArray = block.outputArrayValue ( mOutput )

                # Create a builder to aid in the array construction efficiently.
                #
                bOutArray = hOutArray.builder( )

                # Get the appropriate data array that is being currently evaluated.
                #
                hOut = bOutArray.addElement( multiIndex )

                # Create the data and apply the function set,
                # particle array initialized to length zero, 
                # fnOutput.clear()
                #
                fnOutput = OpenMaya.MFnArrayAttrsData()
                dOutput = fnOutput.create( )

                # Check if the particle object has reached it's maximum,
                # hence is full. If it is full then just return with zero particles.
                #
                beenFull = simpleEmitter.isFullValue( self, multiIndex, block )
                if beenFull == 1:
                    return

                # Get input position and velocity arrays where new particles are from,
                # also known as the owner. An owner is determined if connections exist
                # to the emitter node from a shape such as nurbs, polymesh, curve, 
                # or a lattice shape.
                #
                
                # Get a single position from world transform
                #
                inPosAry = OpenMaya.MVectorArray()
                inPosAry.clear()
            
                worldPos = OpenMaya.MPoint(0.0, 0.0, 0.0)
                simpleEmitter.getWorldPosition( self, worldPos )

                worldV = OpenMaya.MVector(worldPos[0], worldPos[1], worldPos[2])
                inPosAry.append( worldV )

                # Create a single velocity          
                inVelAry = OpenMaya.MVectorArray()
                inVelAry.clear()
                velocity = OpenMaya.MVector(0, 0, 0)
                inVelAry.append( velocity )

                # Get deltaTime, currentTime and startTime.
                # If deltaTime <= 0.0, or currentTime <= startTime,
                # do not emit new pariticles and return.
                #
                cT = simpleEmitter.currentTimeValue( self, block )
                sT = simpleEmitter.startTimeValue( self, multiIndex, block )
                dT = simpleEmitter.deltaTimeValue( self, multiIndex, block )

                dTValue = dT.value()
            
                if cT <= sT or dTValue <= 0.0:
                    # We do not emit particles before the start time, 
                    # and do not emit particles when moving backwards in time.
                    # 

                    # This code is necessary primarily the first time to 
                    # establish the new data arrays allocated, and since we have 
                    # already set the data array to length zero it does 
                    # not generate any new particles.
                    # 
                    hOut.setMObject( dOutput )
                    block.setClean( plug )
                    return

                # Compute and store an emission rate
                #
                emitCountPP = OpenMaya.MIntArray()
                emitCountPP.clear()

                plugIndex = plug.logicalIndex( )

                # Get rate and delta time.
                #
                rate = simpleEmitter.rateValue( self, block )
                dtRate = simpleEmitter.deltaTimeValue( self, plugIndex, block )
                dtRateDbl = dtRate.asUnits( OpenMaya.MTime.kSeconds )
                dblCount = rate * dtRateDbl
                intCount = int(dblCount)
                emitCountPP.append( intCount )

                # Get speed, direction vector, and inheritFactor attributes.
                #           
                speed = simpleEmitter.speedValue( self, block )
                dirV = simpleEmitter.directionValue( self, block )
                inheritFactor = simpleEmitter.inheritFactorValue( self, multiIndex, block )

                # Get the position, velocity, and normalized time arrays to append new particle data.
                #
                fnOutPos = fnOutput.vectorArray( "position" )
                fnOutVel = fnOutput.vectorArray( "velocity" )
                fnOutTime = fnOutput.doubleArray( "timeInStep" )

                # Convert deltaTime into seconds.
                #           
                dt = dT.asUnits( OpenMaya.MTime.kSeconds )

                # Rotate the direction attribute by world transform
                rotatedV = simpleEmitter.useRotation( self, dirV )

                # Start emitting particles.
                #
                simpleEmitter.emit( self, inPosAry, inVelAry, emitCountPP, dt, speed, inheritFactor,\
                        rotatedV, fnOutPos, fnOutVel, fnOutTime)

                # Update the data block with new dOutput and set plug clean.
                #   
                # sys.__stdout__.write( " handle: " + str(hOut) + "\n" )        
                hOut.setMObject( dOutput )
                block.setClean( plug )
            except:
                sys.stderr.write("simpleEmitter compute error\n")
                raise
        else:
            return OpenMaya.kUnknownParameter

    def emit( self, inPosAry, inVelAry, emitCountPP, dt, speed, inheritFactor, dirV, outPosAry, outVelAry, outTimeAry):
        try:
            posLength = inPosAry.length()
            velLength = inVelAry.length()
            countLength = emitCountPP.length()

            if not posLength == velLength or not posLength == countLength:
                return

            totalCount = 0
            for index in range(countLength):
                totalCount += emitCountPP[index]
            if totalCount <= 0:
                return

            dirV.normalize()
            
            for index in range(posLength):
                emitCount = emitCountPP[index]
                if emitCount <= 0:
                    continue
            
                sPos = inPosAry[index]
                sVel = inVelAry[index]
                prePos = sPos - sVel * dt

                for i in range(emitCount):
                    alpha = ( float(i) + random.random() ) / float(emitCount)
                    newPos = prePos * (1.0 - alpha) + sPos * alpha
                    newVel = dirV * speed

                    newPos += newVel * ( dt * (1.0 - alpha) )
                    newVel += sVel * inheritFactor

                    # Add new data into output arrays.
                    #
                    outPosAry.append( newPos )
                    outVelAry.append( newVel )
                    outTimeAry.append( alpha )
        except Exception, e:
            sys.stderr.write( "Error in simpleEmitter.emit\n" )
            raise

    def draw( self, view, path, style, status):
        view.beginGL()

        for j in range(0, SEGMENTS):
            glFT.glPushMatrix()
            glFT.glRotatef(float(360.0 * j/SEGMENTS), 0.0, 1.0, 0.0)
            glFT.glTranslatef( 1.5, 0.0, 0.0)
            
            for i in range(0, EDGES):
                glFT.glBegin(OpenMayaRender.MGL_LINE_STRIP)

                p0 = float(TORUS_2PI * i/EDGES)
                p1 = float(TORUS_2PI * (i+1)/EDGES)
                glFT.glVertex2f(math.cos(p0), math.sin(p0))
                glFT.glVertex2f(math.cos(p1), math.sin(p1))
                
                glFT.glEnd()

            glFT.glPopMatrix()

        view.endGL()


def nodeCreator():
    return OpenMayaMPx.asMPxPtr( simpleEmitter() )

def nodeInitializer():
    return

# initialize the script plug-in
def initializePlugin(mobject):
    mplugin = OpenMayaMPx.MFnPlugin(mobject)

    try:
        mplugin.registerNode( kSimpleEmitterNodeName, kSimpleEmitterNodeID, \
                                nodeCreator, nodeInitializer, OpenMayaMPx.MPxNode.kEmitterNode )
    except:
        sys.stderr.write( "Failed to register node: %s\n" % kSimpleEmitterNodeName )
        raise

# uninitialize the script plug-in
def uninitializePlugin(mobject):
    mplugin = OpenMayaMPx.MFnPlugin(mobject)

    try:
        mplugin.deregisterNode( kSimpleEmitterNodeID )
    except:
        sys.stderr.write( "Failed to unregister node: %s\n" % kSimpleEmitterNodeName )
        raise