========================================
Optimization on Half Edge datastructures
========================================

The :py:mod:`~ddg.optimize.he` sub-package of :py:mod:`~ddg.optimize` contains utilities 
to create functionals on our half edge datastructures and
the means for their minimization (or maximization).
Further it contains some examples that can be found in the :py:mod:`~ddg.optimize.he.functionals` package.
The values of the functional can be located on any single cell type of a given surface,
and a subset of boundary cells with fixed values can be chosen.
For minimization the `scipy.optimize.minimize <https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html>`_ function is used.

.. contents:: Table of contents
   :local:
   :backlinks: none


The HalfEdgeFunctional
----------------------

The :py:class:`~ddg.optimize.he.functional.HalfEdgeFunctional` is an abstract base class for functionals
on half-edge surfaces.
Its init supports the following arguments:

    - surface: The half-edge surface
    - attr_name: The name of the attribute being used for minimization
    - cell_type: Either "verts", "edges" or "faces" depending where the minimization attribute is located
    - boundary_cells (optional, default=set()): Boundary cells for minimization, i.e. cells that will not change their value

Another essential part of the HalfEdgeFunctional is the :py:meth:`~ddg.optimize.he.functional.HalfEdgeFunctional.evaluate`
method that has to be implemented by any inheriting functional.
This method represents the evaluation function of the functional. It has to be of the form:

    ``evaluate(self, x, **kwargs)``

where x is an ndarray of shape (n,), storing the values of interest for the minimization.

In order to show how to work with the :py:class:`~ddg.optimize.he.functional.HalfEdgeFunctional` we will compute and minimize the combinatorial Dirichlet energy of a function :math:`x : V \rightarrow \mathbb{R}` on the vertices :math:`V` of a discrete surface:

.. math::
    D(x) = \sum_{(ij) \in E}^{|E|} (x_i - x_j)^2

where :math:`E` are the edges of the surface and :math:`x_k` the values of :math:`x` located at the vertices.
We use the ``__init__`` function of the ``HalfEdgeFunctional`` and define its ``evaluate`` function.

.. doctest::

    >>> from ddg.optimize.he.functional import HalfEdgeFunctional
    >>> from numpy import linalg as LA
    >>> import numpy as np

    >>> class DirichletEnergy(HalfEdgeFunctional):
    ...     def __init__(
    ...         self, surface, attr_name, attr_dimension=None, boundary_cells=set()
    ...     ):
    ...         super().__init__(
    ...             surface, attr_name, "verts", boundary_cells=boundary_cells
    ...         )
    ...
    ...     def evaluate(self, x):
    ...         E = 0
    ...         half_edges = ddg.halfedge.single_edges(self.surface)
    ...         for e in half_edges:
    ...             i = e.tail
    ...             j = e.head
    ...             x_i = (
    ...                 x[i.interior_cell_index]
    ...                 if i.interior_cell_index is not None
    ...                 else getattr(i, self.attr_name)
    ...             )
    ...             x_j = (
    ...                 x[j.interior_cell_index]
    ...                 if j.interior_cell_index is not None
    ...                 else getattr(j, self.attr_name)
    ...             )
    ...             E += LA.norm(x_i - x_j) ** 2
    ...         return E
    ...

The minimizer function
----------------------
The minimizer function :py:func:`~ddg.optimize.he.minimize.minimizer` minimizes a given functional using
`scipy.optimize.minimize <https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html>`_.
We initialize a half-edge surface, assign zero values to the function on the boundary, and random values to all other vertices.
Then we create an instance of the ``DirichletEnergy`` functional using this surfaces and its boundary and initial data,
and minimize the Dirichlet energy of the function.
The result of the minimization will be a function with constant zero values.
Its total Dirichlet energy is zero.

.. doctest::

    >>> import ddg
    >>> from ddg.optimize.he.minimize import minimizer
    >>> from ddg.halfedge import (
    ...     is_boundary_vertex,
    ...     boundary_vertices,
    ...     some_vertex,
    ... )
    >>> import numpy as np


    >>> s = ddg.halfedge.grid((6, 5, 1))
    >>> attr = s.verts.add_attribute("x")

    >>> for v in s.verts:
    ...     if is_boundary_vertex(v):
    ...         v.x = 0
    ...     else:
    ...         v.x = np.random.random()
    ...

    >>> DE = DirichletEnergy(s, "x", boundary_cells=list(boundary_vertices(s)))

    >>> result = minimizer(DE)

    >>> assert np.allclose([v.x for v in s.verts], 0, atol=1.0e-5)


Note that  :py:func:`~ddg.halfedge.boundary_vertices` returns a generator object and thus has to be converted to a list before handing it over to the DirichletEnergy.
Instead of using the boundary of the surface you can give an arbitrary subset of the vertices as boundary vertices.

Using

.. doctest::

    >>> result, history = minimizer(DE, return_history=True)

the history can be accessed after the minimization process. It contains entries for each step where the entries contain
dictionaries that map indices of interior cells to the values:

.. doctest::

    >>> for i in range(len(history)):
    ...     surface.verts.add_attribute(f"minimization_step_{i}")
    ...     for v in DE.interior_cells:
    ...         (v, f"minimization_step_{i}", history[i][v.index])
    ...


Example functionals
-------------------

Example functionals from the :py:mod:`~ddg.optimize.he.functionals` package can currently be
used to calculate the :py:mod:`~ddg.optimize.he.functionals.dirichlet_energy` or the
:py:mod:`~ddg.optimize.he.functionals.planar_faces_energy` of a given half-edge datastructure.