Skip to content

# Solver

## `` `felice.solver`

### Classes

#### `` `ClipSolver`

Bases: `Module`

Source code in `felice/solver.py`

    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64

    class ClipSolver(eqx.Module):
        solver: AbstractSolver

        def __getattr__(self, name):
            return getattr(self.solver, name)

        def step(
            self,
            terms: PyTree[AbstractTerm],
            t0: RealScalarLike,
            t1: RealScalarLike,
            y0: Y,
            args: Args,
            solver_state: _SolverState,
            made_jump: BoolScalarLike,
        ) -> tuple[Y, Optional[Y], DenseInfo, _SolverState, RESULTS]:
            """Make a single step of the solver.

            Each step is made over the specified interval $[t_0, t_1]$.

            **Arguments:**

            - `terms`: The PyTree of terms representing the vector fields and controls.
            - `t0`: The start of the interval that the step is made over.
            - `t1`: The end of the interval that the step is made over.
            - `y0`: The current value of the solution at `t0`.
            - `args`: Any extra arguments passed to the vector field.
            - `solver_state`: Any evolving state for the solver itself, at `t0`.
            - `made_jump`: Whether there was a discontinuity in the vector field at `t0`.
                Some solvers (notably FSAL Runge--Kutta solvers) usually assume that there
                are no jumps and for efficiency re-use information between steps; this
                indicates that a jump has just occurred and this assumption is not true.

            **Returns:**

            A tuple of several objects:

            - The value of the solution at `t1`.
            - A local error estimate made during the step. (Used by adaptive step size
                controllers to change the step size.) May be `None` if no estimate was
                made.
            - Some dictionary of information that is passed to the solver's interpolation
                routine to calculate dense output. (Used with `SaveAt(ts=...)` or
                `SaveAt(dense=...)`.)
            - The value of the solver state at `t1`.
            - An integer (corresponding to `diffrax.RESULTS`) indicating whether the step
                happened successfully, or if (unusually) it failed for some reason.
            """
            y1, y_error, dense_info, solver_state, result = self.solver.step(
                terms, t0, t1, y0, args, solver_state, made_jump
            )
            y1_clipped = jax.tree_util.tree_map(jax.nn.relu, y1)
            return y1_clipped, y_error, dense_info, solver_state, result

##### Functions

###### `` `step(terms: PyTree[AbstractTerm], t0: RealScalarLike, t1: RealScalarLike, y0: Y, args: Args, solver_state: _SolverState, made_jump: BoolScalarLike) -> tuple[Y, Optional[Y], DenseInfo, _SolverState, RESULTS]`

Make a single step of the solver.

Each step is made over the specified interval \([t_0, t_1]\).

Arguments:

- `terms`: The PyTree of terms representing the vector fields and controls.
- `t0`: The start of the interval that the step is made over.
- `t1`: The end of the interval that the step is made over.
- `y0`: The current value of the solution at `t0`.
- `args`: Any extra arguments passed to the vector field.
- `solver_state`: Any evolving state for the solver itself, at `t0`.
- `made_jump`: Whether there was a discontinuity in the vector field at `t0`. Some solvers (notably FSAL Runge--Kutta solvers) usually assume that there are no jumps and for efficiency re-use information between steps; this indicates that a jump has just occurred and this assumption is not true.

Returns:

A tuple of several objects:

- The value of the solution at `t1`.
- A local error estimate made during the step. (Used by adaptive step size controllers to change the step size.) May be `None` if no estimate was made.
- Some dictionary of information that is passed to the solver's interpolation routine to calculate dense output. (Used with `SaveAt(ts=...)` or `SaveAt(dense=...)`.)
- The value of the solver state at `t1`.
- An integer (corresponding to `diffrax.RESULTS`) indicating whether the step happened successfully, or if (unusually) it failed for some reason.

Source code in `felice/solver.py`

    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64

    def step(
        self,
        terms: PyTree[AbstractTerm],
        t0: RealScalarLike,
        t1: RealScalarLike,
        y0: Y,
        args: Args,
        solver_state: _SolverState,
        made_jump: BoolScalarLike,
    ) -> tuple[Y, Optional[Y], DenseInfo, _SolverState, RESULTS]:
        """Make a single step of the solver.

        Each step is made over the specified interval $[t_0, t_1]$.

        **Arguments:**

        - `terms`: The PyTree of terms representing the vector fields and controls.
        - `t0`: The start of the interval that the step is made over.
        - `t1`: The end of the interval that the step is made over.
        - `y0`: The current value of the solution at `t0`.
        - `args`: Any extra arguments passed to the vector field.
        - `solver_state`: Any evolving state for the solver itself, at `t0`.
        - `made_jump`: Whether there was a discontinuity in the vector field at `t0`.
            Some solvers (notably FSAL Runge--Kutta solvers) usually assume that there
            are no jumps and for efficiency re-use information between steps; this
            indicates that a jump has just occurred and this assumption is not true.

        **Returns:**

        A tuple of several objects:

        - The value of the solution at `t1`.
        - A local error estimate made during the step. (Used by adaptive step size
            controllers to change the step size.) May be `None` if no estimate was
            made.
        - Some dictionary of information that is passed to the solver's interpolation
            routine to calculate dense output. (Used with `SaveAt(ts=...)` or
            `SaveAt(dense=...)`.)
        - The value of the solver state at `t1`.
        - An integer (corresponding to `diffrax.RESULTS`) indicating whether the step
            happened successfully, or if (unusually) it failed for some reason.
        """
        y1, y_error, dense_info, solver_state, result = self.solver.step(
            terms, t0, t1, y0, args, solver_state, made_jump
        )
        y1_clipped = jax.tree_util.tree_map(jax.nn.relu, y1)
        return y1_clipped, y_error, dense_info, solver_state, result