Class version of Spring FEM in Python

Was cleaning the laptop and found this class version of the spring FEM code that I have posted earlier.

Not the best implementation but Posting it here just in case…

"""
Created on Sun May 11 08:55:27 2014

@author: Sukhbinder Singh


FEM Spring Class version

"""

import numpy as np

class SpringProb():
    def __init__(self,elemNodes,presDof,SpringStiff,force):
        self.elemNodes=elemNodes
        self.numElem =elemNodes.shape[0]
        self.numNodes= len(set(elemNodes.flatten()))
        self.presDof=presDof
        self.k=SpringStiff # Spring Stiffness
        self.disp=np.zeros((self.numNodes,1))
        self.force=force.T
        self.stiffness=np.zeros((self.numNodes,self.numNodes))
        self._solve=None


    def _calstiffness(self):
        for e in xrange(self.numElem):
            self.stiffness[np.ix_(self.elemNodes[e,:],self.elemNodes[e,:])]+=np.array([[self.k[e],-self.k[e]],[-self.k[e],self.k[e]]])

    def solve(self):
        self._calstiffness()
        actDof=np.setdiff1d(np.arange(self.numNodes),self.presDof)
        disp1=np.linalg.solve(self.stiffness[np.ix_(actDof,actDof)],self.force[np.ix_(actDof)]);
        self.disp[np.ix_(actDof)]=disp1
    # reactions
        react = np.dot(self.stiffness,self.disp)
        self.reaction = react[self.presDof]
        self._solve = True

    def printResults(self):
        print "\nDisplacement"
        for i,d in enumerate(self.disp):
            print i,d[0]
# Reactions
        if self._solve:
             print "\nReactions"
             for (p,r) in zip(self.presDof,self.reaction):
                print p,r[0]
        
        
if __name__ == '__main__':
    sp1 = SpringProb(np.array([[0,1],[1,2],[1,3]]),np.array([0,2,3]),np.array([10,10,10]),np.array([[0.0,100.0,0.0,0.0]]))
    sp1.solve()
    sp1.printResults()

    
    sp2 = SpringProb(np.array([[0,3],[0,1],[1,2],[1,2],[1,3],[2,3]]),np.array([0,3]),np.array([10,15,20,25,30,35]),np.array([[0.0,100.0,0.0,0.0]]))
    sp2.solve()
    sp2.printResults()