def step_fn(time,
next_time,
coord_grid,
value_grid,
boundary_conditions,
second_order_coeff_fn,
first_order_coeff_fn,
zeroth_order_coeff_fn,
inner_second_order_coeff_fn,
inner_first_order_coeff_fn,
num_steps_performed,
dtype=None,
name=None):
"""Performs the step."""
name = name or 'composite_scheme_step'
def scheme(*args, **kwargs):
return tf.cond(num_steps_performed < first_scheme_steps,
lambda: first_scheme(*args, **kwargs),
lambda: second_scheme(*args, **kwargs))
return parabolic_equation_step(time,
next_time,
coord_grid,
value_grid,
boundary_conditions,
second_order_coeff_fn,
first_order_coeff_fn,
zeroth_order_coeff_fn,
inner_second_order_coeff_fn,
inner_first_order_coeff_fn,
time_marching_scheme=scheme,
dtype=dtype,
name=name)
def step_fn(time,
next_time,
coord_grid,
value_grid,
boundary_conditions,
second_order_coeff_fn,
first_order_coeff_fn,
zeroth_order_coeff_fn,
inner_second_order_coeff_fn,
inner_first_order_coeff_fn,
num_steps_performed,
dtype=None,
name=None):
"""Performs the step."""
del num_steps_performed
name = name or 'weighted_implicit_explicit_scheme'
return parabolic_equation_step(time,
next_time,
coord_grid,
value_grid,
boundary_conditions,
second_order_coeff_fn,
first_order_coeff_fn,
zeroth_order_coeff_fn,
inner_second_order_coeff_fn,
inner_first_order_coeff_fn,
time_marching_scheme=scheme,
dtype=dtype,
name=name)