def step_rk4(self,t,u0,v0,h0,He=None,hbcx=None,hbcy=None,first=False):
#######################
# Init local state #
#######################
u1 = deepcopy(u0)
v1 = deepcopy(v0)
h1 = deepcopy(h0)
#######################
# Init bc #
#######################
if first:
init_bc(self,u1,v1,h1,He,hbcx,hbcy,t0=t)
#######################
# Right hand sides #
#######################
# k1
ku1 = self.rhs_u(self.v_on_u(v1),h1)*self.dt
kv1 = self.rhs_v(self.u_on_v(u1),h1)*self.dt
kh1 = self.rhs_h(u1,v1,He)*self.dt
# k2
ku2 = self.rhs_u(self.v_on_u(v1+0.5*kv1),h1+0.5*kh1)*self.dt
kv2 = self.rhs_v(self.u_on_v(u1+0.5*ku1),h1+0.5*kh1)*self.dt
kh2 = self.rhs_h(u1+0.5*ku1,v1+0.5*kv1,He)*self.dt
# k3
ku3 = self.rhs_u(self.v_on_u(v1+0.5*kv2),h1+0.5*kh2)*self.dt
kv3 = self.rhs_v(self.u_on_v(u1+0.5*ku2),h1+0.5*kh2)*self.dt
kh3 = self.rhs_h(u1+0.5*ku2,v1+0.5*kv2,He)*self.dt
# k4
ku4 = self.rhs_u(self.v_on_u(v1+kv3),h1+kh3)*self.dt
kv4 = self.rhs_v(self.u_on_v(u1+ku3),h1+kh3)*self.dt
kh4 = self.rhs_h(u1+ku3,v1+kv3,He)*self.dt
#######################
# Time propagation #
#######################
u = u1 + 1/6*(ku1+2*ku2+2*ku3+ku4)
v = v1 + 1/6*(kv1+2*kv2+2*kv3+kv4)
h = h1 + 1/6*(kh1+2*kh2+2*kh3+kh4)
#######################
# Boundary conditions #
#######################
if hbcx is not None and hbcy is not None:
obcs(self,t,u,v,h,u1,v1,h1,He,hbcx,hbcy)
return u,v,h
def step_lf(self,t,u0,v0,h0,He=None,hbcx=None,hbcy=None,first=False):
#######################
# Init local state #
#######################
u1 = deepcopy(u0)
v1 = deepcopy(v0)
h1 = deepcopy(h0)
#######################
# Init bc #
#######################
if first:
init_bc(self,u1,v1,h1,He,hbcx,hbcy,t0=t)
#######################
# Right hand sides #
#######################
ku = self.rhs_u(self.v_on_u(v1),h1)
kv = self.rhs_v(self.u_on_v(u1),h1)
kh = self.rhs_h(u1,v1,He)
#######################
# Time propagation #
#######################
if first:
# forward euler
u = u1 + self.dt*ku
v = v1 + self.dt*kv
h = h1 + self.dt*kh
else:
# leap-frog
u = self._pu + 2*self.dt*ku
v = self._pv + 2*self.dt*kv
h = self._ph + 2*self.dt*kh
# For next timestep
self._pu = u1
self._pv = v1
self._ph = h1
#######################
# Boundary conditions #
#######################
if hbcx is not None and hbcy is not None:
obcs(self,t,u,v,h,u1,v1,h1,He,hbcx,hbcy)
return u,v,h
def step_euler(self,t,u0,v0,h0,He=None,hbcx=None,hbcy=None,first=False):
#######################
# Init local state #
#######################
u1 = deepcopy(u0)
v1 = deepcopy(v0)
h1 = deepcopy(h0)
#######################
# Init bc #
#######################
if first:
init_bc(self,u1,v1,h1,He,hbcx,hbcy,t0=t)
#######################
# Right hand sides #
#######################
ku = self.rhs_u(self.v_on_u(v1),h1)
kv = self.rhs_v(self.u_on_v(u1),h1)
kh = self.rhs_h(u1,v1,He)
#######################
# Time propagation #
#######################
u = u1 + self.dt*ku
v = v1 + self.dt*kv
h = h1 + self.dt*kh
#######################
# Boundary conditions #
#######################
if hbcx is not None and hbcy is not None:
obcs(self,t,u,v,h,u1,v1,h1,He,hbcx,hbcy)
return u,v,h
def step_rk4_adj(self,
t,
adu0,
adv0,
adh0,
u0,
v0,
h0,
He=None,
hbcx=None,
hbcy=None,
first=False):
########################
# Init #
########################
u1 = deepcopy(u0)
v1 = deepcopy(v0)
h1 = deepcopy(h0)
adu = 0
adv = 0
adh = 0
adHe = 0
adhbcx = 0
adhbcy = 0
if first:
init_bc(self, u1, v1, h1, He, hbcx, hbcy, t0=t)
#######################
# Boundary conditions #
#######################
adu_tmp,adv_tmp,adh_tmp,adHe_tmp,adhbcx_tmp,adhbcy_tmp = \
obcs_adj(self,t,adu0,adv0,adh0,u1,v1,h1,He,hbcx,hbcy)
adu += adu_tmp
adv += adv_tmp
adh += adh_tmp
adHe += adHe_tmp
adhbcx += adhbcx_tmp
adhbcy += adhbcy_tmp
#######################
# Time propagation #
#######################
# k1
ku1 = self.rhs_u(self.v_on_u(v1), h1) * self.dt
kv1 = self.rhs_v(self.u_on_v(u1), h1) * self.dt
kh1 = self.rhs_h(u1, v1, He) * self.dt
# k2
ku2 = self.rhs_u(self.v_on_u(v1 + 0.5 * kv1), h1 + 0.5 * kh1) * self.dt
kv2 = self.rhs_v(self.u_on_v(u1 + 0.5 * ku1), h1 + 0.5 * kh1) * self.dt
kh2 = self.rhs_h(u1 + 0.5 * ku1, v1 + 0.5 * kv1, He) * self.dt
# k3
ku3 = self.rhs_u(self.v_on_u(v1 + 0.5 * kv2), h1 + 0.5 * kh2) * self.dt
kv3 = self.rhs_v(self.u_on_v(u1 + 0.5 * ku2), h1 + 0.5 * kh2) * self.dt
#######################
# Time propagation #
#######################
# Update
kh1_ad = 1 / 6 * adh0
kh2_ad = 1 / 3 * adh0
kh3_ad = 1 / 3 * adh0
kh4_ad = 1 / 6 * adh0
#adh0 = 0
kv1_ad = 1 / 6 * adv0
kv2_ad = 1 / 3 * adv0
kv3_ad = 1 / 3 * adv0
kv4_ad = 1 / 6 * adv0
#adv0 = 0
ku1_ad = 1 / 6 * adu0
ku2_ad = 1 / 3 * adu0
ku3_ad = 1 / 3 * adu0
ku4_ad = 1 / 6 * adu0
#adu0 = 0
#######################
# Right hand sides #
#######################
# kh4_ad
kh4_ad, adu_incr, adv_incr, adHe, ku3_ad, kv3_ad = self.kh_adj(
t, kh4_ad, adu, adv, adHe, u1 + ku3, v1 + kv3, He, ku3_ad, kv3_ad)
# kv4_ad
kv4_ad, adu, adh, ku3_ad, kh3_ad = self.kv_adj(kv4_ad, adu, adh,
ku3_ad, kh3_ad)
# ku4_ad
ku4_ad, adv, adh, kv3_ad, kh3_ad = self.ku_adj(ku4_ad, adv, adh,
kv3_ad, kh3_ad)
# kh3_ad
kh3_ad, adu, adv, adHe, ku2_ad, kv2_ad = self.kh_adj(
t, kh3_ad, adu, adv, adHe, u1 + 0.5 * ku2, v1 + 0.5 * kv2, He,
ku2_ad, kv2_ad, 1 / 2)
# kv3_ad
kv3_ad, adu, adh, ku2_ad, kh2_ad = self.kv_adj(kv3_ad, adu, adh,
ku2_ad, kh2_ad, 1 / 2)
# ku3_ad
ku3_ad, adv, adh, kv2_ad, kh2_ad = self.ku_adj(ku3_ad, adv, adh,
kv2_ad, kh2_ad, 1 / 2)
# kh2_ad
kh2_ad, adu, adv, adHe, ku1_ad, kv1_ad = self.kh_adj(
t, kh2_ad, adu, adv, adHe, u1 + 0.5 * ku1, v1 + 0.5 * kv1, He,
ku1_ad, kv1_ad, 1 / 2)
# kv2_ad
kv2_ad, adu, adh, ku1_ad, kh1_ad = self.kv_adj(kv2_ad, adu, adh,
ku1_ad, kh1_ad, 1 / 2)
# ku2_ad
ku2_ad, adv, adh, kv1_ad, kh1_ad = self.ku_adj(ku2_ad, adv, adh,
kv1_ad, kh1_ad, 1 / 2)
# kh1_ad
kh1_ad, adu, adv, adHe = self.kh_adj(t, kh1_ad, adu, adv, adHe, u1, v1,
He)
# kv1_ad
kv1_ad, adu, adh = self.kv_adj(kv1_ad, adu, adh, None, None)
# ku1_ad
ku1_ad, adv, adh = self.ku_adj(ku1_ad, adv, adh, None, None)
#######################
# Update #
#######################
adu += adu0
adv += adv0
adh += adh0
#######################
# Init bc #
#######################
if first:
adHe_tmp, adhbcx_tmp, adhbcy_tmp = init_bc_adj(self,
adu,
adv,
adh,
He,
hbcx,
hbcy,
t0=t)
adHe += adHe_tmp
adhbcx += adhbcx_tmp
adhbcy += adhbcy_tmp
return adu, adv, adh, adHe, adhbcx, adhbcy
def step_lf_adj(self,
t,
adu0,
adv0,
adh0,
u0,
v0,
h0,
He=None,
hbcx=None,
hbcy=None,
first=False):
########################
# Init #
########################
u1 = deepcopy(u0)
v1 = deepcopy(v0)
h1 = deepcopy(h0)
adu = 0
adv = 0
adh = 0
adHe = 0
adhbcx = 0
adhbcy = 0
if first:
init_bc(self, u1, v1, h1, He, hbcx, hbcy, t0=t)
#######################
# Boundary conditions #
#######################
adu_tmp,adv_tmp,adh_tmp,adHe_tmp,adhbcx_tmp,adhbcy_tmp = \
obcs_adj(self,t,adu0,adv0,adh0,u1,v1,h1,He,hbcx,hbcy)
adu += adu_tmp
adv += adv_tmp
adh += adh_tmp
adHe += adHe_tmp
adhbcx += adhbcx_tmp
adhbcy += adhbcy_tmp
#######################
# Time propagation #
#######################
if first:
# forward euler
adku = self.dt * adu0
adkv = self.dt * adv0
adkh = self.dt * adh0
else:
# leap-frog
adku = 2 * self.dt * adu0
adkv = 2 * self.dt * adv0
adkh = 2 * self.dt * adh0
#######################
# Right hand sides #
#######################
adu_tmp, adh_tmp = self.rhs_v_adj(adkv)
adu_tmp = self.u_on_v_adj(adu_tmp)
adu += adu_tmp
adh += adh_tmp
adv_tmp, adh_tmp = self.rhs_u_adj(adku)
adv_tmp = self.v_on_u_adj(adv_tmp)
adv += adv_tmp
adh += adh_tmp
adu_tmp, adv_tmp, adHe_tmp = self.rhs_h_adj(t, adkh, u1, v1, He)
adu += adu_tmp
adv += adv_tmp
adHe += adHe_tmp
#######################
# Update #
#######################
adu += self._adpu
adv += self._adpv
adh += self._adph
#######################
# Init bc #
#######################
if first:
# Init
adHe_tmp, adhbcx_tmp, adhbcy_tmp = init_bc_adj(self,
adu,
adv,
adh,
He,
hbcx,
hbcy,
t0=t)
adHe += adHe_tmp
adhbcx += adhbcx_tmp
adhbcy += adhbcy_tmp
if first:
adu += adu0
adv += adv0
adh += adh0
self._adpu = adu0
self._adpv = adv0
self._adph = adh0
return adu, adv, adh, adHe, adhbcx, adhbcy
def step_euler_tgl(self,
t,
du0,
dv0,
dh0,
u0,
v0,
h0,
dHe=None,
He=None,
dhbcx=None,
dhbcy=None,
hbcx=None,
hbcy=None,
first=False):
#######################
# Init local state #
#######################
u1 = deepcopy(u0)
v1 = deepcopy(v0)
h1 = deepcopy(h0)
du1 = deepcopy(du0)
dv1 = deepcopy(dv0)
dh1 = deepcopy(dh0)
#######################
# Init bc #
#######################
if first:
init_bc_tgl(self,
du1,
dv1,
dh1,
dHe,
dhbcx,
dhbcy,
He,
hbcx,
hbcy,
t0=t)
init_bc(self, u1, v1, h1, He, hbcx, hbcy, t0=t)
#######################
# Boundary conditions #
#######################
duS,dvS,dhS,duN,dvN,dhN,duW,dvW,dhW,duE,dvE,dhE = \
obcs_tgl(self,t,du1,dv1,dh1,dHe,dhbcx,dhbcy,
u1,v1,h1,He,hbcx,hbcy)
#######################
# Right hand sides #
#######################
dku = self.rhs_u(self.v_on_u(dv1), dh1)
dkv = self.rhs_v(self.u_on_v(du1), dh1)
dkh = self.rhs_h_tgl(du1, dv1, dHe, u1, v1, He)
#######################
# Time propagation #
#######################
du = du1 + self.dt * dku
dv = dv1 + self.dt * dkv
dh = dh1 + self.dt * dkh
########################
# Update border pixels #
########################
update_borders_tgl(self, du, dv, dh, duS, dvS, dhS, duN, dvN, dhN, duW,
dvW, dhW, duE, dvE, dhE)
return du, dv, dh
def step_rk4_tgl(self,
t,
du0,
dv0,
dh0,
u0,
v0,
h0,
dHe=None,
He=None,
dhbcx=None,
dhbcy=None,
hbcx=None,
hbcy=None,
first=False):
#######################
# Init local state #
#######################
u1 = deepcopy(u0)
v1 = deepcopy(v0)
h1 = deepcopy(h0)
du1 = deepcopy(du0)
dv1 = deepcopy(dv0)
dh1 = deepcopy(dh0)
#######################
# Init bc #
#######################
if first:
init_bc_tgl(self,
du1,
dv1,
dh1,
dHe,
dhbcx,
dhbcy,
He,
hbcx,
hbcy,
t0=t)
init_bc(self, u1, v1, h1, He, hbcx, hbcy, t0=t)
#######################
# Boundary conditions #
#######################
duS,dvS,dhS,duN,dvN,dhN,duW,dvW,dhW,duE,dvE,dhE = \
obcs_tgl(self,t,du1,dv1,dh1,dHe,dhbcx,dhbcy,
u1,v1,h1,He,hbcx,hbcy)
#######################
# Current trajectory #
#######################
# k1
ku1 = self.rhs_u(self.v_on_u(v1), h1) * self.dt
kv1 = self.rhs_v(self.u_on_v(u1), h1) * self.dt
kh1 = self.rhs_h(u1, v1, He) * self.dt
# k2
ku2 = self.rhs_u(self.v_on_u(v1 + 0.5 * kv1), h1 + 0.5 * kh1) * self.dt
kv2 = self.rhs_v(self.u_on_v(u1 + 0.5 * ku1), h1 + 0.5 * kh1) * self.dt
kh2 = self.rhs_h(u1 + 0.5 * ku1, v1 + 0.5 * kv1, He) * self.dt
# k3
ku3 = self.rhs_u(self.v_on_u(v1 + 0.5 * kv2), h1 + 0.5 * kh2) * self.dt
kv3 = self.rhs_v(self.u_on_v(u1 + 0.5 * ku2), h1 + 0.5 * kh2) * self.dt
#######################
# Right hand sides #
#######################
# k1_p
ku1_p = self.rhs_u(self.v_on_u(dv1), dh1) * self.dt
kv1_p = self.rhs_v(self.u_on_v(du1), dh1) * self.dt
kh1_p = self.rhs_h_tgl(du1, dv1, dHe, u1, v1, He) * self.dt
# k2_p
ku2_p = self.rhs_u(self.v_on_u(dv1 + 0.5 * kv1_p),
dh1 + 0.5 * kh1_p) * self.dt
kv2_p = self.rhs_v(self.u_on_v(du1 + 0.5 * ku1_p),
dh1 + 0.5 * kh1_p) * self.dt
kh2_p = self.rhs_h_tgl(du1 + 0.5 * ku1_p, dv1 + 0.5 * kv1_p, dHe,
u1 + 0.5 * ku1, v1 + 0.5 * kv1, He) * self.dt
# k3_p
ku3_p = self.rhs_u(self.v_on_u(dv1 + 0.5 * kv2_p),
dh1 + 0.5 * kh2_p) * self.dt
kv3_p = self.rhs_v(self.u_on_v(du1 + 0.5 * ku2_p),
dh1 + 0.5 * kh2_p) * self.dt
kh3_p = self.rhs_h_tgl(du1 + 0.5 * ku2_p, dv1 + 0.5 * kv2_p, dHe,
u1 + 0.5 * ku2, v1 + 0.5 * kv2, He) * self.dt
# k4_p
ku4_p = self.rhs_u(self.v_on_u(dv1 + kv3_p), dh1 + kh3_p) * self.dt
kv4_p = self.rhs_v(self.u_on_v(du1 + ku3_p), dh1 + kh3_p) * self.dt
kh4_p = self.rhs_h_tgl(du1 + ku3_p, dv1 + kv3_p, dHe, u1 + ku3,
v1 + kv3, He) * self.dt
#######################
# Time propagation #
#######################
du = du1 + 1 / 6 * (ku1_p + 2 * ku2_p + 2 * ku3_p + ku4_p)
dv = dv1 + 1 / 6 * (kv1_p + 2 * kv2_p + 2 * kv3_p + kv4_p)
dh = dh1 + 1 / 6 * (kh1_p + 2 * kh2_p + 2 * kh3_p + kh4_p)
########################
# Update border pixels #
########################
update_borders_tgl(self, du, dv, dh, duS, dvS, dhS, duN, dvN, dhN, duW,
dvW, dhW, duE, dvE, dhE)
return du, dv, dh