Joint Aim

import imath
import AtomsCore
import Atoms
import AtomsUtils


class AgentAimModule(Atoms.BehaviourModule):

    def __init__(self):
        Atoms.BehaviourModule.__init__(self)

        metadata = self.attributes()
        metadata.addEntry("target", AtomsCore.Vector3Metadata())

    def initFrame(self, agents, agroup):
        target = self.attributes().getEntry("target").get()

        for agent in agents:
            agent_position = agent.metadata()["position"].value()
            new_direction = target - agent_position
            new_direction.normalize()
            agent.metadata()["direction"].set(new_direction)


'''
MAYA SETUP 
'''
from maya import cmds

def register():
    Atoms.BehaviourModules.instance().registerBehaviourModule("agentAim",
                                                              AgentAimModule,
                                                              True)

def setup():
    cmds.tcAtoms(init=True)
    agn = cmds.createNode('tcAgentGroupNode')
    cmds.setAttr(agn + ".displayType", 1)

    register()

    cmds.setAttr(agn + ".modules[0].moduleType", "gridLayout", type="string")
    cmds.setAttr(agn + ".modules[0].moduleName", "gridLayout", type="string")

    cmds.setAttr(agn + ".modules[1].moduleType", "stateMachine", type="string")
    cmds.setAttr(agn + ".modules[1].moduleName", "stateMachine", type="string")
    cmds.setAttr(agn + ".atoms_stateMachine_state", 1)

    cmds.playbackOptions(min=1, max=50)

    s = cmds.spaceLocator()[0]
    cmds.setAttr(s + ".t", 200, 50, -120)
    cmds.setKeyframe(s)
    cmds.currentTime(50)
    cmds.setAttr(s + ".t", 420, 50, 100)
    cmds.setKeyframe(s)
    cmds.currentTime(0)

    cmds.setAttr(agn + ".modules[2].moduleType", "agentAim", type="string")
    cmds.setAttr(agn + ".modules[2].moduleName", "agentAim", type="string")
    cmds.connectAttr(s + ".t", agn + ".atoms_agentAim_target")