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