{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0842cd55",
   "metadata": {},
   "source": [
    "# Optimize with Transformations\n",
    "\n",
    "Here we employ `upoqa` with transformations enabled to optimize a house-made problem. Consider a simple objective function with an L1 penalty term:\n",
    "\n",
    "$$\n",
    "f(x, y) = (x + y - 1)^2 + a |xy|\n",
    "$$\n",
    "\n",
    "The optimal solutions are $x = 0, y = 1$ or $x = 1, y = 0$, yielding an optimal function value of zero. Direct optimization proves challenging due to the L1 term. A viable approach involves applying a **smoothing technique**:\n",
    "\n",
    "$$\n",
    "\\widehat{f}(x, y) = (x + y - 1)^2 + a \\sqrt{x^2y^2 + \\tau}\n",
    "$$\n",
    "\n",
    "where $\\tau \\geq 0$. We initialize the optimization of $\\widehat{f}$ using a relatively large $\\tau$ value. The solution obtained from each optimization step serves as the starting point for the subsequent step. This process iterates until $\\tau$ becomes sufficiently small, at which point we obtain the solution to the original problem.\n",
    "\n",
    "For demonstration purposes, we optimize this problem using `upoqa` by treating both terms as element functions. The smoothed L1 term is interpreted as the composition of $(x,y) \\to x^2y^2$ and the transformation $z \\to a\\sqrt{z + \\tau}$. This functional composition can be implemented by configuring the `xform` parameter. The objective function thus becomes:\n",
    "\n",
    "$$\n",
    "\\widehat{f}(x, y) = f_1(x, y) + h_2(f_2(x, y)),\n",
    "$$\n",
    "\n",
    "where \n",
    "$$\n",
    "f_1(x, y) = (x + y - 1)^2, \\quad f_2(x, y) = x^2y^2, \\quad h_2(z) = a \\sqrt{z + \\tau}.\n",
    "$$\n",
    "\n",
    "The problem is then defined using the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40ba7e11",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "f(x0) = 246.00099999600002\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from upoqa.problems import PSProblem\n",
    "import numpy as np\n",
    "\n",
    "a = 5.0\n",
    "start_tau = 1e-2\n",
    "tau = start_tau\n",
    "\n",
    "def f1(x):\n",
    "    return (x[0] + x[1] - 1) ** 2\n",
    "\n",
    "def f2(x):\n",
    "    return (x[0] * x[1]) ** 2\n",
    "\n",
    "def h2(z, tau):\n",
    "    return a * np.sqrt(z + tau)\n",
    "\n",
    "def h2_grad(z, tau):\n",
    "    # :math:`\\nabla h_2(z) = 0.5 * (z + tau) ** (- 0.5)`\n",
    "    return 0.5 * a / np.sqrt(z + tau)\n",
    "\n",
    "def h2_hess(z, tau):\n",
    "    # :math:`\\nabla^2 h_2(z) = - 0.25 * (z + tau) ** (- 1.5)`\n",
    "    return - 0.25 * a/ np.pow(z + tau, 1.5)\n",
    "\n",
    "def obj(x, tau):\n",
    "    # The objective function with an L2 smoothing term (parameterized by tau)\n",
    "    return f1(x) + h2(f2(x), tau) \n",
    "\n",
    "# Start from a far point\n",
    "x0 = [-5, -5]\n",
    "\n",
    "# The minimum function value\n",
    "fopt = 0.0\n",
    "\n",
    "# Define the partially separable structure\n",
    "fun =    {'principal': f1,     'penalty': f2}\n",
    "coords = {'principal': [0, 1], 'penalty': [0, 1]}\n",
    "\n",
    "# The function, gradient and hession of the transformation\n",
    "xforms = {\n",
    "    'penalty': [lambda x: h2(x, tau), lambda x: h2_grad(x, tau), lambda x: h2_hess(x, tau)],\n",
    "}\n",
    "\n",
    "# Bound check for the input of transformation\n",
    "xform_bounds = {'penalty': (0, np.inf)}\n",
    "\n",
    "print(\"f(x0) =\", obj(x0, tau))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b0876cc",
   "metadata": {},
   "source": [
    "We then proceed to optimize the problem. \n",
    "\n",
    "A key technique employed here involves **passing a custom callback function** `callback_func` that UPOQA invokes at each iteration. This callback dynamically reduces the $\\tau$ value in a stepwise manner. When the reduction magnitude is appropriately controlled, this approach enables us to obtain the solution directly within a single optimization run, bypassing the need for multiple algorithm executions.\n",
    "\n",
    "The callback function accepts an `intermediate_result` parameter (of type `upoqa.utils.OptimizeResult`), which contains the current iteration's results. Through this object, we can access:\n",
    "- The current iterate `intermediate_result.x`\n",
    "- The current objective function value `intermediate_result.fun`\n",
    "- And some other things, see the docstring of `upoqa.utils.OptimizeResult` for details.\n",
    "\n",
    "To enable deeper control, we configure `upoqa` with `return_internals=True`. This grants access to the algorithm's internal state, particularly the algorithm manager (`intermediate_result.manager`). Then we can invoke `intermediate_result.manager.update_xforms()` to implement the dynamic $\\tau$ updates described above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "146defc9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Function value obtained by UPOQA: 4.522866417616116e-10\n",
      "Function value obtained by BFGS via direct optimization: 0.5555555555555798\n"
     ]
    }
   ],
   "source": [
    "def callback_func(intermediate_result):\n",
    "    # calculate the new tau\n",
    "    tau = start_tau * 0.9 ** intermediate_result.nit\n",
    "    # update the transformation with the new tau\n",
    "    intermediate_result.manager.update_xforms({\n",
    "        'penalty': [lambda x: h2(x, tau), lambda x: h2_grad(x, tau), lambda x: h2_hess(x, tau)],\n",
    "    })\n",
    "    \n",
    "# optimize the problem with upoqa\n",
    "import upoqa\n",
    "res = upoqa.minimize(\n",
    "    fun,\n",
    "    x0=x0,\n",
    "    coords=coords,\n",
    "    xforms=xforms,\n",
    "    xform_bounds=xform_bounds,\n",
    "    callback=callback_func,\n",
    "    disp=0,\n",
    "    return_internals=True\n",
    ")\n",
    "print(\"Function value obtained by UPOQA:\", obj(res.x, 0.0))\n",
    "\n",
    "# direct optimization fails\n",
    "import scipy\n",
    "sp_res = scipy.optimize.minimize(lambda x: obj(x, tau = 0.0), x0)\n",
    "print(\"Function value obtained by BFGS via direct optimization:\", obj(sp_res.x, 0.0))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
