-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathlinear_stommel_qg.py
85 lines (66 loc) · 2.58 KB
/
linear_stommel_qg.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
from firedrake import *
import numpy as np
<<<<<<< HEAD
import matplotlib.pyplot as plt
=======
#import matplotlib.pyplot as plt
>>>>>>> 61ee2dee86887d2fb8938c6229f63272ab31d0f1
# Geometry
Lx = 1. # Zonal length
Ly = 1. # Meridonal length
n0 = 50 # Spatial resolution
mesh = RectangleMesh(n0, n0, Lx, Ly, reorder=None)
# Function and Vector Spaces
Vcg = FunctionSpace(mesh,"CG",3) # CG elements for Streamfunction
Vdg = FunctionSpace(mesh,"DG",1) # DG elements for Potential Vorticity (PV)
Vcdg = VectorFunctionSpace(mesh,"DG",2) # DG elements for velocity
<<<<<<< HEAD
#psi0 = Function(Vcg, name='Streamfunction')
=======
>>>>>>> 61ee2dee86887d2fb8938c6229f63272ab31d0f1
# Boundary Conditions
bc = DirichletBC(Vcg, 0.0, (1, 2, 3, 4))
# Physical parameters
beta = Constant("1.0") # Beta parameter
F = Constant("1.0") # Burger number
r = Constant("0.2") # Bottom drag
tau = Constant("0.001") # Wind Forcing
# Test Functions
phi, p, v = TestFunction(Vcg), TestFunction(Vdg), TestFunction(Vcdg)
# Trial Functions
psi, q, u = TrialFunction(Vcg), TrialFunction(Vdg), TrialFunction(Vcdg)
# Solution Functions
psi_soln, u_soln = Function(Vcg, name="Streamfunction"), Function(Vcdg)
q0_soln, qn_soln = Function(Vdg), Function(Vcdg)
# Define Winds and Weak Form
Fwinds = Function(Vcg).interpolate(Expression("-tau*cos(pi*(x[1]-0.5))", tau=tau))
a = -r*inner(grad(psi), grad(phi))*dx - F*psi*phi*dx + beta*phi*psi.dx(0)*dx
L = Fwinds*phi*dx
#Question: Why does this not solve this correctly?
"""
# Set up Elliptic inverter
psi_problem = LinearVariationalProblem(a, L, psi_soln, bcs=bc)
psi_solver = LinearVariationalSolver(psi_problem,
solver_parameters={
'ksp_type':'preonly',
'pc_type':'lu'
})
"""
solve(a == L, psi_soln, bcs=bc)
#Question: Why does this not plot the solution?
#plt.plot(psi_soln)
#plt.show()
<<<<<<< HEAD
=======
# solve for streamfunction
psi_solver.solve()
# Plot Solution
p = plot(psi_soln)
p.show()
>>>>>>> 61ee2dee86887d2fb8938c6229f63272ab31d0f1
# Potential Energy
potential_energy = assemble(0.5*psi_soln*psi_soln*dx)
print potential_energy
# Output to a file
outfile = File("outputQG.pvd")
outfile.write(psi_soln)