A* for Maya

a_star_for_maya.py

# -*- coding: utf-8 -*- 
#A*アルゴリズム
from maya import cmds
from collections import OrderedDict
import maya.api.OpenMaya as om2
import math
import copy
import time

#ヒューリスティック関数,ゴールからのユークリッド距離を算出
def heuristic_func(end, vtx_pos_dict):
    heuristic_dict = {}
    end_pos = om2.MPoint(vtx_pos_dict[end])
    for vtx in vtx_pos_dict.keys():
        #dist = culcVtxDist(vtx, end)
        pos = om2.MPoint(vtx_pos_dict[vtx])
        dist = (pos - end_pos).length()
        heuristic_dict[vtx] = dist
    return heuristic_dict

def a_star():
    #開始時間
    print 'A-Star algorithm------------------------------------'
    start = time.time()
    rap_time = time.time()
    
    #メッシュと探索頂点の開始終了位置を取得
    try:
        target_mesh = cmds.ls(hl=True)[0]
        start_vtx = cmds.ls(sl=True, fl=True)[0]
        goal_vtx = cmds.ls(sl=True, fl=True)[1]
    except:
        print '頂点を2つ選択してください'
        return
    
    #すべての頂点の位置を取得
    selection = om2.MGlobal.getActiveSelectionList()
    mDag, __ = selection.getComponent(0)
    if mDag.hasFn(om2.MFn.kMesh):
        mfnMesh = om2.MFnMesh(mDag)
    vertex_array = mfnMesh.getPoints(om2.MSpace.kWorld)
    
    #メッシュの全頂点を取得しておく
    all_vtx = cmds.polyListComponentConversion(target_mesh, tv=True)
    all_vtx = cmds.filterExpand(all_vtx, sm=31)
    
    #頂点位置をすべて格納しておく
    vtx_pos_dict = {vtx:pos for vtx, pos in zip(all_vtx, vertex_array)}

    print u'rap time / メッシュ頂点取得:', time.time() - rap_time
    rap_time = time.time()

    #頂点距離に重みづけ、初期値無限大を設定しておく。
    temp_dict  = {}
    for vtx in all_vtx:
        temp_dict [vtx] = [('', float('inf'))]
        
    print u'rap time / つながり辞書   :', time.time() - rap_time
    rap_time = time.time()
        
    #スタートとゴールの経路を保持するためオーダードディクトを利用
    route_dict = OrderedDict()

    #スタートとゴールを最初と最後のキーと要素に設定する。あとは適当。
    route_dict[start_vtx] = temp_dict [start_vtx]
    for k in temp_dict .keys():
        if k != start_vtx and k != goal_vtx:
            route_dict[k] = temp_dict [k]
    route_dict[goal_vtx] = temp_dict [goal_vtx]

    print 'rap time / ルート格納辞書 :', time.time() - rap_time
    rap_time = time.time()

    #ゴールの周辺を取得
    around = cmds.polyListComponentConversion(goal_vtx, tf=True)
    around = cmds.polyListComponentConversion(around, tv=True)
    around = cmds.filterExpand(around, sm=31)
    around.remove(goal_vtx)
    round_of_goul = around

    #探索用変数を色々初期化
    nodeNum = len(route_dict)#ノードの総数
    distanceDict = {vtx:float('inf') for vtx in all_vtx}#ノードごとの距離のリスト、初期値無限大を与える
    previousNodesDict = {vtx:'' for vtx in all_vtx}#最短経路でそのノードのひとつ前に到達するノードの対応辞書
    distanceDict[start_vtx] = 0#初期のノードの距離は0とする
    unsearchedNodes = copy.copy(distanceDict)#未探索ノード
    real_dist_dict = copy.copy(distanceDict)#実際の距離を格納する辞書

    print 'rap time / データ初期化   :', time.time() - rap_time
    rap_time = time.time()

    #ヒューリスティック関数実行、全頂点のゴールからのユークリッド距離を出しておく
    heuristic_dict = heuristic_func(goal_vtx, vtx_pos_dict)

    print 'rap time / ヒューリスティック:', time.time() - rap_time
    rap_time = time.time()
    i = 0
    while(len(unsearchedNodes) != 0): #未探索ノードがなくなるまで繰り返す
        #未探索ノードのうちで最小のindex番号を取得
        min_dist = min(unsearchedNodes.values())
        min_id = unsearchedNodes.values().index(min_dist)
        target_min_key = unsearchedNodes.keys()[min_id]
               
        #未探索ノードから除去しておく
        unsearchedNodes.pop(target_min_key)
        
        #ターゲットノードの周辺ノード取得
        roundVertex = route_dict[target_min_key]
        
        if roundVertex[0][1] == float('inf'):
            around = cmds.polyListComponentConversion(target_min_key, tf=True)
            around = cmds.polyListComponentConversion(around, tv=True)
            around = cmds.filterExpand(around, sm=31)
            around.remove(target_min_key)
            roundList = []
            a = om2.MPoint(vtx_pos_dict[target_min_key])
            for rVtx in around:
                b = om2.MPoint(vtx_pos_dict[rVtx])
                dist = (a-b).length()
                roundList.append((rVtx, dist))
            route_dict[target_min_key] = roundList
            roundVertex = route_dict[target_min_key]
        #周辺ノードへのスタートからの到達距離を比較、以前設定した距離より小さい場合は更新する。
        for rVtx in roundVertex:
            i += 1
            #rVtx[1]:コスト,heuristic_dict[rVtx[0]]:ユークリッド距離
            real_dist = real_dist_dict[target_min_key] + rVtx[1]#実際の距離
            sum_dist = real_dist + heuristic_dict[rVtx[0]]#ゴールまでのユークリッド距離を加味した推定距離
            if distanceDict[rVtx[0]] > sum_dist:#推定距離が今までの推定距離より小さかったら更新
                distanceDict[rVtx[0]] = sum_dist
                unsearchedNodes[rVtx[0]] = sum_dist
                real_dist_dict[rVtx[0]] = real_dist#rVtxまでのスタートからの距離を格納
                #そのノードに最短で到達するための一つ前のノード辞書を更新しておく
                previousNodesDict[rVtx[0]] = target_min_key
        #ゴール前まで辿り着いたらブレイク
        if target_min_key in round_of_goul:
            break
            
    print u'rap time / ルート探索     :', time.time() - rap_time
    rap_time = time.time()

    #最短ひとつ前辞書をゴールから逆引きしてスタートまでつなげる。
    minimum_route = []
    previous_node = goal_vtx
    while previous_node != start_vtx:
        minimum_route.append(previous_node)
        previous_node = previousNodesDict[previous_node]
    #ゴールへの最短到達距離を取得
    minimum_route.append(previous_node)

    #探索したつながりを選択
    cmds.select(minimum_route, r=True)
    elapsed_time = time.time() - start

    print u'rap time / 最終データ取得 :', time.time() - rap_time
    rap_time = time.time()

    #プリントして探索結果を確認
    print u'総探索回数  :', i
    print u'取得したルート:', minimum_route[::-1]
    print u'ルートの距離 :', distanceDict[goal_vtx]
    print u'総実行時間  :', elapsed_time

a_star()