def heavyside_approx(H, alpha):
return 0.5 * (H / (th_adj.sqrt(H**2 + alpha**2))) + 0.5
def forward_next(eta_bc_const, mesh2d, wd_obs_arr, scale, scale_return=False):
def heavyside_approx(H, alpha):
return 0.5 * (H / (th_adj.sqrt(H**2 + alpha**2))) + 0.5
def update_forcings_hydrodynamics(t_new):
in_fn = (
(1 - dirk22_factor) *
eta_bc_const[int(t_new / options.timestep) - 1]) + (
dirk22_factor * eta_bc_const[int(t_new / options.timestep)])
eta_list.append(in_fn)
#print(int(t_new/options.timestep))
in_c.assign(in_fn)
if np.round(t_new % options.timestep, 2) == 0.00:
elev1 = (solver_obj.fields.solution_2d[2])
in_fn = eta_bc_const[int(t_new / options.timestep)]
eta_list.append(in_fn)
in_c.assign(in_fn)
wd_obs = th_adj.project(wd_obs_arr[int(t_new / dt)], P1_2d)
# Record the simulated wetting and drying front
wd_tmp = th_adj.project(
heavyside_approx(-elev1 - h_static, wetting_alpha), P1_2d)
wd.assign(wd_tmp)
form = 0.5 * th_adj.inner(wd - wd_obs, wd - wd_obs) * th_adj.dx
#form = 0.5*th_adj.inner(wd, wd)*th_adj.dx
#print(J_mc)
#J_list.append(th_adj.assemble(dt*scale*form))
J_list.append(th_adj.assemble(scale * options.timestep * form))
#J_list.append(th_adj.assemble(options.timestep*form))
wetting_and_drying = True
ks = 0.025
average_size = 200 * (10**(-6))
t_end = 24 * 3600.
friction = 'manning'
friction_coef = 0.025
diffusivity = 0.15
viscosity = 10**(-6)
dt = 600
# choose directory to output results
ts = time.time()
st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
outputdir = 'outputs' + st
# export interval in seconds
t_export = np.round(t_end / 40, 0)
th_adj.print_output('Exporting to ' + outputdir)
#mesh2d = th_adj.Mesh('mesh.msh')#th_adj.RectangleMesh(12, 6, 13800, 7200)
# define function spaces
V = th_adj.FunctionSpace(mesh2d, 'CG', 2)
P1_2d = th_adj.FunctionSpace(mesh2d, 'DG', 6)
x, y = th_adj.SpatialCoordinate(mesh2d)
max_depth = 5
basin_x = 13800
h_static = (1. - x / basin_x) * max_depth
bathymetry_2d = th_adj.Function(V)
bathymetry_2d.project((1. - x / basin_x) * max_depth)
#elev1 = th_adj.Function(P1_2d).interpolate(th_adj.Constant(0.0))
# define parameters
# ksp = th_adj.Constant(3*average_size)
# initial condition: assign non-zero velocity
elev_init = th_adj.Constant(0.01)
uv_init = th_adj.Constant((10**(-4), 0.))
wetting_alpha = 0.43
#wd_obs_arr =[(th_adj.project(heavyside_approx(- elev_init - h_static, wetting_alpha), P1_2d, annotate = False))]
wd_obs = th_adj.project(wd_obs_arr[0], P1_2d) #, annotate = False)
# Record the simulated wetting and drying front
wd = th_adj.project(heavyside_approx(-elev_init - h_static, wetting_alpha),
P1_2d)
eta_list = []
J_list = []
th_adj.parameters['form_compiler']['quadrature_degree'] = 20
th_adj.parameters['form_compiler']['cpp_optimize'] = True
th_adj.parameters['form_compiler'][
'cpp_optimize_flags'] = '-O3 -ffast-math -march=native'
th_adj.parameters['form_compiler']['optimize'] = True
# set up solver
solver_obj = th_adj.solver2d.FlowSolver2d(mesh2d, bathymetry_2d)
options = solver_obj.options
options.simulation_export_time = t_export
options.simulation_end_time = t_end
options.output_directory = outputdir
options.norm_smoother = th_adj.Constant(0.43)
options.check_volume_conservation_2d = True
options.fields_to_export = ['uv_2d', 'elev_2d']
options.solve_tracer = False
options.use_lax_friedrichs_tracer = False
if friction == 'nikuradse':
options.quadratic_drag_coefficient = cfactor
elif friction == 'manning':
if friction_coef == 0:
friction_coef = 0.02
options.manning_drag_coefficient = th_adj.Constant(friction_coef)
else:
print('Undefined friction')
# set horizontal diffusivity parameter
options.horizontal_diffusivity = th_adj.Constant(diffusivity)
if viscosity == None:
print('no viscosity')
else:
options.horizontal_viscosity = th_adj.Constant(viscosity)
# crank-nicholson used to integrate in time system of ODEs resulting from application of galerkin FEM
options.timestepper_type = 'DIRK22'
dirk22_factor = (2 - th_adj.sqrt(2)) / 2
#options.timestepper_options.implicitness_theta = 0.5
options.use_wetting_and_drying = wetting_and_drying
options.wetting_and_drying_alpha = th_adj.Constant(wetting_alpha)
if not hasattr(options.timestepper_options, 'use_automatic_timestep'):
options.timestep = dt
# set boundary conditions
in_constant = eta_bc_const[0]
in_c = th_adj.Constant(0.0)
in_c.assign(in_constant)
#in_c.assign(in_constant)
swe_bnd = {}
swe_bnd[1] = {'elev': in_c}
solver_obj.bnd_functions['shallow_water'] = swe_bnd
solver_obj.assign_initial_conditions(uv=uv_init, elev=elev_init)
#solver_obj.assign_initial_conditions(uv = uv_init)
solver_obj.iterate(update_forcings=update_forcings_hydrodynamics)
regularisation_sum_init = sum([
1 / (options.timestep) * ((eta_bc_const[i]) - eta_bc_const[i - 1])**2
for i in range(1,
len(eta_bc_const) - int(2 * 60 * 60 / 600))
])
regularisation_init = th_adj.AdjFloat(fac) * th_adj.AdjFloat(
scale) * regularisation_sum_init
#regularisation_init = fac * scale * sum([1/(options.timestep*len(th_adj.Function(R).dat.data[:]))*(th_adj.Function(R).assign(eta_bc_const[i])-th_adj.Function(R).assign(eta_bc_const[i-1]))**2*th_adj.dx for i in range(3, 10)])
#regularisation_out = fac * scale * sum([1/(options.timestep*len(th_adj.Function(R).dat.data[:]))*(th_adj.Function(R).assign(eta_bc_const[i])-th_adj.Function(R).assign(eta_bc_const[i-1]))**2*th_adj.dx for i in range(len(eta_bc_const)-10, len(eta_bc_const)-3)])
#import ipdb; ipdb.set_trace()
J = sum(
J_list) + regularisation_init # + th_adj.assemble(regularisation_out)
scale_tmp = th_adj.Constant(scale)
if scale_return:
return J, scale
else:
return J
