Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add and implement augmentation with vorticity transport
- Loading branch information
1 parent
cdd2982
commit 50ef06b
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,159 @@ | ||
| from abc import ABCMeta | ||
| from firedrake import ( | ||
| MixedFunctionSpace, Function, TestFunctions, split, inner, dx, grad, | ||
| LinearVariationalProblem, LinearVariationalSolver, lhs, rhs, dot, | ||
| ds_b, ds_v, ds_t, ds, FacetNormal, TestFunction, TrialFunction, | ||
| transpose, nabla_grad, outer, dS, dS_h, dS_v, sign, jump, div, | ||
| Constant, sqrt, cross, curl | ||
| ) | ||
| from firedrake.fml import subject | ||
| from gusto import ( | ||
| time_derivative, transport, transporting_velocity, TransportEquationType, | ||
| logger | ||
| ) | ||
|
|
||
|
|
||
| class Augmentation(object, metaclass=ABCMeta): | ||
| """ | ||
| Augments an equation with another equation to be solved simultaneously. | ||
| """ | ||
|
|
||
|
|
||
| class VorticityTransport(Augmentation): | ||
| """ | ||
| Solves the transport of a velocity field, simultaneously with the vorticity. | ||
| """ | ||
|
|
||
| ### An argument to time discretisation or spatial method?? | ||
| # TODO: this all needs to be generalised | ||
|
|
||
| def __init__(self, domain, V_vel, V_vort, transpose_commutator=True, | ||
| supg=False, min_dx=None): | ||
|
|
||
| self.fs = MixedFunctionSpace((V_vel, V_vort)) | ||
| self.X = Function(self.fs) | ||
| self.tests = TestFunctions(self.fs) | ||
|
|
||
| u = Function(V_vel) | ||
| F, Z = split(self.X) | ||
| test_F, test_Z = self.tests | ||
|
|
||
| if hasattr(domain.mesh, "_base_mesh"): | ||
| self.ds = ds_b + ds_t + ds_v | ||
| self.dS = dS_v + dS_h | ||
| else: | ||
| self.ds = ds | ||
| self.dS = dS | ||
|
|
||
| n = FacetNormal(domain.mesh) | ||
| sign_u = 0.5*(sign(dot(u, n)) + 1) | ||
| upw = lambda f: (sign_u('+')*f('+') + sign_u('-')*f('-')) | ||
|
|
||
| if domain.mesh.topological_dimension() == 2: | ||
| mix_test = test_F - domain.perp(grad(test_Z)) | ||
| F_cross_u = Z*domain.perp(u) | ||
| elif domain.mesh.topological_dimension == 3: | ||
| mix_test = test_F - curl(test_Z) | ||
| F_cross_u = cross(Z, u) | ||
|
|
||
| time_deriv_form = inner(F, test_F)*dx + inner(Z, test_Z)*dx | ||
|
|
||
| # Standard vector invariant transport form ----------------------------- | ||
| transport_form = ( | ||
| # vorticity term | ||
| inner(mix_test, F_cross_u)*dx | ||
| + inner(n, test_Z*Z*u)*self.ds | ||
| # 0.5*grad(v . F) | ||
| - 0.5 * div(mix_test) * inner(u, F)*dx | ||
| + 0.5 * inner(mix_test, n) * inner(u, F)*self.ds | ||
| ) | ||
|
|
||
| # Communtator of tranpose gradient terms ------------------------------- | ||
| # This is needed for general vector transport | ||
| if transpose_commutator: | ||
| u_dot_nabla_F = dot(u, transpose(nabla_grad(F))) | ||
| transport_form += ( | ||
| - inner(n, test_Z*domain.perp(u_dot_nabla_F))*self.ds | ||
| # + 0.5*grad(F).v | ||
| - 0.5 * dot(F, div(outer(u, mix_test)))*dx | ||
| + 0.5 * inner(mix_test('+'), n('+'))*dot(jump(u), upw(F))*self.dS | ||
| # - 0.5*grad(v).F | ||
| + 0.5 * dot(u, div(outer(F, mix_test)))*dx | ||
| - 0.5 * inner(mix_test('+'), n('+'))*dot(jump(F), upw(u))*self.dS | ||
| ) | ||
|
|
||
| # SUPG terms ----------------------------------------------------------- | ||
| # Add the vorticity residual to the transported vorticity, | ||
| # which damps enstrophy | ||
| if supg: | ||
| if min_dx is not None: | ||
| lamda = Constant(0.5) | ||
| #TODO: decide on expression here | ||
| # tau = 0.5 / (lamda/domain.dt + sqrt(dot(u, u))/Constant(min_dx)) | ||
| tau = 0.5*domain.dt*(1.0 + sqrt(dot(u, u))*domain.dt/Constant(min_dx)) | ||
| else: | ||
| tau = 0.5*domain.dt | ||
|
|
||
| dxqp = dx(degree=3) | ||
|
|
||
| if domain.mesh.topological_dimension() == 2: | ||
| time_deriv_form -= inner(mix_test, tau*Z*domain.perp(u)/domain.dt)*dxqp | ||
| transport_form -= inner( | ||
| mix_test, tau*domain.perp(u)*domain.divperp(Z*domain.perp(u)) | ||
| )*dxqp | ||
| if transpose_commutator: | ||
| transport_form -= inner( | ||
| mix_test, | ||
| tau*domain.perp(u)*domain.divperp(u_dot_nabla_F) | ||
| )*dxqp | ||
| elif domain.mesh.topological_dimension() == 3: | ||
| time_deriv_form -= inner(mix_test, tau*cross(Z, u)/domain.dt)*dxqp | ||
| transport_form -= inner( | ||
| mix_test, tau*cross(curl(Z*u), u) | ||
| )*dxqp | ||
| if transpose_commutator: | ||
| transport_form -= inner( | ||
| mix_test, | ||
| tau*cross(curl(u_dot_nabla_F), u) | ||
| )*dxqp | ||
|
|
||
| residual = ( | ||
| time_derivative(time_deriv_form) | ||
| + transport( | ||
| transport_form, TransportEquationType.vector_invariant | ||
| ) | ||
| ) | ||
| residual = transporting_velocity(residual, u) | ||
|
|
||
| self.residual = subject(residual, self.X) | ||
|
|
||
| self.x_in = Function(self.fs) | ||
| self.Z_in = Function(V_vort) | ||
| self.x_out = Function(self.fs) | ||
|
|
||
| vort_test = TestFunction(V_vort) | ||
| vort_trial = TrialFunction(V_vort) | ||
|
|
||
| F_in, _ = split(self.x_in) | ||
|
|
||
| eqn = ( | ||
| inner(vort_trial, vort_test)*dx | ||
| + inner(domain.perp(grad(vort_test)), F_in)*dx | ||
| + vort_test*inner(n, domain.perp(F_in))*self.ds | ||
| ) | ||
| problem = LinearVariationalProblem( | ||
| lhs(eqn), rhs(eqn), self.Z_in, constant_jacobian=True | ||
| ) | ||
| self.solver = LinearVariationalSolver(problem) | ||
|
|
||
| def pre_apply(self, x_in): | ||
| self.x_in.subfunctions[0].assign(x_in) | ||
|
|
||
| def post_apply(self, x_out): | ||
| x_out.assign(self.x_out.subfunctions[0]) | ||
|
|
||
| def update(self, x_in_mixed): | ||
| self.x_in.subfunctions[0].assign(x_in_mixed.subfunctions[0]) | ||
| logger.info('Vorticity solve') | ||
| self.solver.solve() | ||
| self.x_in.subfunctions[1].assign(self.Z_in) | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.