particleAttrNode.cpp

//-
// ==========================================================================
// Copyright 1995,2006,2008 Autodesk, Inc. All rights reserved.
//
// Use of this software is subject to the terms of the Autodesk
// license agreement provided at the time of installation or download,
// or which otherwise accompanies this software in either electronic
// or hard copy form.
// ==========================================================================
//+

//
//      Class:  particleAttrNode
//
//      Author: Lonnie Li
//
//      Description:
//
//              particleAttrNode is an example node that extends the MPxParticleAttributeMapperNode
//      type. These nodes allow us to define custom nodes to map per-particle attribute data
//      to a particle shape.
//
//              In this particular example, we have defined two operations:
//
//              - compute2DTexture()
//              - computeMesh()
//
//              compute2DTexture() replicates the internal behaviour of Maya's internal 'arrayMapper'
//      node at the API level. Given an input texture node and a U/V coord per particle input,
//      this node will evaluate the texture at the given coordinates and map the result back
//      to the outColorPP or outValuePP attributes. See the method description for compute2DTexture()
//      for details on how to setup a proper attribute mapping graph.
//
//              computeMesh() is a user-defined behaviour. It is called when the 'computeNode' attribute
//      is connected to a polygonal mesh. From there, given a particle count, it will map the
//      object space vertex positions of the mesh to a user-defined 'outPositions' vector attribute.
//      This 'outPositions' attribute can then be connected to a vector typed, per-particle attribute
//      on the shape to drive. In our particular example we drive the particle shape's 'rampPosition'
//      attribute. For further details, see the method description for computeMesh() to setup
//      this example.
//

#include <particleAttrNode.h>

#include <maya/MStatus.h>
#include <maya/MObject.h>
#include <maya/MDataBlock.h>
#include <maya/MDataHandle.h>
#include <maya/MPlug.h>
#include <maya/MVector.h>
#include <maya/MDoubleArray.h>
#include <maya/MVectorArray.h>
#include <maya/MPlugArray.h>
#include <maya/MPointArray.h>

#include <maya/MFnDoubleArrayData.h>
#include <maya/MFnVectorArrayData.h>
#include <maya/MFnMesh.h>
#include <maya/MFnTypedAttribute.h>
#include <maya/MFnNumericAttribute.h>

#include <maya/MDynamicsUtil.h>

MTypeId particleAttrNode::id = 0x81036;
MObject particleAttrNode::outPositionPP;
MObject particleAttrNode::particleCount;

particleAttrNode::particleAttrNode()
{}

particleAttrNode::~particleAttrNode()
{}

void* particleAttrNode::creator()
{
        return new particleAttrNode();
}

MStatus particleAttrNode::initialize()
{
        MStatus status = MS::kSuccess;

        // In addition to the standard arrayMapper attributes,
        // create an additional vector attribute.
        //
        MFnTypedAttribute typedAttrFn;
        MVectorArray defaultVectArray;
        MFnVectorArrayData vectArrayDataFn;
        vectArrayDataFn.create( defaultVectArray );
        typedAttrFn.create( "outPosition",
                                                "opos",
                                                MFnData::kVectorArray,
                                                vectArrayDataFn.object(),
                                                &status );
        // Transient output attribute. Not writable nor storable.
        //
        typedAttrFn.setWritable(false);
        typedAttrFn.setStorable(false);
        outPositionPP = typedAttrFn.object();
        addAttribute( outPositionPP );

        MFnNumericAttribute numAttrFn;
        numAttrFn.create( "particleCount", "pc", MFnNumericData::kInt );
        particleCount = numAttrFn.object();
        addAttribute( particleCount );

        attributeAffects( computeNode, outPositionPP );
        attributeAffects( particleCount, outPositionPP );

        return status;
}

MStatus particleAttrNode::compute( const MPlug& plug, MDataBlock& block )
{
        MStatus status = MS::kUnknownParameter;
        if( (plug.attribute() == outColorPP) || (plug.attribute() == outValuePP) )
        {
                status = compute2DTexture( plug, block );
        }
        else if( plug.attribute() == outPositionPP )
        {
                status = computeMesh( plug, block );
        }
        return status;
}

MStatus particleAttrNode::compute2DTexture( const MPlug& plug, MDataBlock& block )
//
//      Description:
//
//              This computes an outputValue/Color based on input U and V coordinates.
//      The code contained herein replicates the compute code of the internal Maya
//      'arrayMapper' node.
//
//      How to use:
//
//              1) Create a ramp
//              2) Create a particleAttr node
//              3) Create an emitter with a particleShape
//              4) Create a per particle color attribute (rgbPP) on the particleShape via the AE.
//              5) connectAttr ramp1.outColor particleAttr1.computeNode
//              6) connectAttr particleShape1.age particleAttr1.vCoordPP
//              7) connectAttr particleAttr1.outColorPP particleShape1.rgbPP
//
{
        MStatus status = MS::kSuccess;

        //
        // We request the computeNodeColor attributes here to make sure that
        // they get marked clean as the arrayMapper's output attribute(s) get
        // computed.
        //
        /* MDataHandle hComputeNodeColor = */ block.inputValue( computeNodeColor );
        /* MDataHandle hComputeNodeColorR = */ block.inputValue( computeNodeColorR );
        /* MDataHandle hComputeNodeColorG = */ block.inputValue( computeNodeColorG );
        /* MDataHandle hComputeNodeColorB = */ block.inputValue( computeNodeColorB );

        bool doUcoord = false;
        bool doVcoord = false;
        bool doOutColor = ( plug.attribute() == outColorPP );
        bool doOutValue = ( plug.attribute() == outValuePP );

        // Get pointers to the source attributes: aUCoordPP, aVCoordPP.
        //
        MFnDoubleArrayData dataDoubleArrayFn;
        MObject uCoordD = block.inputValue( uCoordPP ).data();
        status = dataDoubleArrayFn.setObject( uCoordD );
        MDoubleArray uAry;
        if( status == MS::kSuccess )
        {
                uAry = dataDoubleArrayFn.array();
                if (uAry.length())
                {
                        doUcoord = true;
                }
        }

        status = MS::kSuccess;
        MObject vCoordD = block.inputValue( vCoordPP ).data();
        status = dataDoubleArrayFn.setObject( vCoordD );
        MDoubleArray vAry;
        if( status == MS::kSuccess )
        {
                vAry = dataDoubleArrayFn.array();
                if (vAry.length())
                {
                        doVcoord = true;
                }
        }

        // Get pointers to destination attributes: outColorPP, outValuePP.
        //
        status = MS::kSuccess;
        MFnVectorArrayData dataVectorArrayFn;
        MVectorArray outColorAry;
        if( doOutColor )
        {
                MObject colorD = block.outputValue( outColorPP ).data();
                status = dataVectorArrayFn.setObject( colorD );

                if( status == MS::kSuccess )
                {
                        outColorAry = dataVectorArrayFn.array();
                }
        }

        status = MS::kSuccess;
        MDoubleArray outValueAry;
        if( doOutValue )
        {
                MObject valueD = block.outputValue( outValuePP ).data();
                status = dataDoubleArrayFn.setObject( valueD );

                if( status == MS::kSuccess )
                {
                        outValueAry = dataDoubleArrayFn.array();
                }
        }

        // Setup loop counter.
        //
        unsigned int uCount = ( doUcoord ? uAry.length() : 0 );
        unsigned int vCount = ( doVcoord ? vAry.length() : 0 );
        unsigned int count  = ( (uCount) > (vCount) ? (uCount) : (vCount) );

        // Resize the output arrays to match the max size
        // of the input arrays.
        //
        if( doOutColor ) outColorAry.setLength( count );
        if( doOutValue ) outValueAry.setLength( count );

        // Check for a texture node
        //
        bool hasTextureNode =
                MDynamicsUtil::hasValidDynamics2dTexture( thisMObject(),
                                                                                                computeNode );

        // Call the texture node to compute color values for
        // the input u,v coordinate arrays.
        //
        if( hasTextureNode )
        {
                MDynamicsUtil::evalDynamics2dTexture( thisMObject(),
                                                                                        computeNode,
                                                                                        uAry,
                                                                                        vAry,
                                                                                        &outColorAry,
                                                                                        &outValueAry );

                // Scale the output value and/or color to the min/max range.
                //
                const double& minValue = block.inputValue( outMinValue ).asDouble();
                const double& maxValue = block.inputValue( outMaxValue ).asDouble();

                if( (minValue != 0.0) || (maxValue != 1.0) )
                {
                        double r = maxValue - minValue;
                        if( doOutValue )
                        {
                                for( unsigned int i = 0; i < count; i++ )
                                {
                                        outValueAry[i] = (r * outValueAry[i]) + minValue;
                                }
                        }

                        if( doOutColor )
                        {
                                MVector minVector( minValue, minValue, minValue );
                                for( unsigned int i = 0; i < count; i++ )
                                {
                                        outColorAry[i] = (r * outColorAry[i]) + minVector;
                                }
                        }
                }
        }
        else
        {
                // Not connected to texture node, so simply set output
                // arrays to default values.
                //
                for( unsigned int i = 0; i < count; i++ )
                {
                        MVector clr( 0.0, 0.0, 0.0 );
                        if( doOutColor )
                        {
                                outColorAry[i] = clr;
                        }
                        if( doOutValue )
                        {
                                outValueAry[i] = clr.x;
                        }
                }
        }
        dataVectorArrayFn.set( outColorAry );
        dataDoubleArrayFn.set( outValueAry );
        return MS::kSuccess;
}

MStatus particleAttrNode::computeMesh( const MPlug& plug, MDataBlock& block )
//
//      Description:
//
//              To illustrate a custom use of the arrayMapper node, we'll define this
//      alternative compute that will map particles to mesh vertex positions. If
//      array lengths do not match, we'll wrap around the arrays.
//
//      How to use:
//
//              1) Create a poly surface.
//              2) connectAttr polyShape1.outMesh particleAttr1.computeNode
//              3) connectAttr particleShape1.count particleAttr1.particleCount
//              4) connectAttr particleAttr1.outPosition particleShape1.rampPosition (*)
//
//              * You may choose to drive any vector based per particle attribute.
//                rampPosition was selected here as an example.
//
{
        MStatus status = MS::kSuccess;

        // Verify that computeNode has a mesh connected to computeNode
        //
        MPlug compPlug( thisMObject(), computeNode );
        MPlugArray conns;
        compPlug.connectedTo( conns, true, false, &status );
        if( conns.length() <= 0 )
        {
                return MS::kFailure;
        }
        MPlug conn = conns[0];
        MObject compNode = conn.node();
        MFnMesh meshFn( compNode, &status );
        if( status != MS::kSuccess )
        {
                return MS::kFailure;
        }

        // Retrieve the mesh vertices
        //
        MPointArray points;
        meshFn.getPoints( points );
        unsigned int nPoints = points.length();

        // Retrieve the current particle count.
        //
        // NOTE: Due to the order in which the particle system requests
        // various pieces of data, some attributes are requested prior
        // to the actual emission of particles (eg. rampPosition), whereas
        // other attributes are requested after particles have been emitted.
        //
        // If the driven PP attribute on the particleShape is requested prior
        // to the emission of particles, this compute() method will be called
        // before any particles have been emitted. As a result, the effect
        // will be lagged by one frame.
        //
        unsigned int nParticles = 0;
        int nSignedPart = block.inputValue( particleCount ).asInt();
        if( nSignedPart > 0 )
        {
                nParticles = nSignedPart;
        }

        // Get pointer to destination attribute: outPositionPP
        //
        MFnVectorArrayData dataVectorArrayFn;
        MVectorArray outPosArray;
        MObject posD = block.outputValue( outPositionPP ).data();
        //const char* typeStr = posD.apiTypeStr();
        status = dataVectorArrayFn.setObject( posD );
        if( status == MS::kSuccess )
        {
                outPosArray = dataVectorArrayFn.array();
        }

        outPosArray.setLength( nParticles );
        for( unsigned int i = 0; i < nParticles; i++ )
        {
                unsigned int index = i;
                if( nParticles > nPoints )
                {
                        index = i % nPoints;
                }
                MPoint point = points[index];
                MVector pos( point.x, point.y, point.z );
                outPosArray[i] = pos;
        }
        dataVectorArrayFn.set( outPosArray );
        return MS::kSuccess;
}