00001
00002 from sympy import *
00003
00004
00005 x,y,z = symbols('xyz')
00006
00007 class ReferenceSimplex:
00008 def __init__(self, nsd):
00009 self.nsd = nsd
00010 coords = []
00011 if nsd <= 3:
00012 coords = symbols('xyz')[:nsd]
00013 else:
00014 coords = []
00015 for d in range(0,nsd):
00016 coords.append(Symbol("x_%d" % d))
00017 self.coords = coords
00018
00019 def integrate(self,f):
00020 coords = self.coords
00021 nsd = self.nsd
00022
00023 limit = 1
00024 for p in coords:
00025 limit -= p
00026
00027 intf = f
00028 for d in range(0,nsd):
00029 p = coords[d]
00030 limit += p
00031 intf = integrate(intf, (p, 0, limit))
00032 return intf
00033
00034 def bernstein_space(order, nsd):
00035 if nsd > 3:
00036 raise RuntimeError("Bernstein only implemented in 1D, 2D, and 3D")
00037 sum = 0
00038 basis = []
00039 coeff = []
00040
00041 if nsd == 1:
00042 b1, b2 = x, 1-x
00043 for o1 in range(0,order+1):
00044 for o2 in range(0,order+1):
00045 if o1 + o2 == order:
00046 aij = Symbol("a_%d_%d" % (o1,o2))
00047 sum += aij*binomial(order,o1)*pow(b1, o1)*pow(b2, o2)
00048 basis.append(binomial(order,o1)*pow(b1, o1)*pow(b2, o2))
00049 coeff.append(aij)
00050
00051
00052 if nsd == 2:
00053 b1, b2, b3 = x, y, 1-x-y
00054 for o1 in range(0,order+1):
00055 for o2 in range(0,order+1):
00056 for o3 in range(0,order+1):
00057 if o1 + o2 + o3 == order:
00058 aij = Symbol("a_%d_%d_%d" % (o1,o2,o3))
00059 fac = factorial(order)/(factorial(o1)*factorial(o2)*factorial(o3))
00060 sum += aij*fac*pow(b1, o1)*pow(b2, o2)*pow(b3, o3)
00061 basis.append(fac*pow(b1, o1)*pow(b2, o2)*pow(b3, o3))
00062 coeff.append(aij)
00063
00064 if nsd == 3:
00065 b1, b2, b3, b4 = x, y, z, 1-x-y-z
00066 for o1 in range(0,order+1):
00067 for o2 in range(0,order+1):
00068 for o3 in range(0,order+1):
00069 for o4 in range(0,order+1):
00070 if o1 + o2 + o3 + o4 == order:
00071 aij = Symbol("a_%d_%d_%d_%d" % (o1,o2,o3,o4))
00072 fac = factorial(order)/(factorial(o1)*factorial(o2)*factorial(o3)*factorial(o4))
00073 sum += aij*fac*pow(b1, o1)*pow(b2, o2)*pow(b3, o3)*pow(b4, o4)
00074 basis.append(fac*pow(b1, o1)*pow(b2, o2)*pow(b3, o3)*pow(b4, o4))
00075 coeff.append(aij)
00076
00077
00078 return sum, coeff, basis
00079
00080 def create_point_set(order, nsd):
00081 h = Rational(1,order)
00082 set = []
00083
00084 if nsd == 1:
00085 for i in range(0, order+1):
00086 x = i*h
00087 if x <= 1:
00088 set.append((x,y))
00089
00090 if nsd == 2:
00091 for i in range(0, order+1):
00092 x = i*h
00093 for j in range(0, order+1):
00094 y = j*h
00095 if x + y <= 1:
00096 set.append((x,y))
00097
00098 if nsd == 3:
00099 for i in range(0, order+1):
00100 x = i*h
00101 for j in range(0, order+1):
00102 y = j*h
00103 for k in range(0, order+1):
00104 z = j*h
00105 if x + y + z <= 1:
00106 set.append((x,y,z))
00107
00108 return set
00109
00110
00111
00112 def create_matrix(equations, coeffs):
00113 A = zeronm(len(equations), len(equations))
00114 i = 0; j = 0
00115 for j in range(0, len(coeffs)):
00116 c = coeffs[j]
00117 for i in range(0, len(equations)):
00118 e = equations[i]
00119 d, r = div(e, c)
00120 A[i,j] = d
00121 return A
00122
00123
00124
00125 class Lagrange:
00126 def __init__(self,nsd, order):
00127 self.nsd = nsd
00128 self.order = order
00129 self.compute_basis()
00130
00131 def nbf(self):
00132 return len(self.N)
00133
00134 def compute_basis(self):
00135 order = self.order
00136 nsd = self.nsd
00137 N = []
00138 pol, coeffs, basis = bernstein_space(order, nsd)
00139 points = create_point_set(order, nsd)
00140
00141 equations = []
00142 for p in points:
00143 ex = pol.subs(x, p[0])
00144 if nsd > 1:
00145 ex = ex.subs(y, p[1])
00146 if nsd > 2:
00147 ex = ex.subs(z, p[2])
00148 equations.append(ex )
00149
00150 A = create_matrix(equations, coeffs)
00151 Ainv = A.inv()
00152
00153 b = eye(len(equations))
00154
00155 xx = Ainv*b
00156
00157 for i in range(0,len(equations)):
00158 Ni = pol
00159 for j in range(0,len(coeffs)):
00160 Ni = Ni.subs(coeffs[j], xx[j,i])
00161 N.append(Ni)
00162
00163 self.N = N
00164
00165
00166
00167
00168
00169
00170
00171 if __name__ == '__main__':
00172 t = ReferenceSimplex(2)
00173 x,y,z = symbols('xyz')
00174 f = x+y
00175
00176 print "integral of f = x+y ", t.integrate(f)
00177 print ""
00178 print "linear bernstein in 1D ", bernstein_space(1,1)
00179 print "linear bernstein in 2D ", bernstein_space(1,2)
00180 print "linear bernstein in 3D ", bernstein_space(1,3)
00181 print ""
00182
00183
00184 print "basis functions of a second order Lagrange element"
00185 fe = Lagrange(2,2)
00186 for i in range(0, fe.nbf()):
00187 print "N[%d]= %s " % (i,fe.N[i])
00188
00189
00190 print "The Jacobian matrix of Fi = u * Ni "
00191
00192 u = 0
00193
00194 us = []
00195 for i in range(0, fe.nbf()):
00196 ui = Symbol("u_%d" % i)
00197 us.append(ui)
00198 u += ui*fe.N[i]
00199
00200
00201 J = zeronm(fe.nbf(), fe.nbf())
00202 for i in range(0, fe.nbf()):
00203 Fi = u*fe.N[i]
00204 for j in range(0, fe.nbf()):
00205 uj = us[j]
00206 integrands = diff(Fi, uj)
00207 J[j,i] = t.integrate(integrands)
00208
00209
00210 print J
00211
00212