source: trunk/anuga_work/development/2010-projects/kv/kinematic_viscosity.py @ 8051

Last change on this file since 8051 was 8051, checked in by steve, 9 years ago

Added in Lindon's code

File size: 8.3 KB
Line 
1from anuga.abstract_2d_finite_volumes.domain import Domain
2#from anuga.utilities.sparse import Sparse, Sparse_CSR
3from sparse import Sparse, Sparse_CSR #the new import
4from anuga.utilities.cg_solve import conjugate_gradient
5import anuga.abstract_2d_finite_volumes.neighbour_mesh as neighbour_mesh
6from anuga.abstract_2d_finite_volumes.generic_boundary_conditions import Dirichlet_boundary
7import numpy as num
8import kinematic_viscosity_ext
9import anuga.utilities.log as log
10
11class 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],iprint=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
Note: See TracBrowser for help on using the repository browser.