---
jupytext:
  main_language: python
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.16.2
---

```{eval-rst}
.. include:: ../../../global.rst
```

(demo_pre_post_hooks)=
# Pre- and Post-Callbacks for the optimization

## Problem Formulation

In this demo we show how one can use the flexibility of cashocs
to "inject" their own code into the optimization, which is carried
out before each solve of the state system (`pre_callback`) or after each
gradient computation (`post_callback`). To do so, we investigate the following
problem

$$
\begin{align}
    &\min\; J(u, c) = \frac{1}{2} \int_\Omega \left\lvert u - u_d \right\rvert^2
    \text{ d}x +
    \frac{\alpha}{2} \int_\Omega \left\lvert c \right\rvert^2 \text{ d}x \\
    &\text{ subject to } \qquad
    \begin{alignedat}[t]{2}
        -\Delta u + Re (u\cdot \nabla) u + \nabla p &= c \quad &&\text{ in } \Omega, \\
        \text{div}(u) &= 0 \quad &&\text{ in } \Omega,\\
        u &= u_\text{dir} \quad &&\text{ on } \Gamma^\text{dir},\\
        u &= 0 \quad &&\text{ on } \Gamma^\text{no slip},\\
        p &= 0 \quad &&\text{ at } x^\text{pres}.
    \end{alignedat}
\end{align}
$$

In particular, the setting for this demo is very similar to the one of
{ref}`demo_stokes`, but here we consider the nonlinear incompressible Navier-Stokes
equations.

In the following, we will describe how to solve this problem
using cashocs, where we will use the {python}`pre_callback` functionality to implement
a homotopy method for solving the Navier-Stokes equations.

## Implementation

The complete python code can be found in the file {download}`demo_constraints.py
</../../demos/documented/optimal_control/pre_post_callbacks/demo_pre_post_callbacks.py>`
and the corresponding config can be found in {download}`config.ini
</../../demos/documented/optimal_control/pre_post_callbacks/config.ini>`.

### Initialization

The initial part of the code is nearly identical to the one of {ref}`demo_stokes`, but
here we use the Navier-Stokes equations instead of the linear Stokes system

```python
from fenics import *

import cashocs

config = cashocs.load_config("./config.ini")
mesh, subdomains, boundaries, dx, ds, dS = cashocs.regular_mesh(30)

v_elem = VectorElement("CG", mesh.ufl_cell(), 2)
p_elem = FiniteElement("CG", mesh.ufl_cell(), 1)
V = FunctionSpace(mesh, MixedElement([v_elem, p_elem]))
U = VectorFunctionSpace(mesh, "CG", 1)

up = Function(V)
u, p = split(up)
vq = Function(V)
v, q = split(vq)
c = Function(U)

Re = Constant(1e2)

e = (
    inner(grad(u), grad(v)) * dx
    + Constant(Re) * dot(grad(u) * u, v) * dx
    - p * div(v) * dx
    - q * div(u) * dx
    - inner(c, v) * dx
)


def pressure_point(x, on_boundary):
    return near(x[0], 0) and near(x[1], 0)


no_slip_bcs = cashocs.create_dirichlet_bcs(
    V.sub(0), Constant((0, 0)), boundaries, [1, 2, 3]
)
lid_velocity = Expression(("4*x[0]*(1-x[0])", "0.0"), degree=2)
bc_lid = DirichletBC(V.sub(0), lid_velocity, boundaries, 4)
bc_pressure = DirichletBC(V.sub(1), Constant(0), pressure_point, method="pointwise")
bcs = no_slip_bcs + [bc_lid, bc_pressure]

alpha = 1e-5
u_d = Expression(
    (
        "sqrt(pow(x[0], 2) + pow(x[1], 2))*cos(2*pi*x[1])",
        "-sqrt(pow(x[0], 2) + pow(x[1], 2))*sin(2*pi*x[0])",
    ),
    degree=2,
)
J = cashocs.IntegralFunctional(
    Constant(0.5) * inner(u - u_d, u - u_d) * dx
    + Constant(0.5 * alpha) * inner(c, c) * dx
)
```

### Callbacks

Note, that we have chosen a Reynolds number of {python}`Re = 1e2` for this demo.
In order to solve the Navier-Stokes equations for higher Reynolds numbers, it is
often sensible to first solve the equations for a lower Reynolds number and then
use this solution as initial guess for the original high Reynolds number problem.
We can use this procedure in cashocs with its {python}`pre_callback` functionality.
A {python}`pre_callback` is a function without arguments, which gets called each time
before solving the state equation. In our case, the {python}`pre_callback` should
solve the Navier-Stokes equations for a lower Reynolds number, so we define it as
follows

```python
def pre_callback():
    print("Solving low Re Navier-Stokes equations for homotopy.")
    v, q = TestFunctions(V)
    e = (
        inner(grad(u), grad(v)) * dx
        + Constant(Re / 10.0) * dot(grad(u) * u, v) * dx
        - p * div(v) * dx
        - q * div(u) * dx
        - inner(c, v) * dx
    )
    cashocs.newton_solve(e, up, bcs, verbose=False)
```

where we solve the Navier-Stokes equations with a lower Reynolds number of
{python}`Re / 10.0`. Later on, we use this function as keyword argument for defining
the optimization problem.

Additionally, cashocs implements the functionality of also performing a pre-defined
action after each gradient computation, given by a so-called {python}`post_callback`.
In our case, we just want to print a statement so that we can visualize what is
happening. Therefore, we define our {python}`post_callback` as

```python
def post_callback():
    print("Performing an action after computing the gradient.")
```

Next, we define the optimization and use the keyword arguments to define the callbacks
via

```python
ocp = cashocs.OptimalControlProblem(
    e,
    bcs,
    J,
    up,
    c,
    vq,
    config=config,
    pre_callback=pre_callback,
    post_callback=post_callback,
)
```

::::{note}
Alternatively, the pre- and post-callbacks can be injected to an already defined
optimization problem with the code

:::{code-block} python
ocp.inject_pre_callback(pre_callback)
ocp.inject_post_callback(post_callback)
:::

or, equivalently,

:::{code-block} python
ocp.inject_pre_post_hook(pre_hook, post_hook)
:::

::::

:::{note}
The callback functions are allowed to have one argument. In case an argument is
supplied, the callback function is called with the optimization problem itself
as an argument during runtime. This allows in-depth manipulation of the optimization
algorithms and optimization problems during runtime. This feature should only be used
with care!
:::

+++

In the end, we solve the problem with

```python
ocp.solve()
```

We visualize the results with the lines

```python
u, p = up.split(True)
import matplotlib.pyplot as plt

plt.figure(figsize=(15, 5))

plt.subplot(1, 3, 1)
fig = plot(c)
plt.colorbar(fig, fraction=0.046, pad=0.04)
plt.title("Control variable c")

plt.subplot(1, 3, 2)
fig = plot(u)
plt.colorbar(fig, fraction=0.046, pad=0.04)
plt.title("State variable u")

plt.subplot(1, 3, 3)
fig = plot(u_d, mesh=mesh)
plt.colorbar(fig, fraction=0.046, pad=0.04)
plt.title("Desired state u_d")

plt.tight_layout()
# plt.savefig("./img_pre_post_callbacks.png", dpi=150, bbox_inches="tight")
```

and the results are given below
![](/../../demos/documented/optimal_control/pre_post_callbacks/img_pre_post_callbacks.png)