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)
示例#2
0
 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)