[8052] | 1 | from anuga import Domain |
---|
[8051] | 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) |
---|
[8112] | 209 | x = conjugate_gradient(self,b - self.boundary_vector[:,self.qty_considered-1]) #Call the ANUGA conjugate gradient utility |
---|
[8051] | 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 | |
---|
[8052] | 230 | return next_B |
---|