有限元一阶四面体单元python编程（二）

def nodeData2ElemData(nodeData, ELEM, nodes_per_elem =4):
    #后处理云图是按照单元依次绘制，所有须要把按节点排列的数据转化为按单元排列
    node_qty = nodeData.shape[0] # 按节点排序的 X向应力 等等这样的数组,size  node_qty x 1
    elem_qty = ELEM.shape[0] # ELEM 中 2到5 列包含单元中4个节点的ID
    elemData = zeros((elem_qty,nodes_per_elem),dtype =nodeData.dtype)
    for i in range(elem_qty):
        for j in range(nodes_per_elem):
            nodeID = ELEM[i,2+j] #须为整数 # ELEM 中 2到5 列包含单元中4个节点的ID
            elemData[i,j] = nodeData[nodeID]
    return elemData

X= nodeData2ElemData(X,ELEM)
Y= nodeData2ElemData(Y,ELEM)
Z= nodeData2ElemData(Z,ELEM)
Delta_X = nodeData2ElemData(Delta_X,ELEM)
Delta_Y = nodeData2ElemData(Delta_Y,ELEM)
Delta_Z = nodeData2ElemData(Delta_Z,ELEM)
scale = 2.0e4 #变形放大系数
X_new = X +  scale* Delta_X
Y_new = Y +  scale* Delta_Y
Z_new = Z +  scale* Delta_Z
#计算按单元排列的各个后处理变量
DISP_total = nodeData2ElemData(DISP_total,ELEM) #总位移

XStress = Stress_elems[:,0]# X向正应力
YStress = Stress_elems[:,1]# Y向正应力
ZStress = Stress_elems[:,2]# Z向正应力
XYStress = Stress_elems[:,3]# xy向切应力
YZStress = Stress_elems[:,4]# yz向切应力
ZXStress = Stress_elems[:,5]# zx向切应力
#Von mises 应力
VonmiStress = sqrt(2)/2.0* sqrt((XStress-YStress)**2 + (YStress-ZStress)**2 + (ZStress-XStress)**2 +6*(XYStress**2 + YZStress**2 + ZXStress**2))

from OpenGL.GL import *
from OpenGL.GLUT import *
from OpenGL.GLU import *
import sys

class TetraShow():
    def __init__(self,X,Y,Z,Data):
        self.rtri = 0 # inital Rotation angle for the triangle.
        self.X = X
        self.Y= Y
        self.Z = Z
        self.Data = Data
        USL = Data.max()
        LSL = Data.min()
        
        elem_qty = Data.shape[0]
        self.Color= zeros((elem_qty,4,3))
        for i in range(elem_qty):
            for j in range(4):
                self.Color[i,j,:]= self.num2RGB(Data[i,j],LSL, USL)
        self.main()
       
    # A general OpenGL initialization function.  Sets all of the initial parameters. 
    def InitGL(self,Width, Height):    # We call this right after our OpenGL window is created.
        glClearColor(0.0, 0.0, 0.0, 0.0) # 清除背景颜色并设为黑色
        glClearDepth(1.0)     # Enables Clearing Of The Depth Buffer
        glDepthFunc(GL_LESS)    # The Type Of Depth Test To Do
        glEnable(GL_MULTISAMPLE)#多点采样（可抗锯齿）
        glEnable(GL_DEPTH_TEST) # 激活深度测试
        glShadeModel(GL_SMOOTH) # 激活颜色渐变渲染
            
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()     # 重置投影矩阵
                                                                            
        gluPerspective(45.0, float(Width)/float(Height), 0.1, 100.0)#计算窗口长宽比
        glMatrixMode(GL_MODELVIEW)
    # The function called when our window is resized (which shouldn't happen if you enable fullscreen, below)
    def ReSizeGLScene(self,Width, Height):
        if Height == 0:      # Prevent A Divide By Zero If The Window Is Too Small 
                Height = 1
        glViewport(0, 0, Width, Height)  # Reset The Current Viewport And Perspective Transformation
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(45.0, float(Width)/float(Height), 0.1, 100.0)
        glMatrixMode(GL_MODELVIEW)
    # The main drawing function. 
    def DrawGLScene(self):
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT) # 清除屏幕，和深度缓冲
        glLoadIdentity()     # Reset The View
        glTranslatef(-0,0.0,-10)#视图缩放可以调这里  # Move Left And Into The Screen
        glRotatef(self.rtri,1.0,1.0,0.0)   # Rotate The Pyramid On It's self.X Axis
        glBegin(GL_TRIANGLES)#画三角形
        for i in range(elem_qty):
            color = self.Color[i,0]
            glColor3f(self.Color[i,0,0],self.Color[i,0,1],self.Color[i,0,2])   # Red
            glVertex3f( self.X[i,0], self.Y[i,0], self.Z[i,0])  # Top Of Triangle (Front)
            glColor3f(self.Color[i,1,0],self.Color[i,1,1],self.Color[i,1,2])   # Green
            glVertex3f( self.X[i,1], self.Y[i,1], self.Z[i,1])  # Left Of Triangle (Front)
            glColor3f(self.Color[i,2,0],self.Color[i,2,1],self.Color[i,2,2])   # Blue
            glVertex3f( self.X[i,2], self.Y[i,2], self.Z[i,2])
            glColor3f(self.Color[i,1,0],self.Color[i,1,1],self.Color[i,1,2])
            glVertex3f( self.X[i,1], self.Y[i,1], self.Z[i,1])  # Top Of Triangle (Right)
            glColor3f(self.Color[i,2,0],self.Color[i,2,1],self.Color[i,2,2])
            glVertex3f( self.X[i,2], self.Y[i,2], self.Z[i,2])  # Left Of Triangle (Right)
            glColor3f(self.Color[i,3,0],self.Color[i,3,1],self.Color[i,3,2])
            glVertex3f( self.X[i,3], self.Y[i,3], self.Z[i,3]) # Right
            glColor3f(self.Color[i,2,0],self.Color[i,2,1],self.Color[i,2,2])
            glVertex3f( self.X[i,2], self.Y[i,2], self.Z[i,2])  # Top Of Triangle (Back)
            glColor3f(self.Color[i,3,0],self.Color[i,3,1],self.Color[i,3,2])
            glVertex3f( self.X[i,3], self.Y[i,3], self.Z[i,3]) # Left Of Triangle (Back)
            glColor3f(self.Color[i,0,0],self.Color[i,0,1],self.Color[i,0,2])
            glVertex3f( self.X[i,0], self.Y[i,0], self.Z[i,0])  # Right Of 
                            
            glColor3f(self.Color[i,3,0],self.Color[i,3,1],self.Color[i,3,2])
            glVertex3f( self.X[i,3], self.Y[i,3], self.Z[i,3]) # Top Of Triangle (Left)
            glColor3f(self.Color[i,0,0],self.Color[i,0,1],self.Color[i,0,2])
            glVertex3f( self.X[i,0], self.Y[i,0], self.Z[i,0]) # Left Of Triangle (Left)
            glColor3f(self.Color[i,1,0],self.Color[i,1,1],self.Color[i,1,2])
            glVertex3f( self.X[i,1], self.Y[i,1], self.Z[i,1]) # Right Of Triangle (Left)
        glEnd()
       
        glLineWidth(2.0)#设置线宽
        glBegin(GL_LINE_LOOP)#画线
        glColor3f(1.0,1.0,1.0)#白色（线） 
        for i in range(elem_qty):
            glVertex3f( self.X[i,0], self.Y[i,0], self.Z[i,0])  # Top Of Triangle (Front)
            glVertex3f( self.X[i,1], self.Y[i,1], self.Z[i,1])  # Left Of Triangle (Front)
            glVertex3f( self.X[i,2], self.Y[i,2], self.Z[i,2])
            glVertex3f( self.X[i,0], self.Y[i,0], self.Z[i,0])
            glVertex3f( self.X[i,1], self.Y[i,1], self.Z[i,1])  # Top Of Triangle (Right)
            glVertex3f( self.X[i,2], self.Y[i,2], self.Z[i,2])  # Left Of Triangle (Right)
            glVertex3f( self.X[i,3], self.Y[i,3], self.Z[i,3]) # Right
            glVertex3f( self.X[i,1], self.Y[i,1], self.Z[i,1])
            glVertex3f( self.X[i,2], self.Y[i,2], self.Z[i,2])  # Top Of Triangle (Back)
            glVertex3f( self.X[i,3], self.Y[i,3], self.Z[i,3]) # Left Of Triangle (Back)
            glVertex3f( self.X[i,0], self.Y[i,0], self.Z[i,0])  # Right Of
            glVertex3f( self.X[i,2], self.Y[i,2], self.Z[i,2])
                            
            glVertex3f( self.X[i,3], self.Y[i,3], self.Z[i,3]) # Top Of Triangle (Left)
            glVertex3f( self.X[i,0], self.Y[i,0], self.Z[i,0]) # Left Of Triangle (Left)
            glVertex3f( self.X[i,1], self.Y[i,1], self.Z[i,1]) # Right Of Triangle (Left)
            glVertex3f( self.X[i,3], self.Y[i,3], self.Z[i,3])
        glEnd()
       
        #控制转速
        self.rtri  +=  +2.4 # Increase The Rotation Variable For The Triangle
        #  since this is double buffered, swap the buffers to display what just got drawn. 
        glutSwapBuffers()
        
    def num2RGB(self,x,LSL=0, USL=1.0):
        r=(x-LSL)/(USL-LSL)
        if r>=0.75:
            return (1, 1*(1-r)*4, 0)
        elif r>=0.5:
            return (1*(r-0.5)*4, 1, 0)
        elif r>=0.25:
            return (0, 1, 1*(0.5-r)*4)
        elif r>=0:
            return (0, 1*r*4, 1)
        
    def main(self):
            global window
            glutInit(sys.argv)
            glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
        
            glutInitWindowSize(1000, 680)
            
            # the window starts at the upper left corner of the screen 
            glutInitWindowPosition(0, 0)
            
            window = glutCreateWindow("Tetra3d")#OPENGL 窗口标题
            glutDisplayFunc(self.DrawGLScene)
            
            #glutFullScreen()#全屏
            # When we are doing nothing, redraw the scene.
            glutIdleFunc(self.DrawGLScene)
            
            # Register the function called when our window is resized.
            glutReshapeFunc(self.ReSizeGLScene)
            
            # Register the function called when the keyboard is pressed.  
            #glutKeyboardFunc(keyPressed)
            # Initialize our window. 
            self.InitGL(640, 480)
            # Start Event Processing Engine 
            glutMainLoop()

data = VonmiStress.reshape((-1,1))

#print(data)
data = tile(data,(1,4)) # 单元内为常应力，复制到4个顶点
tp = TetraShow(X_new,Y_new,Z_new,data)
#tp = TetraShow(X_new,Y_new,Z_new,DISP_total)
tp.main()