# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
# ---

# ```{eval-rst}
# .. include:: ../../../global.rst
# ```
#
# (demo_shape_stokes)=
# # Optimization of an Obstacle in Stokes Flow
#
# ## Problem Formulation
#
# For this tutorial, we investigate a problem similarly to the ones considered in, e.g.,
# [Blauth - Nonlinear Conjugate Gradient Methods for PDE Constrained Shape Optimization
# Based on Steklov-Poincaré-Type Metrics](https://arxiv.org/abs/2007.12891) and
# [Schulz and Siebenborn - Computational Comparison of Surface Metric for PDE
# Constrained Shape Optimization](https://doi.org/10.1515/cmam-2016-0009), which reads:
#
# $$
# \begin{align}
#     &\min_\Omega J(u, \Omega) = \int_{\Omega^\text{flow}} Du : Du\ \text{ d}x \\
#     &\text{subject to } \qquad
#     \begin{alignedat}[t]{2}
#         - \Delta u + \nabla p &= 0 \quad &&\text{ in } \Omega, \\
#         \text{div}(u) &= 0 \quad &&\text{ in } \Omega, \\
#         u &= u^\text{in} \quad &&\text{ on } \Gamma^\text{in}, \\
#         u &= 0 \quad &&\text{ on } \Gamma^\text{wall} \cup \Gamma^\text{obs}, \\
#         \partial_n u - p n &= 0 \quad &&\text{ on } \Gamma^\text{out}.
#     \end{alignedat}
# \end{align}
# $$
#
# Here, we have an inflow boundary condition on the inlet $\Gamma^\text{in}$,
# a no-slip boundary condition on the wall boundary $\Gamma^\text{wall}$ as well
# as the obstacle boundary $\Gamma^\text{obs}$, and a do-nothing condition
# on the outlet $\Gamma^\text{out}$.
#
# This problem is supplemented with the geometrical constraints
#
# $$
# \begin{align}
#     \text{vol}(\Omega) = \int_\Omega 1 \text{ d}x &= \int_{\Omega_0} 1 \text{ d}x
#     = \text{vol}(\Omega_0), \\
#     \text{bary}(\Omega) = \frac{1}{\text{vol}(\Omega)} \int_\Omega x \text{ d}x
#     &= \frac{1}{\text{vol}(\Omega_0)} \int_{\Omega_0} x \text{ d}x
#     = \text{bary}(\Omega_0),
# \end{align}
# $$
#
# where $\Omega_0$ denotes the initial geometry. This models that both the volume
# and the barycenter of the geometry should remain fixed during the optimization.
# To treat this problem with cashocs, we regularize the constraints in the sense
# of a quadratic penalty function. See {ref}`the corresponding section in the
# documentation of the config files for shape optimization
# <config_shape_regularization>` as well as {ref}`demo_regularization`. In particular,
# the regularized problem reads
#
# $$
# \begin{align}
#     \min_\Omega J(u, \Omega) = &\int_{\Omega^\text{flow}} Du : Du\ \text{ d}x +
#     \frac{\mu_\text{vol}}{2} \left( \int_\Omega 1 \text{ d}x -
#     \text{vol}(\Omega_0) \right)^2 \\
#     &+ \frac{\mu_\text{bary}}{2} \left\lvert \frac{1}{\text{vol}(\Omega)}
#     \int_\Omega x \text{ d}x - \text{bary}(\Omega_0) \right\rvert^2 \\
#     &\text{subject to } \qquad
#     \begin{alignedat}[t]{2}
#         - \Delta u + \nabla p &= 0 \quad &&\text{ in } \Omega, \\
#         \text{div}(u) &= 0 \quad &&\text{ in } \Omega, \\
#         u &= u^\text{in} \quad &&\text{ on } \Gamma^\text{in}, \\
#         u &= 0 \quad &&\text{ on } \Gamma^\text{wall} \cup \Gamma^\text{obs}, \\
#         \partial_n u - p n &= 0 \quad &&\text{ on } \Gamma^\text{out}.
#     \end{alignedat}
# \end{align}.
# $$
#
# where we have no additional geometrical constraints. However, the parameters
# $\mu_\text{vol}$ and $\mu_\text{bary}$ have to be sufficiently large,
# so that the regularized constraints are satisfied approximately.
#
# Note, that the domain $\Omega$ denotes the domain of the obstacle, but that
# the cost functional involves an integral over the flow domain $\Omega^\text{flow}$.
# Therefore, we switch our point of view and optimize the geometry of the flow instead
# (which of course includes the optimization of the "hole" generated by the obstacle).
# In particular, only the boundary of the obstacle $\Gamma^\text{obs}$ is deformable,
# all remaining boundaries are fixed. Hence, it is equivalent to pose the geometrical
# constraints on $\Omega$ and $\Omega^\text{flow}$, so that we can work completely with
# the flow domain.
#
# As (initial) domain, we consider the rectangle $(-2, 2) \times (-3, 3)$
# without a circle centered in the origin with radius $0.5$. The correspondig
# geometry is created with Gmsh, and can be found in the `./mesh/` directory.
#
# ## Implementation
#
# The complete python code can be found in the file {download}`demo_shape_stokes.py
# </../../demos/documented/shape_optimization/shape_stokes/demo_shape_stokes.py>`,
# and the corresponding config can be found in {download}`config.ini
# </../../demos/documented/shape_optimization/shape_stokes/config.ini>`.
#
# ### Initialization
#
# As for the previous tutorial problems, we start by importing FEniCS and cashocs and
# load the configuration file

# +
from fenics import *

import cashocs

config = cashocs.load_config("./config.ini")
# -

# Afterwards we import the (initial) geometry with the {py:func}`import_mesh
# <cashocs.import_mesh>` command

mesh, subdomains, boundaries, dx, ds, dS = cashocs.import_mesh("./mesh/mesh.xdmf")

# ::::{note}
# As defined in `./mesh/mesh.geo`, we have the following boundaries
#
# - The inlet boundary $\Gamma^\text{in}$, which has index {python}`1`
# - The wall boundary $\Gamma^\text{wall}$, which have index {python}`2`
# - The outlet boundary $\Gamma^\text{out}$, which has index {python}`3`
# - The boundary of the obstacle $\Gamma^\text{obs}$, which has index {python}`4`
#
# To include the fact that only the obstacle boundary, with index {python}`4`, is
# deformable, we have the following lines in the config file
#
# ```{code-block} ini
# :caption: config.ini
# [ShapeGradient]
# shape_bdry_def = [4]
# shape_bdry_fix = [1,2,3]
# ```
#
# which ensures that only $\Gamma^\text{obs}$ is deformable.
# ::::
#
# ### State system
#
# The definition of the state system is analogous to the one we considered in
# {ref}`demo_stokes`. Here, we, too, use LBB stable Taylor-Hood elements, which are
# defined as

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

# For the weak form of the PDE, we have the same code as in {ref}`demo_stokes`

# +
up = Function(V)
u, p = split(up)
vq = Function(V)
v, q = split(vq)

e = inner(grad(u), grad(v)) * dx - p * div(v) * dx - q * div(u) * dx
# -

# The Dirichlet boundary conditions are slightly different, though. For the inlet
# velocity $u^\text{in}$ we use a parabolic profile

u_in = Expression(("-1.0/4.0*(x[1] - 2.0)*(x[1] + 2.0)", "0.0"), degree=2)
bc_in = DirichletBC(V.sub(0), u_in, boundaries, 1)

# The wall and obstacle boundaries get a no-slip boundary condition, each, with the line

bc_no_slip = cashocs.create_dirichlet_bcs(
    V.sub(0), Constant((0, 0)), boundaries, [2, 4]
)

# Finally, all Dirichlet boundary conditions are gathered into the list {python}`bcs`

bcs = [bc_in] + bc_no_slip

# :::{note}
# The outflow boundary condition is of Neumann type and already included in the
# weak form of the problem.
# :::
#
# ### Cost functional and optimization problem
#
# The cost functional is easily defined with the line

J = cashocs.IntegralFunctional(inner(grad(u), grad(u)) * dx)

# ::::{note}
# The additional regularization terms are defined similarly to
# {ref}`demo_regularization`.
# In the config file, we have the following (relevant) lines:
#
# ```{code-block} ini
# :caption: config.ini
# [Regularization]
# factor_volume = 1e4
# use_initial_volume = True
# factor_barycenter = 1e5
# use_initial_barycenter = True
# ```
#
# This ensures that we use $\mu_\text{vol} = 1e4$ and $\mu_\text{bary} = 1e5$,
# which are comparatively large parameters, so that the geometrical constraints are
# satisfied with high accuracy.
# Moreover, the boolean flags {ini}`use_initial_volume` and
# {ini}`use_initial_barycenter`, which are both set to `True`, ensure that we actually
# use the volume and barycenter of the initial geometry.
# ::::
#
# Finally, we solve the shape optimization problem as previously with the commands

sop = cashocs.ShapeOptimizationProblem(e, bcs, J, up, vq, boundaries, config=config)
sop.solve()

# :::{note}
# For the definition of the shape gradient, we use the same parameters for the
# linear elasticity equations as in [Blauth - Nonlinear Conjugate Gradient Methods for
# PDE Constrained Shape Optimization Based on Steklov-Poincaré-Type Metrics](
# https://arxiv.org/abs/2007.12891) and [Schulz and Siebenborn - Computational
# Comparison of Surface Metric for PDE Constrained Shape Optimization](
# https://doi.org/10.1515/cmam-2016-0009). These are defined in the config file
# in the {ref}`ShapeGradient section <config_shape_shape_gradient>`
#
# ```{code-block} ini
# :caption: config.ini
# [ShapeGradient]
# lambda_lame = 0.0
# damping_factor = 0.0
# mu_fix = 1
# mu_def = 5e2
# ```
#
# so that we use a rather high stiffness for the elements at the deformable boundary,
# which is also discretized finely, and a rather low stiffness for the fixed
# boundaries.
# :::
#
# :::{note}
# This demo might take a little longer than the others, depending on the machine
# it is run on. This is normal, and caused by the finer discretization of the geometry
# compared to the previous problems.
# :::
#
# We visualize the optimized geometry with the lines

# +
import matplotlib.pyplot as plt

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

ax_mesh = plt.subplot(1, 3, 1)
fig_mesh = plot(mesh)
plt.title("Discretization of the optimized geometry")

ax_u = plt.subplot(1, 3, 2)
ax_u.set_xlim(ax_mesh.get_xlim())
ax_u.set_ylim(ax_mesh.get_ylim())
fig_u = plot(u)
plt.colorbar(fig_u, fraction=0.046, pad=0.04)
plt.title("State variable u")


ax_p = plt.subplot(1, 3, 3)
ax_p.set_xlim(ax_mesh.get_xlim())
ax_p.set_ylim(ax_mesh.get_ylim())
fig_p = plot(p)
plt.colorbar(fig_p, fraction=0.046, pad=0.04)
plt.title("State variable p")

plt.tight_layout()
# plt.savefig('./img_shape_stokes.png', dpi=150, bbox_inches='tight')
# -

# and the result is shown below
# :::{image} /../../demos/documented/shape_optimization/shape_stokes/img_shape_stokes.png
# :::