Saturday, October 10, 2015

Setting AtMatrix values in Python

When building an Arnold scene in memory, especially when converting from another scene format, one of the most common operations is setting the transform on the new Arnold objects. This is done by setting the matrix attribute which is either an individual AtMatrix, or an array of AtMatrix in the case of several motion steps. 

AtMatrix is defined as a simple 2d array (as of this writing):

typedef float AtMatrix [4][4]
4-by-4 matrix 

The AtMatrix C++ API is pretty straightforward for setting values on an existing matrix:

 AiNodeGetMatrix(node, "matrix", m);  
 // ...
 m[0][0] = 1.0f;  
 m[0][1] = 0.0f;  

The Python API is slightly trickier:

 >>> m[0][0] = 0.0  
 Traceback (most recent call last):  
 File "<stdin>", line 1, in <module>  
 TypeError: 'AtMatrix' object does not support indexing  
   
 >>> m[0] = 0.0  
 Traceback (most recent call last):  
 File "<stdin>", line 1, in <module>  
 TypeError: 'AtMatrix' object does not support item assignment  

So what's going on here?

 >>> help(AtMatrix)  
 Help on class AtMatrix in module arnold.ai_matrix:  
   
 class AtMatrix(_ctypes.Structure)  
 | Method resolution order:  
 | AtMatrix  
 | _ctypes.Structure  
 | _ctypes._CData  
 | builtins.object  
 |  
 | Data descriptors defined here:  
 |  
 | __dict__  
 | dictionary for instance variables (if defined)  
 |  
 | __weakref__  
 | list of weak references to the object (if defined)  
 |  
 | a00  
 | Structure/Union member  
 |  
 | a01  
 | Structure/Union member  
 |  
 | a02  
 | Structure/Union member  
 |  
 | a03  
 | Structure/Union member  
 |  
 | a10  
 | Structure/Union member  
 ...  
 |  
 | a32  
 | Structure/Union member  
 |  
 | a33  
 | Structure/Union member  
 |  

Ah! Each matrix value is represented by its own member variable of the AtMatrix class. No problem! Python can make quick work of this. In my application, my transforms are exported as 16-length arrays, so I wrote two small functions to do the conversion and update the matrix of an AiNode. The xform_to_matrix  function can either update an existing AtMatrix using __setattr__, or create a new matrix using argument expansion.

def set_node_xform(node, transform):  
    """Set an AtNode matrix from one or more transform arrays.  
    
    Set a node transform from a 16-entry python array,  
    or an array of several transform arrays if you want  
    motion blur (max 15 samples)  
     
    Args:  
        node (AtNode): Node to set 'matrix' attribute  
        transform (list(float), or list(list(float))): List, or list of lists  
            containing 16 entries representating a transform matrix.  
    """  
    list_size = len(transform)  
    if list_size < 16:  
        xform_array = AiArrayAllocate(1, list_size, AI_TYPE_MATRIX)  
        for i in range(list_size):  
            AiArraySetMtx(xform_array, i, xform_to_matrix(transform[i]))  

        AiNodeSetArray(node, "matrix", xform_array)     
        AiMsgInfo(b"Setting {0} time samples for {1}".format(  
            list_size,  
            AiNodeGetName(node)  
        ))  
    else:  
        AiNodeSetMatrix(node, "matrix", xform_to_matrix(transform))  
   
       
def xform_to_matrix(transform16, matrix=None):
    """Update or make a new AtMatrix from an array
    
    Create an AtMatrix and set its values from 
    a transform array of 16 entries.
    
    Args:
        transform16 (list(float)): 16-entry python array
        matrix (AtMatrix): existing matrix, or None if creating a new one
        
    Returns:
        AtMatrix: New, initialized matrix
    """
    if matrix is None:
        # Can early-out here with argument expansion if creating
        # a new matrix
        return AtMatrix(*transform16)
    
    for y in range(4):
        for x in range(4):
            matrix.__setattr__("a%d%d"%(y, x), transform16[y*4+x])    
    return matrix 

To demonstrate, I put together a super-complex simulation in Maya and exported the transforms to a text file. A also wrote a short Python script to render all of the frames from the simulation, updating the scene with the correct transforms for each frame:

if __name__ == "__main__":  
    # Read the transfroms dumped from Maya  
    mats_file = open("mats.txt", "r")  
    as_txt = mats_file.read()  
    mats_file.close()  
   
    lines = as_txt.split("\n")  
    mats = []  
    for line in lines:  
        mats.append([float(i) for i in line.split()])  
   
    # Warm up Arnold  
    AiBegin()  
   
    # Show EVERYTING in the log  
    AiMsgSetConsoleFlags(AI_LOG_ALL)  
     
    # Read in the scene file  
    AiASSLoad("cube_bounce.ass")  
   
    cube = AiNodeLookUpByName("cube")  
    driver = AiNodeLookUpByName("exr")  
   
    for frame in range(START, END):  
        # Set the file output name per-frame  
        fileName = "renders/bounce.%03d.exr" % frame  
        AiNodeSetStr(driver, "filename", fileName)  
       
        # Apply the cube transform for this frame  
        xforms = [mats[frame], mats[frame+1]]  
        set_node_xform(cube, xforms)  
       
        # Render this frame  
        result = AiRender()  
        if result != AI_SUCCESS:  
            AiMsgError(b"Quitting because I failed somehow")  
            break  
   
    # All done, tell Arnold to quit  
    AiEnd()  

I think you'll be impressed with the result:



Or generate your very own cube! Grab the code sample from github