FEMの勉強２

# -*- coding: utf-8 -*-

import math

class Material:
    def __init__(self,no):
        self.no = no
        self.young = .0
        self.poisson = .0
        self.thickness = .0

class Element:
    def __init__(self,no):
        self.no = no
        self.node = [0,0,0] # 節点番号
        self.x = [.0,.0,.0] # 節点番号のX座標
        self.y = [.0,.0,.0] 
        self.u = [.0 for i in range(6)] # 節点番号の変位
        self.elNode = [0,0,0,0,0,0]
        self.matNo = 0
        self.area = .0
        self.Bmat = [[.0 for i in range(6)] for j in range(3)]
        self.Smat = [[.0 for i in range(6)] for j in range(3)]
        self.Dmat = [[.0 for i in range(3)] for j in range(3)]
        self.Kmat = [[.0 for i in range(6)] for j in range(6)]
        self.stress = [.0 for i in range(9)]   # 要素の応力

    def setDmat(self,mat):
        a = mat.young / (1+mat.poisson) / (1-mat.poisson)
        b = a * mat.poisson
        c = 0.5 * mat.young / (1+mat.poisson)
        self.Dmat = [[a,b,0,],
                     [b,a,0,],
                     [0,0,c]]

    def setBmat(self):
        self.elNode = [2*self.node[0],2*self.node[0]+1,
                       2*self.node[1],2*self.node[1]+1,
                       2*self.node[2],2*self.node[2]+1]
        Ai = self.y[1] - self.y[2]
        Aj = self.y[2] - self.y[0]
        Ak = self.y[0] - self.y[1]
        Bi = self.x[2] - self.x[1]
        Bj = self.x[0] - self.x[2]
        Bk = self.x[1] - self.x[0]
        dl = Bk * Aj - Ak * Bj
        self.area = 0.5 * dl
        self.Bmat = [[Ai/dl,    0,Aj/dl,    0,Ak/dl,    0],
                     [    0,Bi/dl,    0,Bj/dl,    0,Bk/dl],
                     [Bi/dl,Ai/dl,Bj/dl,Aj/dl,Bk/dl,Ak/dl]]
                     
    def setSmat(self):
        for i in range(3):
            for j in range(6):
                s = 0
                for k in range(3):
                    s += self.Dmat[i][k]*self.Bmat[k][j]
                self.Smat[i][j] = s
        
    def setKmat(self,mat): #要素剛性行列の設定
        v = mat.thickness * self.area
        for i in range(6):
            for j in range(6):
                s = 0
                for k in range(3):
                    s += self.Bmat[k][i]*self.Smat[k][j]
                self.Kmat[i][j] = v * s

    def calcStress(self):
        for i in range(3):
            s = 0
            for j in range(6):
                s += self.Smat[i][j] * self.u[j]
            self.stress[i] = s
        cc = 0.5 * (self.stress[0] + self.stress[1])
        ss = 0.5 * (self.stress[0] - self.stress[1])
        cr = math.sqrt(ss*ss+self.stress[2]*self.stress[2])
        self.stress[3] = cc + cr
        self.stress[4] = cc - cr
        self.stress[5] = cr
        if cr <= ss:
            self.stress[6]  = 0
        else:
            wk = abs((ss+cr)/(ss-cr))
            self.stress[6] = 57.29577951 * math.atan(math.sqrt(wk))
            if self.stress[6] < 0:
                self.stress[6] = -self.stress[6]
            self.stress[6] = 90 - self.stress[6]
        self.stress[7] = math.sqrt(self.stress[3]*self.stress[3]
                                  - self.stress[3] * self.stress[4]
                                  + self.stress[4] * self.stress[4])
        if self.stress[4] > 0:
            self.stress[8] = self.stress[3]
        elif self.stress[3] > 0:
            self.stress[8] = self.stress[3] - self.stress[4]
        else:
            self.stress[8] = - self.stress[4]

class Fem2D():
    
    def setNodeData(self):
        '''節点データ設定'''
        f = open('node.txt')
        lines = f.readlines()
        f.close()
        lines = lines[1:]
        
        self.x = []
        self.y = []
        self.force = []
        self.nodeCond = []
        self.disp = []
        
        n = 0
        for line in lines:
            line = line.strip('\n')
            d = line.split('\t')
            d = [x if x is not '' else 0 for x in d] + [0,0,0,0]
            
            self.x.append(float(d[1]))
            self.y.append(float(d[2]))
            self.nodeCond.append(int(d[3]))
            self.nodeCond.append(int(d[4]))
            self.disp.append(float(d[5]))
            self.disp.append(float(d[6]))
            self.force.append(float(d[7]))
            self.force.append(float(d[8]))
            n += 1
        self.numberOfNode = n
    
    def setMaterialData(self):
        '''材料データ設定'''
        self.mats = []
        m = Material(0)
        m.young = 20000.0
        m.poisson = 0.3
        m.thickness = 1.0
        self.mats.append(m)
        
    def setElementData(self):
        '''要素データ設定'''
        f = open('elem.txt')
        lines = f.readlines()
        lines = lines[1:]
        f.close()
        
        self.elems = []
        for line in lines:
            d = line.split('\t')
            no = int(d[0])-1
            e = Element(no)
            n0 = int(d[1])-1
            n1 = int(d[2])-1
            n2 = int(d[3])-1
            e.node = [n0,n1,n2]
            e.matNo = int(d[4])-1
            e.x[0] = self.x[n0] 
            e.x[1] = self.x[n1] 
            e.x[2] = self.x[n2] 
            e.y[0] = self.y[n0] 
            e.y[1] = self.y[n1] 
            e.y[2] = self.y[n2] 
            self.elems.append(e)
        self.numberOfElement = len(self.elems)
    
    def setElementMatrix(self):
        '''要素剛性行列設定'''
        for i in range(self.numberOfElement):
            n = self.elems[i].matNo
            self.elems[i].setDmat(self.mats[n])
            self.elems[i].setBmat()
            self.elems[i].setSmat()
            self.elems[i].setKmat(self.mats[n])

    def setBandWidth(self):
        n = 0
        for k in range(self.numberOfElement):
            for i in range(6):
                for j in range(6):
                    ii = self.elems[k].elNode[i]
                    jj = self.elems[k].elNode[j] - ii
                    if jj > n:
                        n = jj
        self.bandWidth = n
    
    def setTotalMatrix(self):
        '''全体剛性行列設定'''
        i = self.numberOfNode*2
        j = self.bandWidth+1
        self.totalMat = [[0 for a in range(j)] for b in range(i)]
        
        for k in range(self.numberOfElement):
            for i in range(6):
                for j in range(6):
                    ii = self.elems[k].elNode[i]
                    jj = self.elems[k].elNode[j]-ii
                    if jj >= 0:
                        self.totalMat[ii][jj] += self.elems[k].Kmat[i][j]
     
    def setBoundaryCondition(self):
        '''境界条件による全体剛性行列，荷重ベクトルの変更'''
        for i in range(self.numberOfNode*2):
            if self.nodeCond[i] != 0:
                self.totalMat[i][0] = 1e30
                self.force[i] = self.totalMat[i][0]*self.disp[i]
    
    def saveTotalMatrix(self,fp='matrix.txt'):
        '''全体剛性行列の保存'''
        f = open(fp,'w')
        for i in range(self.numberOfNode*2):
            for j in range(self.bandWidth+1):
                f.write(str(self.totalMat[i][j])+'\t')
            f.write('\n')          
        f.close()
        print('saveTotalMat: %s' % fp)

    def solver(self):
        '''連立方程式の解（ガウスの消去法）'''
        nfreedom = self.numberOfNode*2
        akj = ['' for i in range(nfreedom)]
        aik = ['' for i in range(nfreedom)]
        n1 = self.bandWidth + 1
        n2 = nfreedom 
        for k in range(n2):
            m = self.bandWidth + 1
            if k+self.bandWidth+1 > n2:
                m = n2 - k
            for j in range(1,m):
                akj[j] = self.totalMat[k][j]/self.totalMat[k][0]
                aik[j] = self.totalMat[k][j]
            for i in range(1,m):
                it = i + k
                im1 = i
                for j in range(i,m):
                    jt = j - im1
                    self.totalMat[it][jt] -= aik[i]*akj[j]
        for k in range(n2-1):
            self.force[k] = self.force[k]/self.totalMat[k][0]
            m = self.bandWidth + 1
            if (k+self.bandWidth+1)>n2:
                m = n2 - k
            for i in range(1,m):
                it = i + k
                self.force[it] -= self.totalMat[k][i]*self.force[k]
        self.disp[n2-1] = self.force[n2-1] / self.totalMat[n2-1][0]
        for k in range(n2-2,-1,-1):
            s = 0
            m = self.bandWidth + 1
            if (k+self.bandWidth+1)> n2:
                m = n2-k
            for j in range(1,m):
                jt = j + k
                s += self.totalMat[k][j]*self.disp[jt]
            self.disp[k] = self.force[k] - s / self.totalMat[k][0]
            
    def saveNodeDisp(self):
        '''求まった節点変位データの保存'''
        f = open('disp.txt','w')
        for i in range(self.numberOfNode):
            f.write('%d\t' % (i+1))
            f.write('%f\t' % self.disp[2*i])
            f.write('%f\n' % self.disp[2*i+1])
        f.close()
        print('saveNodeDisp: disp.txt')
    
    def setStress(self):
        '''応力の設定'''
        for i in range(self.numberOfElement):
            e = self.elems[i]
            for j in range(6):
                k = e.elNode[j]
                e.u[j] = self.disp[k]
            e.calcStress()

    def saveStress(self):
        '''求まった応力データの保存'''
        f = open('stress.txt','w')
        for i in range(self.numberOfElement):
            for j in range(9):
                f.write(str(self.elems[i].stress[j])+'\t')
            f.write('\n')
        f.close()
        print('saveStress: stress.txt')

    def main(self):    
        self.setNodeData()
        self.setMaterialData()
        self.setElementData()
        self.setElementMatrix()
        self.setBandWidth()
        self.setTotalMatrix()
        self.saveTotalMatrix()
        self.setBoundaryCondition()
        self.saveTotalMatrix('matrix2.txt')
        self.solver()
        self.saveNodeDisp()
        self.setStress()
        self.saveStress()

def main():
    fem = Fem2D()
    fem.main()

if __name__ == '__main__':
    main()