1 | from anuga import Domain |
---|
2 | #from anuga.utilities.sparse import Sparse, Sparse_CSR |
---|
3 | from sparse import Sparse, Sparse_CSR #the new import |
---|
4 | from anuga.utilities.cg_solve import conjugate_gradient |
---|
5 | import anuga.abstract_2d_finite_volumes.neighbour_mesh as neighbour_mesh |
---|
6 | from anuga.abstract_2d_finite_volumes.generic_boundary_conditions import Dirichlet_boundary |
---|
7 | import numpy as num |
---|
8 | import kinematic_viscosity_ext |
---|
9 | import anuga.utilities.log as log |
---|
10 | |
---|
11 | class Kinematic_Viscosity_Operator: |
---|
12 | |
---|
13 | def __init__(self, domain, triangle_areas=True, verbose=False): |
---|
14 | if verbose: log.critical('Kinematic Viscosity: Beginning Initialisation') |
---|
15 | #Expose the domain attributes |
---|
16 | self.domain = domain |
---|
17 | self.mesh = domain.mesh |
---|
18 | self.boundary = domain.boundary |
---|
19 | self.n = len(self.domain) |
---|
20 | self.dt = 1e-6 #Need to set to domain.timestep |
---|
21 | self.boundary_len = len(domain.boundary) |
---|
22 | self.tot_len = self.n+self.boundary_len |
---|
23 | self.boundary_enum = self.enumerate_boundary() |
---|
24 | self.verbose = verbose |
---|
25 | |
---|
26 | #Geometric Information |
---|
27 | if verbose: log.critical('Kinematic Viscosity: Building geometric structure') |
---|
28 | self.geo_structure_indices = num.zeros((self.n,3),num.int) |
---|
29 | self.geo_structure_values = num.zeros((self.n,3),num.float) |
---|
30 | kinematic_viscosity_ext.build_geo_structure(self,self.n,self.tot_len) |
---|
31 | self.apply_triangle_areas = triangle_areas |
---|
32 | self.triangle_areas = Sparse(self.n,self.n) |
---|
33 | for i in range(self.n): |
---|
34 | self.triangle_areas[i,i] = 1.0 / self.mesh.areas[i] |
---|
35 | self.triangle_areas = Sparse_CSR(self.triangle_areas) |
---|
36 | |
---|
37 | #Application Information |
---|
38 | self.qty_considered = 1 #1 or 2 (uh or vh respectively) |
---|
39 | self.operator_matrix = Sparse(self.n,self.tot_len) |
---|
40 | self.operator_data = num.zeros((4*self.n,),num.float) #Sparse_CSR.data |
---|
41 | self.operator_colind = num.zeros((4*self.n,),num.int) #Sparse_CSR.colind |
---|
42 | self.operator_rowptr = 4*num.indices((self.n+1,))[0,:] #Sparse_CSR.rowptr (4 entries in every row, we know this already) = [0,4,8,...,4*n] |
---|
43 | self.stage_heights = num.zeros((self.n,1),num.float) |
---|
44 | self.stage_heights_scaling = Sparse(self.n,self.n) |
---|
45 | self.boundary_vector = num.zeros((self.n,),num.float) |
---|
46 | |
---|
47 | self.parabolic_solve = False #Are we doing a parabolic solve at the moment? |
---|
48 | |
---|
49 | if verbose: log.critical('Kinematic Viscosity: Initialisation Done') |
---|
50 | |
---|
51 | # Enumerate the boundary conditions in some way |
---|
52 | def enumerate_boundary(self): |
---|
53 | #Just enumerate by the triangle number, then edge number |
---|
54 | enumeration = {} |
---|
55 | enum = 0 |
---|
56 | for i in range(self.n): |
---|
57 | for edge in range(3): |
---|
58 | if self.mesh.neighbours[i,edge] == -1: |
---|
59 | enumeration[(i,edge)] = enum |
---|
60 | enum += 1 |
---|
61 | return enumeration |
---|
62 | |
---|
63 | def set_qty_considered(self,qty): |
---|
64 | if qty == 1 or qty == 'u': |
---|
65 | self.qty_considered = 1 |
---|
66 | elif qty == 2 or qty == 'v': |
---|
67 | self.qty_considered = 2 |
---|
68 | else: #Raise an exception |
---|
69 | msg = "Incorrect input qty" |
---|
70 | assert 0==1, msg |
---|
71 | |
---|
72 | def apply_stage_heights(self,h): |
---|
73 | msg = "(KV_Operator.apply_stage_heights) h vector has incorrect length" |
---|
74 | assert h.size == self.n, msg |
---|
75 | |
---|
76 | self.operator_matrix = Sparse(self.n,self.tot_len) |
---|
77 | |
---|
78 | #Evaluate the boundary stage heights - use generic_domain.update_boundary? (check enumeration) |
---|
79 | boundary_heights = num.zeros((self.boundary_len,1),num.float) |
---|
80 | for key in self.boundary_enum.keys(): |
---|
81 | boundary_heights[self.boundary_enum[key]] = self.boundary[key].evaluate()[0] |
---|
82 | |
---|
83 | kinematic_viscosity_ext.build_operator_matrix(self,self.n,self.tot_len,h,boundary_heights) |
---|
84 | self.operator_matrix = Sparse_CSR(None,self.operator_data,self.operator_colind,self.operator_rowptr,self.n,self.tot_len) |
---|
85 | |
---|
86 | self.stage_heights = h |
---|
87 | #Set up the scaling matrix |
---|
88 | data = h |
---|
89 | num.putmask(data, data!=0, 1/data) #take the reciprocal of each entry unless it is zero |
---|
90 | self.stage_heights_scaling = Sparse_CSR(None,data,num.arange(self.n),num.append(num.arange(self.n),self.n),self.n,self.n) |
---|
91 | |
---|
92 | self.build_boundary_vector() |
---|
93 | |
---|
94 | return self.operator_matrix |
---|
95 | |
---|
96 | def build_boundary_vector(self): |
---|
97 | #Require stage heights applied |
---|
98 | X = num.zeros((self.tot_len,2),num.float) |
---|
99 | for key in self.boundary_enum.keys(): |
---|
100 | quantities = self.boundary[key].evaluate() |
---|
101 | h_i = quantities[0] |
---|
102 | if h_i == 0.0: |
---|
103 | X[self.n + self.boundary_enum[key],:] = num.array([0.0, 0.0]) |
---|
104 | else: |
---|
105 | X[self.n + self.boundary_enum[key],:] = quantities[1:] / h_i |
---|
106 | self.boundary_vector = self.operator_matrix * X |
---|
107 | #Tidy up |
---|
108 | if self.apply_triangle_areas: |
---|
109 | self.boundary_vector = self.triangle_areas * self.boundary_vector |
---|
110 | |
---|
111 | return self.boundary_vector |
---|
112 | |
---|
113 | def elliptic_multiply(self,V,qty_considered=None,include_boundary=True): |
---|
114 | msg = "(KV_Operator.apply_vector) V vector has incorrect dimensions" |
---|
115 | assert V.shape == (self.n,) or V.shape == (self.n,1), msg |
---|
116 | |
---|
117 | if qty_considered!=None: |
---|
118 | print "Quantity Considered changed!" |
---|
119 | self.set_qty_considered(qty_considered) |
---|
120 | |
---|
121 | D = num.zeros((self.n,),num.float) |
---|
122 | #Change (uh) to (u), (vh) to (v) |
---|
123 | X = self.stage_heights_scaling * V |
---|
124 | |
---|
125 | if self.apply_triangle_areas: |
---|
126 | X = self.triangle_areas * X |
---|
127 | |
---|
128 | X = num.append(X,num.zeros((self.boundary_len,)),axis=0) |
---|
129 | D = self.operator_matrix * X |
---|
130 | |
---|
131 | if include_boundary: |
---|
132 | D += self.boundary_vector[:,self.qty_considered-1] |
---|
133 | |
---|
134 | D = self.stage_heights * D #make sure we return (uh) not (u), for the solver's benefit |
---|
135 | |
---|
136 | return num.array(D).reshape(self.n,) |
---|
137 | |
---|
138 | #parabolic_multiply(V) = identity - dt*elliptic_multiply(V) |
---|
139 | #We do not need include boundary, as we will only use this method for solving (include_boundary = False) |
---|
140 | #Here V is already either (u) or (v), not (uh) or (vh) |
---|
141 | def parabolic_multiply(self,V): |
---|
142 | msg = "(KV_Operator.parabolic_multiply) V vector has incorrect dimensions" |
---|
143 | assert V.shape == (self.n,) or V.shape == (self.n,1), msg |
---|
144 | |
---|
145 | D = V #The return value |
---|
146 | X = V #a temporary array |
---|
147 | |
---|
148 | #Apply triangle areas |
---|
149 | if self.apply_triangle_areas: |
---|
150 | X = self.triangle_areas * X |
---|
151 | |
---|
152 | #Multiply out |
---|
153 | X = num.append(X,num.zeros((self.boundary_len,)),axis=0) |
---|
154 | D = D - self.dt * (self.operator_matrix * X) |
---|
155 | |
---|
156 | return num.array(D).reshape(self.n,) |
---|
157 | |
---|
158 | def __mul__(self,other): |
---|
159 | try: |
---|
160 | B = num.array(other) |
---|
161 | except: |
---|
162 | msg = "Trying to multiply the Kinematic Viscosity Operator against a non-numeric type" |
---|
163 | raise msg |
---|
164 | |
---|
165 | if len(B.shape) == 0: |
---|
166 | #Scalar |
---|
167 | R = B*self |
---|
168 | elif len(B.shape) == 1 or (len(B.shape) == 2 and B.shape[1]==1): |
---|
169 | #Vector |
---|
170 | if self.parabolic_solve: |
---|
171 | R = self.parabolic_multiply(other) |
---|
172 | else: |
---|
173 | #include_boundary=False is this is *only* used for cg_solve() |
---|
174 | R = self.elliptic_multiply(other,include_boundary=False) |
---|
175 | else: |
---|
176 | raise ValueError, 'Dimension too high: d=%d' %len(B.shape) |
---|
177 | return R |
---|
178 | |
---|
179 | def __rmul__(self, other): |
---|
180 | #Right multiply with scalar |
---|
181 | try: |
---|
182 | other = float(other) |
---|
183 | except: |
---|
184 | msg = 'Sparse matrix can only "right-multiply" onto a scalar' |
---|
185 | raise TypeError, msg |
---|
186 | else: |
---|
187 | new = self.operator_matrix * new |
---|
188 | return new |
---|
189 | |
---|
190 | def cg_solve(self,B,qty_to_consider=None): |
---|
191 | if len(B.shape) == 1: |
---|
192 | return self.cg_solve_vector(B,qty_to_consider) |
---|
193 | elif len(B.shape) == 2: |
---|
194 | return self.cg_solve_matrix(B) |
---|
195 | else: |
---|
196 | raise ValueError, 'Dimension too high: d=%d' %len(B.shape) |
---|
197 | |
---|
198 | def cg_solve_matrix(self,B): |
---|
199 | assert B.shape[1] < 3, "Input matrix has too many columns (max 2)" |
---|
200 | |
---|
201 | X = num.zeros(B.shape,num.float) |
---|
202 | for i in range(B.shape[1]): |
---|
203 | X[:,i] = self.cg_solve_vector(B[:,i],i+1) #assuming B columns are (uh) and (vh) |
---|
204 | return X |
---|
205 | |
---|
206 | def cg_solve_vector(self,b,qty_to_consider=None): |
---|
207 | if not qty_to_consider==None: |
---|
208 | self.set_qty_considered(qty_to_consider) |
---|
209 | x = conjugate_gradient(self,b - self.boundary_vector[:,self.qty_considered-1]) #Call the ANUGA conjugate gradient utility |
---|
210 | return x |
---|
211 | |
---|
212 | #Solve the parabolic equation to find u^(n+1), where u = u^n |
---|
213 | #Here B = [uh vh] where uh,vh = n*1 column vectors |
---|
214 | def parabolic_solver(self,B): |
---|
215 | assert B.shape == (self.n,2), "Input matrix has incorrect dimensions" |
---|
216 | |
---|
217 | next_B = num.zeros((self.n,2),num.float) |
---|
218 | #The equation is self.parabolic_multiply*u^(n+1) = u^n + (dt)*self.boundary_vector |
---|
219 | self.parabolic_solve = True |
---|
220 | #Change (uh) and (vh) to (u) and (v) |
---|
221 | b = self.stage_heights_scaling * B |
---|
222 | |
---|
223 | next_B = conjugate_gradient(self, b + (self.dt * self.boundary_vector), iprint=1) |
---|
224 | |
---|
225 | #Return (u) and (v) back to (uh) and (vh) |
---|
226 | next_B = self.stage_heights_scaling * next_B |
---|
227 | |
---|
228 | self.parabolic_solve = False |
---|
229 | |
---|
230 | return next_B |
---|