Lagrangians.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lagrangians"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 561,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy import *\n",
    "(x, y, X, Y, Z, B, C, n, m, l, j, q, N, K, Q, S, k, epsilon, z, p1, p2, p3, g2, g3,\n",
    " z0, w1, w2, w3, e1, e2, e3, Delta, omega1, omega2, omega3, eta1, eta2, eta3, d1, d2, d3) = symbols(\n",
    "    '''x, y, X, Y, Z, B, C, n, m, l, j, q, N, K, Q, S, k, epsilon, z, p1, p2, p3, g2, g3,\n",
    "    z0, w1, w2, w3, e1, e2, e3, Delta, omega1, omega2, omega3, eta1, eta2, eta3, d1, d2, d3'''\n",
    ")\n",
    "t, nu, epsilon, theta = symbols('t, nu, epsilon, theta')\n",
    "pw = Function('pw') # Weierstrass P function\n",
    "pwp = Function('pwp') # Derivative of Weierstrass P function\n",
    "zw = Function('zw') # Weierstrass Zeta function\n",
    "sigma = Function('sigma') # Weierstrass Sigma function\n",
    "f = Function('f')\n",
    "h = Function('h')\n",
    "P = Function('P') # Polynomial\n",
    "s = Function(\"s\")\n",
    "Det = Function(\"Det\")\n",
    "kappa = IndexedBase('kappa')\n",
    "beta = IndexedBase('beta')\n",
    "mu = IndexedBase('mu')\n",
    "xi = IndexedBase('xi')\n",
    "rho = IndexedBase('rho')\n",
    "delta = IndexedBase('delta')\n",
    "b = IndexedBase('b')\n",
    "B = IndexedBase('B')\n",
    "c = IndexedBase('c')\n",
    "d = IndexedBase('d')\n",
    "p = IndexedBase('p')\n",
    "alpha = IndexedBase('alpha')\n",
    "Omega = IndexedBase('Omega')\n",
    "T = Function('T')\n",
    "V = Function('V')\n",
    "Dz = Function('Dz') # Derivate w.r.t z\n",
    "L = Function('L')\n",
    "U = Function(\"U\")\n",
    "W = Function(\"W\")\n",
    "from math import prod"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modes and Modal Powers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "Nmodes = 4\n",
    "A = [Function(f\"A{i}\") for i in range(Nmodes)] # Mode Functions of z\n",
    "Ac = [Function(f\"Ac{i}\") for i in range(Nmodes)] # Conjugate mode functions of z. May be complex conjugates but not always\n",
    "# Uvec = [Function(f\"U{i}\") for i in range(Nmodes)] # Modal Power Functions of z. May be positive real but not always\n",
    "Az = [A[i](z) for i in range(Nmodes)]\n",
    "Acz = [Ac[i](z) for i in range(Nmodes)]\n",
    "Uz = [Function(f\"U{i}\")(z) for i in range(Nmodes)]\n",
    "DAz = [diff(A[i](z),z) for i in range(Nmodes)] \n",
    "DAcz = [diff(Ac[i](z),z) for i in range(Nmodes)]\n",
    "phi = [Function(f\"phi{i}\") for i in range(Nmodes)] # Phases\n",
    "phiz = [phi[i](z) for i in range(Nmodes)] # Phases"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The terminology from optics will be used in which the complex function $A_i(z)$ is refered to as a mode. Each mode has a corresponding conjugate mode $Ac_i(z)$ where conjugate will often, but not always, mean complex conjugate. A modal power function will be denoted $P_i(z)$ and defined as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{U_{0}}{\\left(z \\right)} = \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(U0(z), A0(z)*Ac0(z))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(Uz[0], Az[0]*Acz[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{U_{1}}{\\left(z \\right)} = \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(U1(z), A1(z)*Ac1(z))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(Uz[1], Az[1]*Acz[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And so on. When the conjugation is complex conjugation $P_i(z)$ is a positive real function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lagrangians"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lagrangians will be considered involving $2N$ modes made from $N$ pairs of conjugate modes. Lagrangians will be considered that are defined as kinetic terms coming from some function of the modes and their derivatives, minus some potential terms coming from some function of the modes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Kinetic terms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Kinetic terms will be considered formed from each mode multiplied by the derivative of its conjugate, minus the corresponding conjugate pair, and weighted by the reciprocal of some constant $\\Omega_i$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle - \\frac{i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{0}} - \\frac{i \\left(\\operatorname{A_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{1}}{\\left(z \\right)} - \\operatorname{Ac_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{1}} - \\frac{i \\left(\\operatorname{A_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{2}}{\\left(z \\right)} - \\operatorname{Ac_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{2}} - \\frac{i \\left(\\operatorname{A_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{3}}{\\left(z \\right)} - \\operatorname{Ac_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{3}}$"
      ],
      "text/plain": [
       "-I*(A0(z)*Derivative(Ac0(z), z) - Ac0(z)*Derivative(A0(z), z))/(2*Omega[0]) - I*(A1(z)*Derivative(Ac1(z), z) - Ac1(z)*Derivative(A1(z), z))/(2*Omega[1]) - I*(A2(z)*Derivative(Ac2(z), z) - Ac2(z)*Derivative(A2(z), z))/(2*Omega[2]) - I*(A3(z)*Derivative(Ac3(z), z) - Ac3(z)*Derivative(A3(z), z))/(2*Omega[3])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kinetic = ((I/2*sum([(Acz[_i]*DAz[_i] - Az[_i]*DAcz[_i])/Omega[_i] for _i in range(Nmodes)]))\n",
    "           .expand()\n",
    "           .collect([Omega[_i] for _i in range(Nmodes)], factor))\n",
    "kinetic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Potential terms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Potential terms will be considered that are made from two parts. The first part provides a nonlinear coupling between modes in the form of a product of all modes plus a product of all the conjugate modes. The second part is some function, $P$ of pairs of conjugate modes multiplied together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right) \\kappa + P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)}$"
      ],
      "text/plain": [
       "(A0(z)*A1(z)*A2(z)*A3(z) + Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))*kappa + P(A0(z)*Ac0(z), A1(z)*Ac1(z), A2(z)*Ac2(z), A3(z)*Ac3(z))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "potential = kappa*(prod(Az)+prod(Acz)) + P(*[Az[_i]*Acz[_i] for _i in range(Nmodes)])\n",
    "# potential = kappa*(prod(Az)+prod(Acz)) + alpha*P(*[*Az, *Acz])\n",
    "potential"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An example potential is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{A_{0}}^{2}{\\left(z \\right)} \\operatorname{Ac_{0}}^{2}{\\left(z \\right)} + \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} \\kappa + \\operatorname{A_{1}}^{2}{\\left(z \\right)} \\operatorname{Ac_{1}}^{2}{\\left(z \\right)} + \\operatorname{A_{2}}^{2}{\\left(z \\right)} \\operatorname{Ac_{2}}^{2}{\\left(z \\right)} + \\operatorname{A_{3}}^{2}{\\left(z \\right)} \\operatorname{Ac_{3}}^{2}{\\left(z \\right)} + \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\kappa$"
      ],
      "text/plain": [
       "A0(z)**2*Ac0(z)**2 + A0(z)*A1(z)*A2(z)*A3(z)*kappa + A1(z)**2*Ac1(z)**2 + A2(z)**2*Ac2(z)**2 + A3(z)**2*Ac3(z)**2 + Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z)*kappa"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_potential = (kappa*(prod(Az)+prod(Acz)) + sum((Az[_i]*Acz[_i])**2 for _i in range(Nmodes))).expand()\n",
    "example_potential"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Equations of Motion from Euler-Lagrange Equations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The equation of motion for each mode will be obtained from the Lagrangian using the Euler-Lagrange equations for each mode. Note that the conjugate mode is considered distinct and unrelated for the Euler-Lagrange equations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle - \\frac{i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{0}} - \\frac{i \\left(\\operatorname{A_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{1}}{\\left(z \\right)} - \\operatorname{Ac_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{1}} - \\frac{i \\left(\\operatorname{A_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{2}}{\\left(z \\right)} - \\operatorname{Ac_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{2}} - \\frac{i \\left(\\operatorname{A_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{3}}{\\left(z \\right)} - \\operatorname{Ac_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{3}} - \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right) \\kappa - P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)}$"
      ],
      "text/plain": [
       "-I*(A0(z)*Derivative(Ac0(z), z) - Ac0(z)*Derivative(A0(z), z))/(2*Omega[0]) - I*(A1(z)*Derivative(Ac1(z), z) - Ac1(z)*Derivative(A1(z), z))/(2*Omega[1]) - I*(A2(z)*Derivative(Ac2(z), z) - Ac2(z)*Derivative(A2(z), z))/(2*Omega[2]) - I*(A3(z)*Derivative(Ac3(z), z) - Ac3(z)*Derivative(A3(z), z))/(2*Omega[3]) - (A0(z)*A1(z)*A2(z)*A3(z) + Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))*kappa - P(A0(z)*Ac0(z), A1(z)*Ac1(z), A2(z)*Ac2(z), A3(z)*Ac3(z))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lagrangian = kinetic - potential\n",
    "lagrangian"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an example, the below Euler-Lagrange equation leads to the equation of motion for $A_0(z)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d \\operatorname{Ac_{0}}{\\left(z \\right)}} L{\\left(\\operatorname{Ac_{0}}{\\left(z \\right)},\\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)} \\right)} - \\frac{d^{2}}{d zd \\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)}} L{\\left(\\operatorname{Ac_{0}}{\\left(z \\right)},\\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)} \\right)} = 0$"
      ],
      "text/plain": [
       "Eq(Derivative(L(Ac0(z), Derivative(Ac0(z), z)), Ac0(z)) - Derivative(L(Ac0(z), Derivative(Ac0(z), z)), Derivative(Ac0(z), z), z), 0)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(Derivative(L(Ac[0](z),Derivative(Ac[0](z),z)),Ac[0](z)) - \n",
    "Derivative(Derivative(L(Ac[0](z),Derivative(Ac[0](z),z)),Derivative(Ac[0](z),z)),z),0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "euler_lagrangian_eqs = [\n",
    "    Eq(\n",
    "        diff([*Az,*Acz][_i],z),\n",
    "        solve(diff(lagrangian, [*Acz,*Az][_i]) - diff(diff(lagrangian, [*DAcz,*DAz][_i]),z), diff([*Az,*Acz][_i],z))[0]\n",
    "    ).expand()\n",
    "    for _i in range(2*Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This produces equations of the motion for the modes that are of the form:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)} = - i \\operatorname{A_{0}}{\\left(z \\right)} {\\Omega}_{0} \\left. \\frac{\\partial}{\\partial \\xi_{1}} P{\\left(\\xi_{1},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{1}=\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} }} - i \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} {\\Omega}_{0} \\kappa$"
      ],
      "text/plain": [
       "Eq(Derivative(A0(z), z), -I*A0(z)*Omega[0]*Subs(Derivative(P(_xi_1, A1(z)*Ac1(z), A2(z)*Ac2(z), A3(z)*Ac3(z)), _xi_1), _xi_1, A0(z)*Ac0(z)) - I*Ac1(z)*Ac2(z)*Ac3(z)*Omega[0]*kappa)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)} = - i \\operatorname{A_{1}}{\\left(z \\right)} {\\Omega}_{1} \\left. \\frac{\\partial}{\\partial \\xi_{2}} P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\xi_{2},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{2}=\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} }} - i \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} {\\Omega}_{1} \\kappa$"
      ],
      "text/plain": [
       "Eq(Derivative(A1(z), z), -I*A1(z)*Omega[1]*Subs(Derivative(P(A0(z)*Ac0(z), _xi_2, A2(z)*Ac2(z), A3(z)*Ac3(z)), _xi_2), _xi_2, A1(z)*Ac1(z)) - I*Ac0(z)*Ac2(z)*Ac3(z)*Omega[1]*kappa)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)} = - i \\operatorname{A_{2}}{\\left(z \\right)} {\\Omega}_{2} \\left. \\frac{\\partial}{\\partial \\xi_{3}} P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\xi_{3},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{3}=\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} }} - i \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} {\\Omega}_{2} \\kappa$"
      ],
      "text/plain": [
       "Eq(Derivative(A2(z), z), -I*A2(z)*Omega[2]*Subs(Derivative(P(A0(z)*Ac0(z), A1(z)*Ac1(z), _xi_3, A3(z)*Ac3(z)), _xi_3), _xi_3, A2(z)*Ac2(z)) - I*Ac0(z)*Ac1(z)*Ac3(z)*Omega[2]*kappa)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)} = - i \\operatorname{A_{3}}{\\left(z \\right)} {\\Omega}_{3} \\left. \\frac{\\partial}{\\partial \\xi_{4}} P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\xi_{4} \\right)} \\right|_{\\substack{ \\xi_{4}=\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} }} - i \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} {\\Omega}_{3} \\kappa$"
      ],
      "text/plain": [
       "Eq(Derivative(A3(z), z), -I*A3(z)*Omega[3]*Subs(Derivative(P(A0(z)*Ac0(z), A1(z)*Ac1(z), A2(z)*Ac2(z), _xi_4), _xi_4), _xi_4, A3(z)*Ac3(z)) - I*Ac0(z)*Ac1(z)*Ac2(z)*Omega[3]*kappa)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)} = i \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} {\\Omega}_{0} \\kappa + i \\operatorname{Ac_{0}}{\\left(z \\right)} {\\Omega}_{0} \\left. \\frac{\\partial}{\\partial \\xi_{1}} P{\\left(\\xi_{1},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{1}=\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} }}$"
      ],
      "text/plain": [
       "Eq(Derivative(Ac0(z), z), I*A1(z)*A2(z)*A3(z)*Omega[0]*kappa + I*Ac0(z)*Omega[0]*Subs(Derivative(P(_xi_1, A1(z)*Ac1(z), A2(z)*Ac2(z), A3(z)*Ac3(z)), _xi_1), _xi_1, A0(z)*Ac0(z)))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{Ac_{1}}{\\left(z \\right)} = i \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} {\\Omega}_{1} \\kappa + i \\operatorname{Ac_{1}}{\\left(z \\right)} {\\Omega}_{1} \\left. \\frac{\\partial}{\\partial \\xi_{2}} P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\xi_{2},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{2}=\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} }}$"
      ],
      "text/plain": [
       "Eq(Derivative(Ac1(z), z), I*A0(z)*A2(z)*A3(z)*Omega[1]*kappa + I*Ac1(z)*Omega[1]*Subs(Derivative(P(A0(z)*Ac0(z), _xi_2, A2(z)*Ac2(z), A3(z)*Ac3(z)), _xi_2), _xi_2, A1(z)*Ac1(z)))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs[5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{Ac_{2}}{\\left(z \\right)} = i \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} {\\Omega}_{2} \\kappa + i \\operatorname{Ac_{2}}{\\left(z \\right)} {\\Omega}_{2} \\left. \\frac{\\partial}{\\partial \\xi_{3}} P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\xi_{3},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{3}=\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} }}$"
      ],
      "text/plain": [
       "Eq(Derivative(Ac2(z), z), I*A0(z)*A1(z)*A3(z)*Omega[2]*kappa + I*Ac2(z)*Omega[2]*Subs(Derivative(P(A0(z)*Ac0(z), A1(z)*Ac1(z), _xi_3, A3(z)*Ac3(z)), _xi_3), _xi_3, A2(z)*Ac2(z)))"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs[6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{Ac_{3}}{\\left(z \\right)} = i \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} {\\Omega}_{3} \\kappa + i \\operatorname{Ac_{3}}{\\left(z \\right)} {\\Omega}_{3} \\left. \\frac{\\partial}{\\partial \\xi_{4}} P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\xi_{4} \\right)} \\right|_{\\substack{ \\xi_{4}=\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} }}$"
      ],
      "text/plain": [
       "Eq(Derivative(Ac3(z), z), I*A0(z)*A1(z)*A2(z)*Omega[3]*kappa + I*Ac3(z)*Omega[3]*Subs(Derivative(P(A0(z)*Ac0(z), A1(z)*Ac1(z), A2(z)*Ac2(z), _xi_4), _xi_4), _xi_4, A3(z)*Ac3(z)))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs[7]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and equations of motion for the modal powers that are of the form:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "power_eqns = [\n",
    "    Eq(diff(Uz[_i],z), diff(Az[_i]*Acz[_i],z)).expand().subs([\n",
    "        (euler_lagrangian_eqs[_i].lhs, euler_lagrangian_eqs[_i].rhs) for _i in range(2*Nmodes)\n",
    "    ]).simplify()\n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)} = i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right) {\\Omega}_{0} \\kappa$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z), I*(A0(z)*A1(z)*A2(z)*A3(z) - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))*Omega[0]*kappa)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "power_eqns[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{1}}{\\left(z \\right)} = i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right) {\\Omega}_{1} \\kappa$"
      ],
      "text/plain": [
       "Eq(Derivative(U1(z), z), I*(A0(z)*A1(z)*A2(z)*A3(z) - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))*Omega[1]*kappa)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "power_eqns[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{2}}{\\left(z \\right)} = i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right) {\\Omega}_{2} \\kappa$"
      ],
      "text/plain": [
       "Eq(Derivative(U2(z), z), I*(A0(z)*A1(z)*A2(z)*A3(z) - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))*Omega[2]*kappa)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "power_eqns[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{3}}{\\left(z \\right)} = i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right) {\\Omega}_{3} \\kappa$"
      ],
      "text/plain": [
       "Eq(Derivative(U3(z), z), I*(A0(z)*A1(z)*A2(z)*A3(z) - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))*Omega[3]*kappa)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "power_eqns[3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conservation Laws"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conservation of total power"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The total modal power is defined as the sum of all modal powers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle s{\\left(z \\right)} = \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} + \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} + \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} + \\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(s(z), A0(z)*Ac0(z) + A1(z)*Ac1(z) + A2(z)*Ac2(z) + A3(z)*Ac3(z))"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sz_eq = Eq(s(z), sum([Az[_i]*Acz[_i] for _i in range(Nmodes)]))\n",
    "sz_eq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the equations of motion, it evolves accordingly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} s{\\left(z \\right)} = \\operatorname{A_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)} + \\operatorname{A_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{1}}{\\left(z \\right)} + \\operatorname{A_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{2}}{\\left(z \\right)} + \\operatorname{A_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{3}}{\\left(z \\right)} + \\operatorname{Ac_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)} + \\operatorname{Ac_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)} + \\operatorname{Ac_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)} + \\operatorname{Ac_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(s(z), z), A0(z)*Derivative(Ac0(z), z) + A1(z)*Derivative(Ac1(z), z) + A2(z)*Derivative(Ac2(z), z) + A3(z)*Derivative(Ac3(z), z) + Ac0(z)*Derivative(A0(z), z) + Ac1(z)*Derivative(A1(z), z) + Ac2(z)*Derivative(A2(z), z) + Ac3(z)*Derivative(A3(z), z))"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DPz_eq = Eq(diff(s(z),z), diff(s(z),z).subs(sz_eq.lhs, sz_eq.rhs)).doit().subs([\n",
    "    (power_eqn.lhs, power_eqn.rhs) for power_eqn in power_eqns\n",
    "])\n",
    "DPz_eq = Eq(DPz_eq.lhs, DPz_eq.rhs.factor())\n",
    "DPz_eq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For it to be a constant, a constraint must be placed on $\\Omega_i$ in the kinetic terms, that is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} s{\\left(z \\right)} = 0$"
      ],
      "text/plain": [
       "Eq(Derivative(s(z), z), 0)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(diff(s(z),z),0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "requires:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sum_{n=0}^{N} {\\Omega}_{n} = 0$"
      ],
      "text/plain": [
       "Eq(Sum(Omega[n], (n, 0, N)), 0)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(Sum(Omega[n],(n,0,N)),0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conservation of modal power combinations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By making $P$ a function of the modes multiplied by their conjugates, conserved quantities manifest for combinations of modal powers. These conservation laws mean that any modal power can be expressed in terms of another."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)} {\\Omega}_{1} - \\frac{d}{d z} \\operatorname{U_{1}}{\\left(z \\right)} {\\Omega}_{0} = 0$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z)*Omega[1] - Derivative(U1(z), z)*Omega[0], 0)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(power_eqns[0].lhs*Omega[1]-power_eqns[1].lhs*Omega[0], power_eqns[0].rhs*Omega[1]-power_eqns[1].rhs*Omega[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)} {\\Omega}_{2} - \\frac{d}{d z} \\operatorname{U_{2}}{\\left(z \\right)} {\\Omega}_{0} = 0$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z)*Omega[2] - Derivative(U2(z), z)*Omega[0], 0)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(power_eqns[0].lhs*Omega[2]-power_eqns[2].lhs*Omega[0], power_eqns[0].rhs*Omega[2]-power_eqns[2].rhs*Omega[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)} {\\Omega}_{3} - \\frac{d}{d z} \\operatorname{U_{3}}{\\left(z \\right)} {\\Omega}_{0} = 0$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z)*Omega[3] - Derivative(U3(z), z)*Omega[0], 0)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(power_eqns[0].lhs*Omega[3]-power_eqns[3].lhs*Omega[0], power_eqns[0].rhs*Omega[3]-power_eqns[3].rhs*Omega[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{1}}{\\left(z \\right)} {\\Omega}_{2} - \\frac{d}{d z} \\operatorname{U_{2}}{\\left(z \\right)} {\\Omega}_{1} = 0$"
      ],
      "text/plain": [
       "Eq(Derivative(U1(z), z)*Omega[2] - Derivative(U2(z), z)*Omega[1], 0)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(power_eqns[1].lhs*Omega[2]-power_eqns[2].lhs*Omega[1], power_eqns[1].rhs*Omega[2]-power_eqns[2].rhs*Omega[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{1}}{\\left(z \\right)} {\\Omega}_{3} - \\frac{d}{d z} \\operatorname{U_{3}}{\\left(z \\right)} {\\Omega}_{1} = 0$"
      ],
      "text/plain": [
       "Eq(Derivative(U1(z), z)*Omega[3] - Derivative(U3(z), z)*Omega[1], 0)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(power_eqns[1].lhs*Omega[3]-power_eqns[3].lhs*Omega[1], power_eqns[1].rhs*Omega[3]-power_eqns[3].rhs*Omega[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{2}}{\\left(z \\right)} {\\Omega}_{3} - \\frac{d}{d z} \\operatorname{U_{3}}{\\left(z \\right)} {\\Omega}_{2} = 0$"
      ],
      "text/plain": [
       "Eq(Derivative(U2(z), z)*Omega[3] - Derivative(U3(z), z)*Omega[2], 0)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(power_eqns[2].lhs*Omega[3]-power_eqns[3].lhs*Omega[2], power_eqns[2].rhs*Omega[3]-power_eqns[3].rhs*Omega[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {B}_{0,1} = \\operatorname{U_{0}}{\\left(z \\right)} {\\Omega}_{1} - \\operatorname{U_{1}}{\\left(z \\right)} {\\Omega}_{0}$"
      ],
      "text/plain": [
       "Eq(B[0, 1], U0(z)*Omega[1] - U1(z)*Omega[0])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B01 = Eq(B[0,1],Uz[0]*Omega[1] - Uz[1]*Omega[0])\n",
    "B01"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {B}_{0,2} = \\operatorname{U_{0}}{\\left(z \\right)} {\\Omega}_{2} - \\operatorname{U_{2}}{\\left(z \\right)} {\\Omega}_{0}$"
      ],
      "text/plain": [
       "Eq(B[0, 2], U0(z)*Omega[2] - U2(z)*Omega[0])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B02 = Eq(B[0,2],Uz[0]*Omega[2] - Uz[2]*Omega[0])\n",
    "B02"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {B}_{0,3} = \\operatorname{U_{0}}{\\left(z \\right)} {\\Omega}_{3} - \\operatorname{U_{3}}{\\left(z \\right)} {\\Omega}_{0}$"
      ],
      "text/plain": [
       "Eq(B[0, 3], U0(z)*Omega[3] - U3(z)*Omega[0])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B03 = Eq(B[0,3],Uz[0]*Omega[3] - Uz[3]*Omega[0])\n",
    "B03"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {B}_{1,0} = - \\operatorname{U_{0}}{\\left(z \\right)} {\\Omega}_{1} + \\operatorname{U_{1}}{\\left(z \\right)} {\\Omega}_{0}$"
      ],
      "text/plain": [
       "Eq(B[1, 0], -U0(z)*Omega[1] + U1(z)*Omega[0])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B10 = Eq(B[1,0],Uz[1]*Omega[0] - Uz[0]*Omega[1])\n",
    "B10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {B}_{1,2} = \\operatorname{U_{1}}{\\left(z \\right)} {\\Omega}_{2} - \\operatorname{U_{2}}{\\left(z \\right)} {\\Omega}_{1}$"
      ],
      "text/plain": [
       "Eq(B[1, 2], U1(z)*Omega[2] - U2(z)*Omega[1])"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B12 = Eq(B[1,2],Uz[1]*Omega[2] - Uz[2]*Omega[1])\n",
    "B12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {B}_{2,1} = - \\operatorname{U_{1}}{\\left(z \\right)} {\\Omega}_{2} + \\operatorname{U_{2}}{\\left(z \\right)} {\\Omega}_{1}$"
      ],
      "text/plain": [
       "Eq(B[2, 1], -U1(z)*Omega[2] + U2(z)*Omega[1])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B21 = Eq(B[2,1],Uz[2]*Omega[1] - Uz[1]*Omega[2])\n",
    "B21"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conservation of Energy (Hamiltonian)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The conservation of the potential (also sometimes referred to as conservation of energy or conservation of the Hamiltonian) is a consequence of any potential that lacks an explicit $z$ dependence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 0$"
      ],
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "diff(potential,z).subs([\n",
    "        (euler_lagrangian_eqs[_i].lhs, euler_lagrangian_eqs[_i].rhs) for _i in range(2*Nmodes)\n",
    "    ]).simplify()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle K = \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right) \\kappa + P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)}$"
      ],
      "text/plain": [
       "Eq(K, (A0(z)*A1(z)*A2(z)*A3(z) + Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))*kappa + P(A0(z)*Ac0(z), A1(z)*Ac1(z), A2(z)*Ac2(z), A3(z)*Ac3(z)))"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "K_eq = Eq(K, potential)\n",
    "K_eq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Squaring the equations of motion for the modal powers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Squaring the equations of motion for the powers of the modes, combined with the square of the product term in the conservation of $K$, enables the derivation of equations that are functions of $U_j$ and thus, throught the modal power combination conservation laws, expressible as a function of one modal power exclusively, i.e., they are separable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)}\\right)^{2} = - \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right)^{2} {\\Omega}_{0}^{2} \\kappa^{2}$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z)**2, -(A0(z)*A1(z)*A2(z)*A3(z) - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))**2*Omega[0]**2*kappa**2)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(power_eqns[0].lhs**2, power_eqns[0].rhs**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)}\\right)^{2} = - \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right)^{2} {\\Omega}_{0}^{2} \\kappa^{2} - \\left(- K + \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right) \\kappa + P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)}\\right)^{2}$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z)**2, -(A0(z)*A1(z)*A2(z)*A3(z) - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))**2*Omega[0]**2*kappa**2 - (-K + (A0(z)*A1(z)*A2(z)*A3(z) + Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))*kappa + P(A0(z)*Ac0(z), A1(z)*Ac1(z), A2(z)*Ac2(z), A3(z)*Ac3(z)))**2)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(power_eqns[0].lhs**2, power_eqns[0].rhs**2 - (K_eq.rhs - K)**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle K - P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)} = \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right) \\kappa$"
      ],
      "text/plain": [
       "Eq(K - P(A0(z)*Ac0(z), A1(z)*Ac1(z), A2(z)*Ac2(z), A3(z)*Ac3(z)), (A0(z)*A1(z)*A2(z)*A3(z) + Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))*kappa)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "K_eq_2 = Eq(K - K_eq.rhs + kappa*(prod(Az)+prod(Acz)) , K - K_eq.lhs + kappa*(prod(Az)+prod(Acz)))\n",
    "K_eq_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "squared_power_eqs = [\n",
    "    Eq(power_eqns[_i].lhs**2, (power_eqns[_i].rhs**2 + Omega[_i]**2*K_eq_2.rhs**2).expand()- Omega[_i]**2*K_eq_2.lhs**2)\n",
    "    .subs([(Acz[_j], Uz[_j]/Az[_j]) for _j in range(Nmodes)]) for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)}\\right)^{2} = - \\left(K - P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)}\\right)^{2} {\\Omega}_{0}^{2} + 4 \\operatorname{U_{0}}{\\left(z \\right)} \\operatorname{U_{1}}{\\left(z \\right)} \\operatorname{U_{2}}{\\left(z \\right)} \\operatorname{U_{3}}{\\left(z \\right)} {\\Omega}_{0}^{2} \\kappa^{2}$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z)**2, -(K - P(U0(z), U1(z), U2(z), U3(z)))**2*Omega[0]**2 + 4*U0(z)*U1(z)*U2(z)*U3(z)*Omega[0]**2*kappa**2)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squared_power_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{1}}{\\left(z \\right)}\\right)^{2} = - \\left(K - P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)}\\right)^{2} {\\Omega}_{1}^{2} + 4 \\operatorname{U_{0}}{\\left(z \\right)} \\operatorname{U_{1}}{\\left(z \\right)} \\operatorname{U_{2}}{\\left(z \\right)} \\operatorname{U_{3}}{\\left(z \\right)} {\\Omega}_{1}^{2} \\kappa^{2}$"
      ],
      "text/plain": [
       "Eq(Derivative(U1(z), z)**2, -(K - P(U0(z), U1(z), U2(z), U3(z)))**2*Omega[1]**2 + 4*U0(z)*U1(z)*U2(z)*U3(z)*Omega[1]**2*kappa**2)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squared_power_eqs[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{2}}{\\left(z \\right)}\\right)^{2} = - \\left(K - P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)}\\right)^{2} {\\Omega}_{2}^{2} + 4 \\operatorname{U_{0}}{\\left(z \\right)} \\operatorname{U_{1}}{\\left(z \\right)} \\operatorname{U_{2}}{\\left(z \\right)} \\operatorname{U_{3}}{\\left(z \\right)} {\\Omega}_{2}^{2} \\kappa^{2}$"
      ],
      "text/plain": [
       "Eq(Derivative(U2(z), z)**2, -(K - P(U0(z), U1(z), U2(z), U3(z)))**2*Omega[2]**2 + 4*U0(z)*U1(z)*U2(z)*U3(z)*Omega[2]**2*kappa**2)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squared_power_eqs[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{3}}{\\left(z \\right)}\\right)^{2} = - \\left(K - P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)}\\right)^{2} {\\Omega}_{3}^{2} + 4 \\operatorname{U_{0}}{\\left(z \\right)} \\operatorname{U_{1}}{\\left(z \\right)} \\operatorname{U_{2}}{\\left(z \\right)} \\operatorname{U_{3}}{\\left(z \\right)} {\\Omega}_{3}^{2} \\kappa^{2}$"
      ],
      "text/plain": [
       "Eq(Derivative(U3(z), z)**2, -(K - P(U0(z), U1(z), U2(z), U3(z)))**2*Omega[3]**2 + 4*U0(z)*U1(z)*U2(z)*U3(z)*Omega[3]**2*kappa**2)"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squared_power_eqs[3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Derivatives of the modal phases"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If mode conjugates are complex conjugates, the phase of mode $A_0$ would be:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\phi_{0}{\\left(z \\right)} = - \\frac{i \\log{\\left(\\frac{\\operatorname{A_{0}}{\\left(z \\right)}}{\\operatorname{Ac_{0}}{\\left(z \\right)}} \\right)}}{2}$"
      ],
      "text/plain": [
       "Eq(phi0(z), -I*log(A0(z)/Ac0(z))/2)"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phi0_eq = Eq(phiz[0],-I*ln(Az[0]/Acz[0])/2)\n",
    "phi0_eq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspired by this concept, equations for the modal phases are derived:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "phase_eqs = [\n",
    "    Eq(\n",
    "        diff(phiz[_i],z), \n",
    "        diff(-I*ln(Az[_i]/Acz[_i])/2,z).simplify().subs([\n",
    "            (euler_lagrangian_eqs[_i].lhs, euler_lagrangian_eqs[_i].rhs) for _i in range(2*Nmodes)\n",
    "        ]).simplify() + (\n",
    "            K_eq_2.rhs/(Az[_i]*Acz[_i])*Omega[_i] - K_eq_2.lhs/(Az[_i]*Acz[_i])*Omega[_i]\n",
    "        ).expand()/2\n",
    "      ).expand().subs([(Acz[_i], Uz[_i]/Az[_i]) for _i in range(Nmodes)]) for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\phi_{0}{\\left(z \\right)} = - \\frac{K {\\Omega}_{0}}{2 \\operatorname{U_{0}}{\\left(z \\right)}} + \\frac{P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} {\\Omega}_{0}}{2 \\operatorname{U_{0}}{\\left(z \\right)}} - {\\Omega}_{0} \\left. \\frac{\\partial}{\\partial \\xi_{1}} P{\\left(\\xi_{1},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{1}=\\operatorname{U_{0}}{\\left(z \\right)} }}$"
      ],
      "text/plain": [
       "Eq(Derivative(phi0(z), z), -K*Omega[0]/(2*U0(z)) + P(U0(z), U1(z), U2(z), U3(z))*Omega[0]/(2*U0(z)) - Omega[0]*Subs(Derivative(P(_xi_1, U1(z), U2(z), U3(z)), _xi_1), _xi_1, U0(z)))"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phase_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\phi_{1}{\\left(z \\right)} = - \\frac{K {\\Omega}_{1}}{2 \\operatorname{U_{1}}{\\left(z \\right)}} + \\frac{P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} {\\Omega}_{1}}{2 \\operatorname{U_{1}}{\\left(z \\right)}} - {\\Omega}_{1} \\left. \\frac{\\partial}{\\partial \\xi_{2}} P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\xi_{2},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{2}=\\operatorname{U_{1}}{\\left(z \\right)} }}$"
      ],
      "text/plain": [
       "Eq(Derivative(phi1(z), z), -K*Omega[1]/(2*U1(z)) + P(U0(z), U1(z), U2(z), U3(z))*Omega[1]/(2*U1(z)) - Omega[1]*Subs(Derivative(P(U0(z), _xi_2, U2(z), U3(z)), _xi_2), _xi_2, U1(z)))"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phase_eqs[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\phi_{2}{\\left(z \\right)} = - \\frac{K {\\Omega}_{2}}{2 \\operatorname{U_{2}}{\\left(z \\right)}} + \\frac{P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} {\\Omega}_{2}}{2 \\operatorname{U_{2}}{\\left(z \\right)}} - {\\Omega}_{2} \\left. \\frac{\\partial}{\\partial \\xi_{3}} P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\xi_{3},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{3}=\\operatorname{U_{2}}{\\left(z \\right)} }}$"
      ],
      "text/plain": [
       "Eq(Derivative(phi2(z), z), -K*Omega[2]/(2*U2(z)) + P(U0(z), U1(z), U2(z), U3(z))*Omega[2]/(2*U2(z)) - Omega[2]*Subs(Derivative(P(U0(z), U1(z), _xi_3, U3(z)), _xi_3), _xi_3, U2(z)))"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phase_eqs[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\phi_{3}{\\left(z \\right)} = - \\frac{K {\\Omega}_{3}}{2 \\operatorname{U_{3}}{\\left(z \\right)}} + \\frac{P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} {\\Omega}_{3}}{2 \\operatorname{U_{3}}{\\left(z \\right)}} - {\\Omega}_{3} \\left. \\frac{\\partial}{\\partial \\xi_{4}} P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\xi_{4} \\right)} \\right|_{\\substack{ \\xi_{4}=\\operatorname{U_{3}}{\\left(z \\right)} }}$"
      ],
      "text/plain": [
       "Eq(Derivative(phi3(z), z), -K*Omega[3]/(2*U3(z)) + P(U0(z), U1(z), U2(z), U3(z))*Omega[3]/(2*U3(z)) - Omega[3]*Subs(Derivative(P(U0(z), U1(z), U2(z), _xi_4), _xi_4), _xi_4, U3(z)))"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phase_eqs[3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Logarithmic derivatives of the modes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If mode conjugates are complex conjugates, the logarithmic derivative of the $A_0$ mode would satisfy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)}}{\\operatorname{A_{0}}{\\left(z \\right)}} = i \\frac{d}{d z} \\phi_{0}{\\left(z \\right)} + \\frac{\\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)}}{2 \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)}}$"
      ],
      "text/plain": [
       "Eq(Derivative(A0(z), z)/A0(z), I*Derivative(phi0(z), z) + Derivative(A0(z)*Ac0(z), z)/(2*A0(z)*Ac0(z)))"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(\n",
    "    (Derivative(Az[0]*Acz[0],z)/(Az[0]*Acz[0])/2 + Derivative(Az[0]/Acz[0],z)/(Az[0]/Acz[0])/2).doit().simplify(),\n",
    "    Derivative(Az[0]*Acz[0],z)/(Az[0]*Acz[0])/2 + diff(exp(2*I*phiz[0]),z)/(exp(2*I*phiz[0]))/2\n",
    "  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_diff_eqns = [\n",
    "    Eq(\n",
    "        diff(Az[_i],z)/Az[_i], (I*diff(phiz[_i],z) + diff(Uz[_i],z)/Uz[_i]/2).subs(phase_eqs[_i].lhs, phase_eqs[_i].rhs)\n",
    "    ).expand()\n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)}}{\\operatorname{A_{0}}{\\left(z \\right)}} = - \\frac{i K {\\Omega}_{0}}{2 \\operatorname{U_{0}}{\\left(z \\right)}} + \\frac{i P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} {\\Omega}_{0}}{2 \\operatorname{U_{0}}{\\left(z \\right)}} - i {\\Omega}_{0} \\left. \\frac{\\partial}{\\partial \\xi_{1}} P{\\left(\\xi_{1},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{1}=\\operatorname{U_{0}}{\\left(z \\right)} }} + \\frac{\\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)}}{2 \\operatorname{U_{0}}{\\left(z \\right)}}$"
      ],
      "text/plain": [
       "Eq(Derivative(A0(z), z)/A0(z), -I*K*Omega[0]/(2*U0(z)) + I*P(U0(z), U1(z), U2(z), U3(z))*Omega[0]/(2*U0(z)) - I*Omega[0]*Subs(Derivative(P(_xi_1, U1(z), U2(z), U3(z)), _xi_1), _xi_1, U0(z)) + Derivative(U0(z), z)/(2*U0(z)))"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_eqns[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)}}{\\operatorname{A_{1}}{\\left(z \\right)}} = - \\frac{i K {\\Omega}_{1}}{2 \\operatorname{U_{1}}{\\left(z \\right)}} + \\frac{i P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} {\\Omega}_{1}}{2 \\operatorname{U_{1}}{\\left(z \\right)}} - i {\\Omega}_{1} \\left. \\frac{\\partial}{\\partial \\xi_{2}} P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\xi_{2},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{2}=\\operatorname{U_{1}}{\\left(z \\right)} }} + \\frac{\\frac{d}{d z} \\operatorname{U_{1}}{\\left(z \\right)}}{2 \\operatorname{U_{1}}{\\left(z \\right)}}$"
      ],
      "text/plain": [
       "Eq(Derivative(A1(z), z)/A1(z), -I*K*Omega[1]/(2*U1(z)) + I*P(U0(z), U1(z), U2(z), U3(z))*Omega[1]/(2*U1(z)) - I*Omega[1]*Subs(Derivative(P(U0(z), _xi_2, U2(z), U3(z)), _xi_2), _xi_2, U1(z)) + Derivative(U1(z), z)/(2*U1(z)))"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_eqns[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)}}{\\operatorname{A_{2}}{\\left(z \\right)}} = - \\frac{i K {\\Omega}_{2}}{2 \\operatorname{U_{2}}{\\left(z \\right)}} + \\frac{i P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} {\\Omega}_{2}}{2 \\operatorname{U_{2}}{\\left(z \\right)}} - i {\\Omega}_{2} \\left. \\frac{\\partial}{\\partial \\xi_{3}} P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\xi_{3},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} \\right|_{\\substack{ \\xi_{3}=\\operatorname{U_{2}}{\\left(z \\right)} }} + \\frac{\\frac{d}{d z} \\operatorname{U_{2}}{\\left(z \\right)}}{2 \\operatorname{U_{2}}{\\left(z \\right)}}$"
      ],
      "text/plain": [
       "Eq(Derivative(A2(z), z)/A2(z), -I*K*Omega[2]/(2*U2(z)) + I*P(U0(z), U1(z), U2(z), U3(z))*Omega[2]/(2*U2(z)) - I*Omega[2]*Subs(Derivative(P(U0(z), U1(z), _xi_3, U3(z)), _xi_3), _xi_3, U2(z)) + Derivative(U2(z), z)/(2*U2(z)))"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_eqns[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)}}{\\operatorname{A_{3}}{\\left(z \\right)}} = - \\frac{i K {\\Omega}_{3}}{2 \\operatorname{U_{3}}{\\left(z \\right)}} + \\frac{i P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} {\\Omega}_{3}}{2 \\operatorname{U_{3}}{\\left(z \\right)}} - i {\\Omega}_{3} \\left. \\frac{\\partial}{\\partial \\xi_{4}} P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\xi_{4} \\right)} \\right|_{\\substack{ \\xi_{4}=\\operatorname{U_{3}}{\\left(z \\right)} }} + \\frac{\\frac{d}{d z} \\operatorname{U_{3}}{\\left(z \\right)}}{2 \\operatorname{U_{3}}{\\left(z \\right)}}$"
      ],
      "text/plain": [
       "Eq(Derivative(A3(z), z)/A3(z), -I*K*Omega[3]/(2*U3(z)) + I*P(U0(z), U1(z), U2(z), U3(z))*Omega[3]/(2*U3(z)) - I*Omega[3]*Subs(Derivative(P(U0(z), U1(z), U2(z), _xi_4), _xi_4), _xi_4, U3(z)) + Derivative(U3(z), z)/(2*U3(z)))"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_eqns[3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And thus, through the conservation laws that relate $U_j$, the logarithmic derivatives of the modes are separable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finding an elliptic example with at most cubic terms for a Lagrangian with 4 modes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "# P_quadratic_eq = Eq(\n",
    "#     P(*[Uz[_i] for _i in range(Nmodes)]),\n",
    "#     Sum(Sum(Sum(Sum(c[j,k,l,m]*U(j,z)**l*U(k,z)**m,(j,0,3)), (k,0,3)), (l,0,1)), (m,0,1))\n",
    "# )\n",
    "# P_quadratic_eq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The aim is to find the most general form of P that will leave the right hand side of the following equation cubic in $U_0(z)$, so the aim is to choose P to cancel the quartic term coming from the modal power product term. The motivation is to make it solvable quite simply in terms of the Weierstrass P function. Quartic terms could also be solved but it would involve having to find the roots of a quartic and performing a bilinear transformation which is not the focus of this work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)}\\right)^{2} = - \\left(K - P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)}\\right)^{2} {\\Omega}_{0}^{2} + 4 \\operatorname{U_{0}}{\\left(z \\right)} \\operatorname{U_{1}}{\\left(z \\right)} \\operatorname{U_{2}}{\\left(z \\right)} \\operatorname{U_{3}}{\\left(z \\right)} {\\Omega}_{0}^{2} \\kappa^{2}$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z)**2, -(K - P(U0(z), U1(z), U2(z), U3(z)))**2*Omega[0]**2 + 4*U0(z)*U1(z)*U2(z)*U3(z)*Omega[0]**2*kappa**2)"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squared_power_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{U_{0}}{\\left(z \\right)} \\operatorname{U_{1}}{\\left(z \\right)} \\operatorname{U_{2}}{\\left(z \\right)} \\operatorname{U_{3}}{\\left(z \\right)} = \\frac{\\left(\\operatorname{U_{0}}{\\left(z \\right)} {\\Omega}_{1} - {B}_{0,1}\\right) \\left(\\operatorname{U_{0}}{\\left(z \\right)} {\\Omega}_{2} - {B}_{0,2}\\right) \\left(\\operatorname{U_{0}}{\\left(z \\right)} {\\Omega}_{3} - {B}_{0,3}\\right) \\operatorname{U_{0}}{\\left(z \\right)}}{{\\Omega}_{0}^{3}}$"
      ],
      "text/plain": [
       "Eq(U0(z)*U1(z)*U2(z)*U3(z), (U0(z)*Omega[1] - B[0, 1])*(U0(z)*Omega[2] - B[0, 2])*(U0(z)*Omega[3] - B[0, 3])*U0(z)/Omega[0]**3)"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u_4_term_prod = Eq(Uz[0]*Uz[1]*Uz[2]*Uz[3], (Uz[0]*Uz[1]*Uz[2]*Uz[3]).subs([\n",
    "                         (Uz[1],solve(B01,Uz[1])[0]), (Uz[2],solve(B02,Uz[2])[0]), (Uz[3],solve(B03,Uz[3])[0])\n",
    "                     ]))\n",
    "u_4_term_prod"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As $P$ will get squared, if when doing so it is to have no powers of $U$ greater than quartic, $P$ can be no greater than quadratic in $U$ terms. Thus without loss of generality, let:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle P{\\left(\\operatorname{U_{0}}{\\left(z \\right)},\\operatorname{U_{1}}{\\left(z \\right)},\\operatorname{U_{2}}{\\left(z \\right)},\\operatorname{U_{3}}{\\left(z \\right)} \\right)} = \\left(\\operatorname{U_{0}}{\\left(z \\right)} - {d}_{1}\\right) \\left(\\operatorname{U_{0}}{\\left(z \\right)} - {d}_{2}\\right) \\alpha$"
      ],
      "text/plain": [
       "Eq(P(U0(z), U1(z), U2(z), U3(z)), (U0(z) - d[1])*(U0(z) - d[2])*alpha)"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "P_u_2_term_prod = Eq(P(Uz[0],Uz[1],Uz[2],Uz[3]), alpha*(Uz[0] - d[1])*(Uz[0] - d[2]))\n",
    "P_u_2_term_prod"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "then to cancel the quartic terms in $U_0$, $\\alpha$ should be fixed such that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle - {\\Omega}_{0}^{2} \\alpha^{2} + \\frac{4 {\\Omega}_{1} {\\Omega}_{2} {\\Omega}_{3} \\kappa^{2}}{{\\Omega}_{0}} = 0$"
      ],
      "text/plain": [
       "Eq(-Omega[0]**2*alpha**2 + 4*Omega[1]*Omega[2]*Omega[3]*kappa**2/Omega[0], 0)"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "no_quartic_alpha_eq = Eq(squared_power_eqs[0].rhs.subs([\n",
    "    (u_4_term_prod.lhs, u_4_term_prod.rhs),\n",
    "    (P_u_2_term_prod.lhs, P_u_2_term_prod.rhs)\n",
    "]).expand().coeff(Uz[0]**4),0)\n",
    "no_quartic_alpha_eq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which has two solutions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\alpha = - 2 \\sqrt{\\frac{{\\Omega}_{1} {\\Omega}_{2} {\\Omega}_{3}}{{\\Omega}_{0}^{3}}} \\kappa$"
      ],
      "text/plain": [
       "Eq(alpha, -2*sqrt(Omega[1]*Omega[2]*Omega[3]/Omega[0]**3)*kappa)"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alpha_minus_eq = Eq(alpha, solve(no_quartic_alpha_eq,alpha)[0])\n",
    "alpha_minus_eq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\alpha = 2 \\sqrt{\\frac{{\\Omega}_{1} {\\Omega}_{2} {\\Omega}_{3}}{{\\Omega}_{0}^{3}}} \\kappa$"
      ],
      "text/plain": [
       "Eq(alpha, 2*sqrt(Omega[1]*Omega[2]*Omega[3]/Omega[0]**3)*kappa)"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alpha_plus_eq = Eq(alpha, solve(no_quartic_alpha_eq,alpha)[1])\n",
    "alpha_plus_eq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One Lagrangian that woud yield a cubic in $U_0$ is thus:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle - \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} - {d}_{1}\\right) \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} - {d}_{2}\\right) \\alpha - \\frac{i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{0}} - \\frac{i \\left(\\operatorname{A_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{1}}{\\left(z \\right)} - \\operatorname{Ac_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{1}} - \\frac{i \\left(\\operatorname{A_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{2}}{\\left(z \\right)} - \\operatorname{Ac_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{2}} - \\frac{i \\left(\\operatorname{A_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{3}}{\\left(z \\right)} - \\operatorname{Ac_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{3}} - \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right) \\kappa$"
      ],
      "text/plain": [
       "-(A0(z)*Ac0(z) - d[1])*(A0(z)*Ac0(z) - d[2])*alpha - I*(A0(z)*Derivative(Ac0(z), z) - Ac0(z)*Derivative(A0(z), z))/(2*Omega[0]) - I*(A1(z)*Derivative(Ac1(z), z) - Ac1(z)*Derivative(A1(z), z))/(2*Omega[1]) - I*(A2(z)*Derivative(Ac2(z), z) - Ac2(z)*Derivative(A2(z), z))/(2*Omega[2]) - I*(A3(z)*Derivative(Ac3(z), z) - Ac3(z)*Derivative(A3(z), z))/(2*Omega[3]) - (A0(z)*A1(z)*A2(z)*A3(z) + Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))*kappa"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lagrangian_2_term_alpha = lagrangian.subs(\n",
    "    P_u_2_term_prod.subs([(Uz[_i], Az[_i]*Acz[_i]) for _i in range(Nmodes)]).lhs,\n",
    "    P_u_2_term_prod.subs([(Uz[_i], Az[_i]*Acz[_i]) for _i in range(Nmodes)]).rhs\n",
    ")\n",
    "lagrangian_2_term_alpha"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle - 2 \\sqrt{\\frac{{\\Omega}_{1} {\\Omega}_{2} {\\Omega}_{3}}{{\\Omega}_{0}^{3}}} \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} - {d}_{1}\\right) \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} - {d}_{2}\\right) \\kappa - \\frac{i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{0}} - \\frac{i \\left(\\operatorname{A_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{1}}{\\left(z \\right)} - \\operatorname{Ac_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{1}} - \\frac{i \\left(\\operatorname{A_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{2}}{\\left(z \\right)} - \\operatorname{Ac_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{2}} - \\frac{i \\left(\\operatorname{A_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{3}}{\\left(z \\right)} - \\operatorname{Ac_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)}\\right)}{2 {\\Omega}_{3}} - \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right) \\kappa$"
      ],
      "text/plain": [
       "-2*sqrt(Omega[1]*Omega[2]*Omega[3]/Omega[0]**3)*(A0(z)*Ac0(z) - d[1])*(A0(z)*Ac0(z) - d[2])*kappa - I*(A0(z)*Derivative(Ac0(z), z) - Ac0(z)*Derivative(A0(z), z))/(2*Omega[0]) - I*(A1(z)*Derivative(Ac1(z), z) - Ac1(z)*Derivative(A1(z), z))/(2*Omega[1]) - I*(A2(z)*Derivative(Ac2(z), z) - Ac2(z)*Derivative(A2(z), z))/(2*Omega[2]) - I*(A3(z)*Derivative(Ac3(z), z) - Ac3(z)*Derivative(A3(z), z))/(2*Omega[3]) - (A0(z)*A1(z)*A2(z)*A3(z) + Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))*kappa"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lagrangian_2_term_alpha_no_quartic = lagrangian_2_term_alpha.subs(alpha_plus_eq.lhs, alpha_plus_eq.rhs)\n",
    "lagrangian_2_term_alpha_no_quartic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which after a suitable rescaling of the fields and parameters can be written:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle - 2 \\operatorname{A_{0}}^{2}{\\left(z \\right)} \\operatorname{Ac_{0}}^{2}{\\left(z \\right)} - \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} {d}_{0} - \\frac{i \\operatorname{A_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)}}{2} + \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} {d}_{1} - \\frac{i \\operatorname{A_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{1}}{\\left(z \\right)}}{2} + \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} {d}_{2} - \\frac{i \\operatorname{A_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{2}}{\\left(z \\right)}}{2} + \\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} {d}_{3} - \\frac{i \\operatorname{A_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{3}}{\\left(z \\right)}}{2} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} + \\frac{i \\operatorname{Ac_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)}}{2} + \\frac{i \\operatorname{Ac_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)}}{2} + \\frac{i \\operatorname{Ac_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)}}{2} + \\frac{i \\operatorname{Ac_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)}}{2}$"
      ],
      "text/plain": [
       "-2*A0(z)**2*Ac0(z)**2 - A0(z)*A1(z)*A2(z)*A3(z) + A0(z)*Ac0(z)*d[0] - I*A0(z)*Derivative(Ac0(z), z)/2 + A1(z)*Ac1(z)*d[1] - I*A1(z)*Derivative(Ac1(z), z)/2 + A2(z)*Ac2(z)*d[2] - I*A2(z)*Derivative(Ac2(z), z)/2 + A3(z)*Ac3(z)*d[3] - I*A3(z)*Derivative(Ac3(z), z)/2 - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z) + I*Ac0(z)*Derivative(A0(z), z)/2 + I*Ac1(z)*Derivative(A1(z), z)/2 + I*Ac2(z)*Derivative(A2(z), z)/2 + I*Ac3(z)*Derivative(A3(z), z)/2"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "L_4m_ex1_a = ((lagrangian_2_term_alpha_no_quartic.subs([    \n",
    "    *[(Az[_i], sqrt(Omega[_i])*Az[_i]) for _i in range(Nmodes)],\n",
    "    *[(Acz[_i], sqrt(Omega[_i])*Acz[_i]) for _i in range(Nmodes)],\n",
    "]).simplify().expand().subs([\n",
    "    (sqrt(Omega[1]*Omega[2]*Omega[3]/Omega[0]**3), \n",
    "     sqrt(Omega[1])*sqrt(Omega[2])*sqrt(Omega[3])/sqrt(Omega[0])**3\n",
    "    )\n",
    "]).subs([\n",
    "    (kappa, kappa/(sqrt(Omega[0])*sqrt(Omega[1])*sqrt(Omega[2])*sqrt(Omega[3]))),\n",
    "    (d[1],d[1]*Omega[0]), (d[2],d[2]*Omega[0])\n",
    "])/kappa).expand().subs(kappa,1).subs(2*d[1]*d[2],0).subs(d[2],d[0]/2-d[1]).expand() + \n",
    "           Az[1]*Acz[1]*d[1] + Az[2]*Acz[2]*d[2] + Az[3]*Acz[3]*d[3]\n",
    "           )\n",
    "L_4m_ex1_a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which can then be made symmetric in modal powers by appropriately adding some squares of the other modal powers and compensating with a reduction in the weight of $|A_0|^2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle - \\frac{\\operatorname{A_{0}}^{2}{\\left(z \\right)} \\operatorname{Ac_{0}}^{2}{\\left(z \\right)}}{2} - \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} {d}_{0} - \\frac{i \\operatorname{A_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)}}{2} - \\frac{\\operatorname{A_{1}}^{2}{\\left(z \\right)} \\operatorname{Ac_{1}}^{2}{\\left(z \\right)}}{2} + \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} {d}_{1} - \\frac{i \\operatorname{A_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{1}}{\\left(z \\right)}}{2} - \\frac{\\operatorname{A_{2}}^{2}{\\left(z \\right)} \\operatorname{Ac_{2}}^{2}{\\left(z \\right)}}{2} + \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} {d}_{2} - \\frac{i \\operatorname{A_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{2}}{\\left(z \\right)}}{2} - \\frac{\\operatorname{A_{3}}^{2}{\\left(z \\right)} \\operatorname{Ac_{3}}^{2}{\\left(z \\right)}}{2} + \\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} {d}_{3} - \\frac{i \\operatorname{A_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{Ac_{3}}{\\left(z \\right)}}{2} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} + \\frac{i \\operatorname{Ac_{0}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)}}{2} + \\frac{i \\operatorname{Ac_{1}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)}}{2} + \\frac{i \\operatorname{Ac_{2}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)}}{2} + \\frac{i \\operatorname{Ac_{3}}{\\left(z \\right)} \\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)}}{2}$"
      ],
      "text/plain": [
       "-A0(z)**2*Ac0(z)**2/2 - A0(z)*A1(z)*A2(z)*A3(z) + A0(z)*Ac0(z)*d[0] - I*A0(z)*Derivative(Ac0(z), z)/2 - A1(z)**2*Ac1(z)**2/2 + A1(z)*Ac1(z)*d[1] - I*A1(z)*Derivative(Ac1(z), z)/2 - A2(z)**2*Ac2(z)**2/2 + A2(z)*Ac2(z)*d[2] - I*A2(z)*Derivative(Ac2(z), z)/2 - A3(z)**2*Ac3(z)**2/2 + A3(z)*Ac3(z)*d[3] - I*A3(z)*Derivative(Ac3(z), z)/2 - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z) + I*Ac0(z)*Derivative(A0(z), z)/2 + I*Ac1(z)*Derivative(A1(z), z)/2 + I*Ac2(z)*Derivative(A2(z), z)/2 + I*Ac3(z)*Derivative(A3(z), z)/2"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "L_4m_ex1 = L_4m_ex1_a - ((Az[0]*Acz[0])**2 + (Az[1]*Acz[1])**2 + (Az[2]*Acz[2])**2 + (Az[3]*Acz[3])**2)/2 + 2*(Az[0]*Acz[0])**2\n",
    "L_4m_ex1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will conserve the constant $K$ (potential terms):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle K = - \\frac{\\operatorname{A_{0}}^{2}{\\left(z \\right)} \\operatorname{Ac_{0}}^{2}{\\left(z \\right)}}{2} - \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} {d}_{0} - \\frac{\\operatorname{A_{1}}^{2}{\\left(z \\right)} \\operatorname{Ac_{1}}^{2}{\\left(z \\right)}}{2} + \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} {d}_{1} - \\frac{\\operatorname{A_{2}}^{2}{\\left(z \\right)} \\operatorname{Ac_{2}}^{2}{\\left(z \\right)}}{2} + \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} {d}_{2} - \\frac{\\operatorname{A_{3}}^{2}{\\left(z \\right)} \\operatorname{Ac_{3}}^{2}{\\left(z \\right)}}{2} + \\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} {d}_{3} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(K, -A0(z)**2*Ac0(z)**2/2 - A0(z)*A1(z)*A2(z)*A3(z) + A0(z)*Ac0(z)*d[0] - A1(z)**2*Ac1(z)**2/2 + A1(z)*Ac1(z)*d[1] - A2(z)**2*Ac2(z)**2/2 + A2(z)*Ac2(z)*d[2] - A3(z)**2*Ac3(z)**2/2 + A3(z)*Ac3(z)*d[3] - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "K_ex1 = Eq(K, L_4m_ex1 - kinetic.subs([(Omega[_i],1) for _i in range(Nmodes)]).expand())\n",
    "K_ex1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle K + \\frac{\\operatorname{A_{0}}^{2}{\\left(z \\right)} \\operatorname{Ac_{0}}^{2}{\\left(z \\right)}}{2} - \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} {d}_{0} + \\frac{\\operatorname{A_{1}}^{2}{\\left(z \\right)} \\operatorname{Ac_{1}}^{2}{\\left(z \\right)}}{2} - \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} {d}_{1} + \\frac{\\operatorname{A_{2}}^{2}{\\left(z \\right)} \\operatorname{Ac_{2}}^{2}{\\left(z \\right)}}{2} - \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} {d}_{2} + \\frac{\\operatorname{A_{3}}^{2}{\\left(z \\right)} \\operatorname{Ac_{3}}^{2}{\\left(z \\right)}}{2} - \\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} {d}_{3} = - \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(K + A0(z)**2*Ac0(z)**2/2 - A0(z)*Ac0(z)*d[0] + A1(z)**2*Ac1(z)**2/2 - A1(z)*Ac1(z)*d[1] + A2(z)**2*Ac2(z)**2/2 - A2(z)*Ac2(z)*d[2] + A3(z)**2*Ac3(z)**2/2 - A3(z)*Ac3(z)*d[3], -A0(z)*A1(z)*A2(z)*A3(z) - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z))"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "K_eq_2_ex1 = Eq(K - K_ex1.rhs - (prod(Az)+prod(Acz)) , K - K_ex1.lhs - (prod(Az)+prod(Acz)))\n",
    "K_eq_2_ex1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "such that the polynomial part $P$ is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle P{\\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)},\\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)},\\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)},\\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} \\right)} = - \\frac{\\operatorname{A_{0}}^{2}{\\left(z \\right)} \\operatorname{Ac_{0}}^{2}{\\left(z \\right)}}{2} + \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} {d}_{0} - \\frac{\\operatorname{A_{1}}^{2}{\\left(z \\right)} \\operatorname{Ac_{1}}^{2}{\\left(z \\right)}}{2} + \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} {d}_{1} - \\frac{\\operatorname{A_{2}}^{2}{\\left(z \\right)} \\operatorname{Ac_{2}}^{2}{\\left(z \\right)}}{2} + \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} {d}_{2} - \\frac{\\operatorname{A_{3}}^{2}{\\left(z \\right)} \\operatorname{Ac_{3}}^{2}{\\left(z \\right)}}{2} + \\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} {d}_{3}$"
      ],
      "text/plain": [
       "Eq(P(A0(z)*Ac0(z), A1(z)*Ac1(z), A2(z)*Ac2(z), A3(z)*Ac3(z)), -A0(z)**2*Ac0(z)**2/2 + A0(z)*Ac0(z)*d[0] - A1(z)**2*Ac1(z)**2/2 + A1(z)*Ac1(z)*d[1] - A2(z)**2*Ac2(z)**2/2 + A2(z)*Ac2(z)*d[2] - A3(z)**2*Ac3(z)**2/2 + A3(z)*Ac3(z)*d[3])"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "P_ex1 = Eq(P(*[Az[_i]*Acz[_i] for _i in range(Nmodes)]), K_ex1.rhs + (prod(Az)+prod(Acz)))\n",
    "P_ex1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The equations of motion coming from the Euler-Lagrange equations are thus:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "euler_lagrangian_eqs_ex1 = [\n",
    "    Eq(\n",
    "        diff([*Az,*Acz][_i],z),\n",
    "        solve(diff(L_4m_ex1, [*Acz,*Az][_i]) - diff(diff(L_4m_ex1, [*DAcz,*DAz][_i]),z), diff([*Az,*Acz][_i],z))[0]\n",
    "    ).expand()\n",
    "    for _i in range(2*Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)} = - i \\operatorname{A_{0}}^{2}{\\left(z \\right)} \\operatorname{Ac_{0}}{\\left(z \\right)} + i \\operatorname{A_{0}}{\\left(z \\right)} {d}_{0} - i \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(A0(z), z), -I*A0(z)**2*Ac0(z) + I*A0(z)*d[0] - I*Ac1(z)*Ac2(z)*Ac3(z))"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs_ex1[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)} = - i \\operatorname{A_{1}}^{2}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} + i \\operatorname{A_{1}}{\\left(z \\right)} {d}_{1} - i \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(A1(z), z), -I*A1(z)**2*Ac1(z) + I*A1(z)*d[1] - I*Ac0(z)*Ac2(z)*Ac3(z))"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs_ex1[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)} = - i \\operatorname{A_{2}}^{2}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} + i \\operatorname{A_{2}}{\\left(z \\right)} {d}_{2} - i \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(A2(z), z), -I*A2(z)**2*Ac2(z) + I*A2(z)*d[2] - I*Ac0(z)*Ac1(z)*Ac3(z))"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs_ex1[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)} = - i \\operatorname{A_{3}}^{2}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)} + i \\operatorname{A_{3}}{\\left(z \\right)} {d}_{3} - i \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(A3(z), z), -I*A3(z)**2*Ac3(z) + I*A3(z)*d[3] - I*Ac0(z)*Ac1(z)*Ac2(z))"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs_ex1[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{Ac_{0}}{\\left(z \\right)} = i \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{Ac_{0}}^{2}{\\left(z \\right)} + i \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - i \\operatorname{Ac_{0}}{\\left(z \\right)} {d}_{0}$"
      ],
      "text/plain": [
       "Eq(Derivative(Ac0(z), z), I*A0(z)*Ac0(z)**2 + I*A1(z)*A2(z)*A3(z) - I*Ac0(z)*d[0])"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs_ex1[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{Ac_{1}}{\\left(z \\right)} = i \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + i \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{Ac_{1}}^{2}{\\left(z \\right)} - i \\operatorname{Ac_{1}}{\\left(z \\right)} {d}_{1}$"
      ],
      "text/plain": [
       "Eq(Derivative(Ac1(z), z), I*A0(z)*A2(z)*A3(z) + I*A1(z)*Ac1(z)**2 - I*Ac1(z)*d[1])"
      ]
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs_ex1[5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{Ac_{2}}{\\left(z \\right)} = i \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} + i \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{Ac_{2}}^{2}{\\left(z \\right)} - i \\operatorname{Ac_{2}}{\\left(z \\right)} {d}_{2}$"
      ],
      "text/plain": [
       "Eq(Derivative(Ac2(z), z), I*A0(z)*A1(z)*A3(z) + I*A2(z)*Ac2(z)**2 - I*Ac2(z)*d[2])"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs_ex1[6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{Ac_{3}}{\\left(z \\right)} = i \\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} + i \\operatorname{A_{3}}{\\left(z \\right)} \\operatorname{Ac_{3}}^{2}{\\left(z \\right)} - i \\operatorname{Ac_{3}}{\\left(z \\right)} {d}_{3}$"
      ],
      "text/plain": [
       "Eq(Derivative(Ac3(z), z), I*A0(z)*A1(z)*A2(z) + I*A3(z)*Ac3(z)**2 - I*Ac3(z)*d[3])"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "euler_lagrangian_eqs_ex1[7]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "power_eqns_ex1 = [\n",
    "    Eq(diff(Uz[_i],z), diff(Az[_i]*Acz[_i],z)).expand().subs([\n",
    "        (euler_lagrangian_eqs_ex1[_i].lhs, euler_lagrangian_eqs_ex1[_i].rhs) for _i in range(2*Nmodes)\n",
    "    ]).simplify()\n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)} = i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right)$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z), I*(A0(z)*A1(z)*A2(z)*A3(z) - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z)))"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "power_eqns_ex1[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{1}}{\\left(z \\right)} = i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right)$"
      ],
      "text/plain": [
       "Eq(Derivative(U1(z), z), I*(A0(z)*A1(z)*A2(z)*A3(z) - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z)))"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "power_eqns_ex1[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{2}}{\\left(z \\right)} = i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right)$"
      ],
      "text/plain": [
       "Eq(Derivative(U2(z), z), I*(A0(z)*A1(z)*A2(z)*A3(z) - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z)))"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "power_eqns_ex1[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\operatorname{U_{3}}{\\left(z \\right)} = i \\left(\\operatorname{A_{0}}{\\left(z \\right)} \\operatorname{A_{1}}{\\left(z \\right)} \\operatorname{A_{2}}{\\left(z \\right)} \\operatorname{A_{3}}{\\left(z \\right)} - \\operatorname{Ac_{0}}{\\left(z \\right)} \\operatorname{Ac_{1}}{\\left(z \\right)} \\operatorname{Ac_{2}}{\\left(z \\right)} \\operatorname{Ac_{3}}{\\left(z \\right)}\\right)$"
      ],
      "text/plain": [
       "Eq(Derivative(U3(z), z), I*(A0(z)*A1(z)*A2(z)*A3(z) - Ac0(z)*Ac1(z)*Ac2(z)*Ac3(z)))"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "power_eqns_ex1[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "squared_power_eqs_ex1 = [\n",
    "    Eq(\n",
    "        power_eqns_ex1[_i].lhs**2, \n",
    "        (power_eqns_ex1[_i].rhs**2 + K_eq_2_ex1.rhs**2).expand()- K_eq_2_ex1.lhs**2\n",
    "    ).subs([(Acz[_j], Uz[_j]/Az[_j]) for _j in range(Nmodes)]) for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)}\\right)^{2} = - \\left(K + \\frac{\\operatorname{U_{0}}^{2}{\\left(z \\right)}}{2} - \\operatorname{U_{0}}{\\left(z \\right)} {d}_{0} + \\frac{\\operatorname{U_{1}}^{2}{\\left(z \\right)}}{2} - \\operatorname{U_{1}}{\\left(z \\right)} {d}_{1} + \\frac{\\operatorname{U_{2}}^{2}{\\left(z \\right)}}{2} - \\operatorname{U_{2}}{\\left(z \\right)} {d}_{2} + \\frac{\\operatorname{U_{3}}^{2}{\\left(z \\right)}}{2} - \\operatorname{U_{3}}{\\left(z \\right)} {d}_{3}\\right)^{2} + 4 \\operatorname{U_{0}}{\\left(z \\right)} \\operatorname{U_{1}}{\\left(z \\right)} \\operatorname{U_{2}}{\\left(z \\right)} \\operatorname{U_{3}}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z)**2, -(K + U0(z)**2/2 - U0(z)*d[0] + U1(z)**2/2 - U1(z)*d[1] + U2(z)**2/2 - U2(z)*d[2] + U3(z)**2/2 - U3(z)*d[3])**2 + 4*U0(z)*U1(z)*U2(z)*U3(z))"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squared_power_eqs_ex1[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [],
   "source": [
    "B_U_subs = [ ( B[_j, _k,], Uz[_j] - Uz[_k] ) for _j in range(Nmodes) for _k in range(Nmodes) ]\n",
    "Ac_U_subs = [(Acz[_j], Uz[_j]/Az[_j]) for _j in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "phase_mod_terms = [\n",
    "    Eq(\n",
    "        K_eq_2_ex1.lhs\n",
    "        .subs(Ac_U_subs),\n",
    "        K_eq_2_ex1.lhs\n",
    "        .subs(Ac_U_subs)\n",
    "        .subs([(Uz[_i], solve(Eq(B[_k,_i],Uz[_k]-Uz[_i]),Uz[_i])[0]) for _i in range(Nmodes) if _i != _k])\n",
    "        .expand()\n",
    "        .collect(Uz[_k],factor)\n",
    "        .subs([\n",
    "            (\n",
    "                ( K + Sum(d[j]*B[_k,j],(j,0,Nmodes-1)) + Sum(B[_k,j]**2,(j,0,Nmodes-1))/2 )\n",
    "                .doit().subs(B[_k,_k],0).factor(),\n",
    "                ( K + Sum(d[j]*B[_k,j],(j,0,Nmodes-1)) + Sum(B[_k,j]**2,(j,0,Nmodes-1))/2 )\n",
    "            ),\n",
    "            (\n",
    "                ( -Sum(d[j],(j,0,Nmodes-1)) - Sum(B[_k,j],(j,0,Nmodes-1)) )\n",
    "                .doit().subs(B[_k,_k], 0).factor(),\n",
    "                ( -Sum(d[j],(j,0,Nmodes-1)) - Sum(B[_k,j],(j,0,Nmodes-1)) )\n",
    "            )\n",
    "        ])\n",
    "    ) \n",
    "    for _k in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [],
   "source": [
    "fwm_terms = [\n",
    "    Eq(\n",
    "        4*prod([Uz[_i] for _i in range(Nmodes)]),\n",
    "        (4*prod([Uz[_i] for _i in range(Nmodes)]))\n",
    "        .subs([\n",
    "            (Uz[_i], solve(Eq(B[_k,_i],Uz[_k]-Uz[_i]),Uz[_i])[0]) for _i in range(Nmodes) if _i != _k\n",
    "        ])\n",
    "        .expand()\n",
    "        .collect(Uz[_k],factor)\n",
    "        .subs([\n",
    "            (\n",
    "                ( 2*(Sum(B[_k,j],(j,0,Nmodes-1))**2 - Sum(B[_k,j]**2,(j,0,Nmodes-1))) )\n",
    "                .doit().subs(B[_k,_k],0).factor(),\n",
    "                ( 2*(Sum(B[_k,j],(j,0,Nmodes-1))**2 - Sum(B[_k,j]**2,(j,0,Nmodes-1))) )\n",
    "            ),\n",
    "            (\n",
    "                (Sum(B[_k,j],(j,0,Nmodes-1) )).doit().subs(B[_k,_k],0).factor(),\n",
    "                Sum(B[_k,j],(j,0,Nmodes-1))\n",
    "            )\n",
    "        ])\n",
    "    )\n",
    "    for _k in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[True, True, True, True]"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[(phase_mod_terms[_j].lhs - phase_mod_terms[_j].rhs.doit().subs(B_U_subs)).simplify() == 0 for _j in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[True, True, True, True]"
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[(fwm_terms[_j].lhs - fwm_terms[_j].rhs.doit().subs(B_U_subs)).simplify() == 0 for _j in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle K + \\frac{\\operatorname{U_{0}}^{2}{\\left(z \\right)}}{2} - \\operatorname{U_{0}}{\\left(z \\right)} {d}_{0} + \\frac{\\operatorname{U_{1}}^{2}{\\left(z \\right)}}{2} - \\operatorname{U_{1}}{\\left(z \\right)} {d}_{1} + \\frac{\\operatorname{U_{2}}^{2}{\\left(z \\right)}}{2} - \\operatorname{U_{2}}{\\left(z \\right)} {d}_{2} + \\frac{\\operatorname{U_{3}}^{2}{\\left(z \\right)}}{2} - \\operatorname{U_{3}}{\\left(z \\right)} {d}_{3} = K + \\left(- \\sum_{j=0}^{3} {B}_{0,j} - \\sum_{j=0}^{3} {d}_{j}\\right) \\operatorname{U_{0}}{\\left(z \\right)} + 2 \\operatorname{U_{0}}^{2}{\\left(z \\right)} + \\sum_{j=0}^{3} {B}_{0,j} {d}_{j} + \\frac{\\sum_{j=0}^{3} {B}_{0,j}^{2}}{2}$"
      ],
      "text/plain": [
       "Eq(K + U0(z)**2/2 - U0(z)*d[0] + U1(z)**2/2 - U1(z)*d[1] + U2(z)**2/2 - U2(z)*d[2] + U3(z)**2/2 - U3(z)*d[3], K + (-Sum(B[0, j], (j, 0, 3)) - Sum(d[j], (j, 0, 3)))*U0(z) + 2*U0(z)**2 + Sum(B[0, j]*d[j], (j, 0, 3)) + Sum(B[0, j]**2, (j, 0, 3))/2)"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phase_mod_terms[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 4 \\operatorname{U_{0}}{\\left(z \\right)} \\operatorname{U_{1}}{\\left(z \\right)} \\operatorname{U_{2}}{\\left(z \\right)} \\operatorname{U_{3}}{\\left(z \\right)} = \\left(2 \\left(\\sum_{j=0}^{3} {B}_{0,j}\\right)^{2} - 2 \\sum_{j=0}^{3} {B}_{0,j}^{2}\\right) \\operatorname{U_{0}}^{2}{\\left(z \\right)} + 4 \\operatorname{U_{0}}^{4}{\\left(z \\right)} - 4 \\operatorname{U_{0}}^{3}{\\left(z \\right)} \\sum_{j=0}^{3} {B}_{0,j} - 4 \\operatorname{U_{0}}{\\left(z \\right)} {B}_{0,1} {B}_{0,2} {B}_{0,3}$"
      ],
      "text/plain": [
       "Eq(4*U0(z)*U1(z)*U2(z)*U3(z), (2*Sum(B[0, j], (j, 0, 3))**2 - 2*Sum(B[0, j]**2, (j, 0, 3)))*U0(z)**2 + 4*U0(z)**4 - 4*U0(z)**3*Sum(B[0, j], (j, 0, 3)) - 4*U0(z)*B[0, 1]*B[0, 2]*B[0, 3])"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fwm_terms[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)}\\right)^{2} = \\left(2 \\left(\\sum_{j=0}^{3} {B}_{0,j}\\right)^{2} - 2 \\sum_{j=0}^{3} {B}_{0,j}^{2}\\right) \\operatorname{U_{0}}^{2}{\\left(z \\right)} - \\left(K + \\left(- \\sum_{j=0}^{3} {B}_{0,j} - \\sum_{j=0}^{3} {d}_{j}\\right) \\operatorname{U_{0}}{\\left(z \\right)} + 2 \\operatorname{U_{0}}^{2}{\\left(z \\right)} + \\sum_{j=0}^{3} {B}_{0,j} {d}_{j} + \\frac{\\sum_{j=0}^{3} {B}_{0,j}^{2}}{2}\\right)^{2} + 4 \\operatorname{U_{0}}^{4}{\\left(z \\right)} - 4 \\operatorname{U_{0}}^{3}{\\left(z \\right)} \\sum_{j=0}^{3} {B}_{0,j} - 4 \\operatorname{U_{0}}{\\left(z \\right)} {B}_{0,1} {B}_{0,2} {B}_{0,3}$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z)**2, (2*Sum(B[0, j], (j, 0, 3))**2 - 2*Sum(B[0, j]**2, (j, 0, 3)))*U0(z)**2 - (K + (-Sum(B[0, j], (j, 0, 3)) - Sum(d[j], (j, 0, 3)))*U0(z) + 2*U0(z)**2 + Sum(B[0, j]*d[j], (j, 0, 3)) + Sum(B[0, j]**2, (j, 0, 3))/2)**2 + 4*U0(z)**4 - 4*U0(z)**3*Sum(B[0, j], (j, 0, 3)) - 4*U0(z)*B[0, 1]*B[0, 2]*B[0, 3])"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "squared_power_eqs_ex1[0].subs([\n",
    "    (phase_mod_terms[0].lhs, phase_mod_terms[0].rhs),\n",
    "    (fwm_terms[0].lhs, fwm_terms[0].rhs),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "seperated_squared_power_eqs = [\n",
    "    Eq(\n",
    "        squared_power_eqs_ex1[_k].lhs,\n",
    "        squared_power_eqs_ex1[_k].rhs.subs([\n",
    "            (phase_mod_terms[_k].lhs, phase_mod_terms[_k].rhs),\n",
    "            (fwm_terms[_k].lhs, fwm_terms[_k].rhs),\n",
    "        ]).expand().collect(Uz[_k],factor)\n",
    "    ).subs(Sum(d[j],(j,0,3)),1)\n",
    "    for _k in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "U_coeffs_reversed = [Poly(seperated_squared_power_eqs[_k].rhs,Uz[_k]).coeffs()[::-1] for _k in range(Nmodes)]\n",
    "p_vals = [[Eq(p[_k,_i], U_coeffs_reversed[_k][_i].factor()) for _i in range(Nmodes)] for _k in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {p}_{0,0} = - \\frac{\\left(2 K + 2 \\sum_{j=0}^{3} {B}_{0,j} {d}_{j} + \\sum_{j=0}^{3} {B}_{0,j}^{2}\\right)^{2}}{4}$"
      ],
      "text/plain": [
       "Eq(p[0, 0], -(2*K + 2*Sum(B[0, j]*d[j], (j, 0, 3)) + Sum(B[0, j]**2, (j, 0, 3)))**2/4)"
      ]
     },
     "execution_count": 103,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p_vals[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {p}_{0,1} = 2 K \\sum_{j=0}^{3} {B}_{0,j} + 2 K - 4 {B}_{0,1} {B}_{0,2} {B}_{0,3} + 2 \\left(\\sum_{j=0}^{3} {B}_{0,j} {d}_{j}\\right) \\sum_{j=0}^{3} {B}_{0,j} + 2 \\sum_{j=0}^{3} {B}_{0,j} {d}_{j} + \\left(\\sum_{j=0}^{3} {B}_{0,j}\\right) \\sum_{j=0}^{3} {B}_{0,j}^{2} + \\sum_{j=0}^{3} {B}_{0,j}^{2}$"
      ],
      "text/plain": [
       "Eq(p[0, 1], 2*K*Sum(B[0, j], (j, 0, 3)) + 2*K - 4*B[0, 1]*B[0, 2]*B[0, 3] + 2*Sum(B[0, j]*d[j], (j, 0, 3))*Sum(B[0, j], (j, 0, 3)) + 2*Sum(B[0, j]*d[j], (j, 0, 3)) + Sum(B[0, j], (j, 0, 3))*Sum(B[0, j]**2, (j, 0, 3)) + Sum(B[0, j]**2, (j, 0, 3)))"
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p_vals[0][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {p}_{0,2} = - 4 K - 4 \\sum_{j=0}^{3} {B}_{0,j} {d}_{j} + \\left(\\sum_{j=0}^{3} {B}_{0,j}\\right)^{2} - 2 \\sum_{j=0}^{3} {B}_{0,j} - 4 \\sum_{j=0}^{3} {B}_{0,j}^{2} - 1$"
      ],
      "text/plain": [
       "Eq(p[0, 2], -4*K - 4*Sum(B[0, j]*d[j], (j, 0, 3)) + Sum(B[0, j], (j, 0, 3))**2 - 2*Sum(B[0, j], (j, 0, 3)) - 4*Sum(B[0, j]**2, (j, 0, 3)) - 1)"
      ]
     },
     "execution_count": 105,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p_vals[0][2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {p}_{0,3} = 4$"
      ],
      "text/plain": [
       "Eq(p[0, 3], 4)"
      ]
     },
     "execution_count": 106,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p_vals[0][3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [],
   "source": [
    "seperated_U_p_eqs = [\n",
    "    seperated_squared_power_eqs[_k].subs([(p_val.rhs, p_val.lhs) for p_val in p_vals[_k]]) for _k in range(Nmodes) \n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{0}}{\\left(z \\right)}\\right)^{2} = \\operatorname{U_{0}}^{3}{\\left(z \\right)} {p}_{0,3} + \\operatorname{U_{0}}^{2}{\\left(z \\right)} {p}_{0,2} + \\operatorname{U_{0}}{\\left(z \\right)} {p}_{0,1} + {p}_{0,0}$"
      ],
      "text/plain": [
       "Eq(Derivative(U0(z), z)**2, U0(z)**3*p[0, 3] + U0(z)**2*p[0, 2] + U0(z)*p[0, 1] + p[0, 0])"
      ]
     },
     "execution_count": 108,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seperated_U_p_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{1}}{\\left(z \\right)}\\right)^{2} = \\operatorname{U_{1}}^{3}{\\left(z \\right)} {p}_{1,3} + \\operatorname{U_{1}}^{2}{\\left(z \\right)} {p}_{1,2} + \\operatorname{U_{1}}{\\left(z \\right)} {p}_{1,1} + {p}_{1,0}$"
      ],
      "text/plain": [
       "Eq(Derivative(U1(z), z)**2, U1(z)**3*p[1, 3] + U1(z)**2*p[1, 2] + U1(z)*p[1, 1] + p[1, 0])"
      ]
     },
     "execution_count": 109,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seperated_U_p_eqs[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{2}}{\\left(z \\right)}\\right)^{2} = \\operatorname{U_{2}}^{3}{\\left(z \\right)} {p}_{2,3} + \\operatorname{U_{2}}^{2}{\\left(z \\right)} {p}_{2,2} + \\operatorname{U_{2}}{\\left(z \\right)} {p}_{2,1} + {p}_{2,0}$"
      ],
      "text/plain": [
       "Eq(Derivative(U2(z), z)**2, U2(z)**3*p[2, 3] + U2(z)**2*p[2, 2] + U2(z)*p[2, 1] + p[2, 0])"
      ]
     },
     "execution_count": 110,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seperated_U_p_eqs[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} \\operatorname{U_{3}}{\\left(z \\right)}\\right)^{2} = \\operatorname{U_{3}}^{3}{\\left(z \\right)} {p}_{3,3} + \\operatorname{U_{3}}^{2}{\\left(z \\right)} {p}_{3,2} + \\operatorname{U_{3}}{\\left(z \\right)} {p}_{3,1} + {p}_{3,0}$"
      ],
      "text/plain": [
       "Eq(Derivative(U3(z), z)**2, U3(z)**3*p[3, 3] + U3(z)**2*p[3, 2] + U3(z)*p[3, 1] + p[3, 0])"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seperated_U_p_eqs[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{U_{0}}{\\left(z \\right)} = W{\\left(z \\right)} + {\\alpha}_{0}$"
      ],
      "text/plain": [
       "Eq(U0(z), W(z) + alpha[0])"
      ]
     },
     "execution_count": 112,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "U0_W_alpha = Eq(Uz[0],W(z) + alpha[0])\n",
    "U0_W_alpha"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "U_W_alpha_eqs = [Eq(Uz[_k],W(z) + alpha[_k]) for _k in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [],
   "source": [
    "W_alpha_p = [\n",
    "    Eq(seperated_U_p_eqs[_k].subs(U_W_alpha_eqs[_k].lhs, U_W_alpha_eqs[_k].rhs).lhs.simplify(), \n",
    "       seperated_U_p_eqs[_k].subs(Uz[_k],W(z) + alpha[_k]).expand().rhs.collect(W(z))\n",
    "    ).subs(p_vals[_k][3].lhs, p_vals[_k][3].rhs)\n",
    "    for _k in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} W{\\left(z \\right)}\\right)^{2} = \\left(12 {\\alpha}_{0} + {p}_{0,2}\\right) W^{2}{\\left(z \\right)} + \\left(12 {\\alpha}_{0}^{2} + 2 {\\alpha}_{0} {p}_{0,2} + {p}_{0,1}\\right) W{\\left(z \\right)} + 4 W^{3}{\\left(z \\right)} + 4 {\\alpha}_{0}^{3} + {\\alpha}_{0}^{2} {p}_{0,2} + {\\alpha}_{0} {p}_{0,1} + {p}_{0,0}$"
      ],
      "text/plain": [
       "Eq(Derivative(W(z), z)**2, (12*alpha[0] + p[0, 2])*W(z)**2 + (12*alpha[0]**2 + 2*alpha[0]*p[0, 2] + p[0, 1])*W(z) + 4*W(z)**3 + 4*alpha[0]**3 + alpha[0]**2*p[0, 2] + alpha[0]*p[0, 1] + p[0, 0])"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W_alpha_p[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} W{\\left(z \\right)}\\right)^{2} = \\left(12 {\\alpha}_{1} + {p}_{1,2}\\right) W^{2}{\\left(z \\right)} + \\left(12 {\\alpha}_{1}^{2} + 2 {\\alpha}_{1} {p}_{1,2} + {p}_{1,1}\\right) W{\\left(z \\right)} + 4 W^{3}{\\left(z \\right)} + 4 {\\alpha}_{1}^{3} + {\\alpha}_{1}^{2} {p}_{1,2} + {\\alpha}_{1} {p}_{1,1} + {p}_{1,0}$"
      ],
      "text/plain": [
       "Eq(Derivative(W(z), z)**2, (12*alpha[1] + p[1, 2])*W(z)**2 + (12*alpha[1]**2 + 2*alpha[1]*p[1, 2] + p[1, 1])*W(z) + 4*W(z)**3 + 4*alpha[1]**3 + alpha[1]**2*p[1, 2] + alpha[1]*p[1, 1] + p[1, 0])"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W_alpha_p[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha_eqs = [Eq(alpha[_k], solve(W_alpha_p[_k].rhs.coeff(W(z)**2),alpha[_k])[0]) for _k in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "W_p_eqs = [W_alpha_p[_k].subs(alpha_eqs[_k].lhs, alpha_eqs[_k].rhs) for _k in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} W{\\left(z \\right)}\\right)^{2} = \\left({p}_{0,1} - \\frac{{p}_{0,2}^{2}}{12}\\right) W{\\left(z \\right)} + 4 W^{3}{\\left(z \\right)} + {p}_{0,0} - \\frac{{p}_{0,1} {p}_{0,2}}{12} + \\frac{{p}_{0,2}^{3}}{216}$"
      ],
      "text/plain": [
       "Eq(Derivative(W(z), z)**2, (p[0, 1] - p[0, 2]**2/12)*W(z) + 4*W(z)**3 + p[0, 0] - p[0, 1]*p[0, 2]/12 + p[0, 2]**3/216)"
      ]
     },
     "execution_count": 119,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W_p_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} W{\\left(z \\right)}\\right)^{2} = \\left({p}_{1,1} - \\frac{{p}_{1,2}^{2}}{12}\\right) W{\\left(z \\right)} + 4 W^{3}{\\left(z \\right)} + {p}_{1,0} - \\frac{{p}_{1,1} {p}_{1,2}}{12} + \\frac{{p}_{1,2}^{3}}{216}$"
      ],
      "text/plain": [
       "Eq(Derivative(W(z), z)**2, (p[1, 1] - p[1, 2]**2/12)*W(z) + 4*W(z)**3 + p[1, 0] - p[1, 1]*p[1, 2]/12 + p[1, 2]**3/216)"
      ]
     },
     "execution_count": 120,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W_p_eqs[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} W{\\left(z \\right)}\\right)^{2} = \\left({p}_{2,1} - \\frac{{p}_{2,2}^{2}}{12}\\right) W{\\left(z \\right)} + 4 W^{3}{\\left(z \\right)} + {p}_{2,0} - \\frac{{p}_{2,1} {p}_{2,2}}{12} + \\frac{{p}_{2,2}^{3}}{216}$"
      ],
      "text/plain": [
       "Eq(Derivative(W(z), z)**2, (p[2, 1] - p[2, 2]**2/12)*W(z) + 4*W(z)**3 + p[2, 0] - p[2, 1]*p[2, 2]/12 + p[2, 2]**3/216)"
      ]
     },
     "execution_count": 121,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W_p_eqs[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} W{\\left(z \\right)}\\right)^{2} = \\left({p}_{3,1} - \\frac{{p}_{3,2}^{2}}{12}\\right) W{\\left(z \\right)} + 4 W^{3}{\\left(z \\right)} + {p}_{3,0} - \\frac{{p}_{3,1} {p}_{3,2}}{12} + \\frac{{p}_{3,2}^{3}}{216}$"
      ],
      "text/plain": [
       "Eq(Derivative(W(z), z)**2, (p[3, 1] - p[3, 2]**2/12)*W(z) + 4*W(z)**3 + p[3, 0] - p[3, 1]*p[3, 2]/12 + p[3, 2]**3/216)"
      ]
     },
     "execution_count": 122,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W_p_eqs[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {},
   "outputs": [],
   "source": [
    "g2_eqs = [Eq(g2,-W_p_eqs[_k].rhs.coeff(W(z))) for _k in range(Nmodes)]\n",
    "g3_eqs = [Eq(g3,-W_p_eqs[_k].rhs.subs(W(z),0)) for _k in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle g_{2} = - {p}_{0,1} + \\frac{{p}_{0,2}^{2}}{12}$"
      ],
      "text/plain": [
       "Eq(g2, -p[0, 1] + p[0, 2]**2/12)"
      ]
     },
     "execution_count": 124,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g2_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle g_{2} = - {p}_{1,1} + \\frac{{p}_{1,2}^{2}}{12}$"
      ],
      "text/plain": [
       "Eq(g2, -p[1, 1] + p[1, 2]**2/12)"
      ]
     },
     "execution_count": 125,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g2_eqs[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle g_{3} = - {p}_{0,0} + \\frac{{p}_{0,1} {p}_{0,2}}{12} - \\frac{{p}_{0,2}^{3}}{216}$"
      ],
      "text/plain": [
       "Eq(g3, -p[0, 0] + p[0, 1]*p[0, 2]/12 - p[0, 2]**3/216)"
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g3_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {},
   "outputs": [],
   "source": [
    "W_g2_g3_eqs = [\n",
    "    W_p_eqs[_k].subs([(g2_eqs[_k].rhs, g2_eqs[_k].lhs),(g3_eqs[_k].rhs, g3_eqs[_k].lhs)]) \n",
    "    for _k in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} W{\\left(z \\right)}\\right)^{2} = - g_{2} W{\\left(z \\right)} - g_{3} + 4 W^{3}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(W(z), z)**2, -g2*W(z) - g3 + 4*W(z)**3)"
      ]
     },
     "execution_count": 128,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W_g2_g3_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left(\\frac{d}{d z} W{\\left(z \\right)}\\right)^{2} = - g_{2} W{\\left(z \\right)} - g_{3} + 4 W^{3}{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(W(z), z)**2, -g2*W(z) - g3 + 4*W(z)**3)"
      ]
     },
     "execution_count": 129,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W_g2_g3_eqs[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [],
   "source": [
    "U_pw_eqs = [\n",
    "    Eq(U_W_alpha_eqs[_k].lhs, U_W_alpha_eqs[_k].rhs.subs([(alpha_eqs[_k].lhs, alpha_eqs[_k].rhs),(W(z),pw(z-z0,g2,g3))]))\n",
    "    for _k in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{U_{0}}{\\left(z \\right)} = \\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\frac{{p}_{0,2}}{12}$"
      ],
      "text/plain": [
       "Eq(U0(z), pw(z - z0, g2, g3) - p[0, 2]/12)"
      ]
     },
     "execution_count": 131,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "U_pw_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{U_{1}}{\\left(z \\right)} = \\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\frac{{p}_{1,2}}{12}$"
      ],
      "text/plain": [
       "Eq(U1(z), pw(z - z0, g2, g3) - p[1, 2]/12)"
      ]
     },
     "execution_count": 132,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "U_pw_eqs[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\alpha}_{0} = - \\frac{{p}_{0,2}}{12}$"
      ],
      "text/plain": [
       "Eq(alpha[0], -p[0, 2]/12)"
      ]
     },
     "execution_count": 137,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alpha_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "metadata": {},
   "outputs": [],
   "source": [
    "xi_p_eqs = [Eq(pw(xi[_k],g2,g3),-alpha_eqs[_k].rhs) for _k in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{pw}{\\left({\\xi}_{0},g_{2},g_{3} \\right)} = \\frac{{p}_{0,2}}{12}$"
      ],
      "text/plain": [
       "Eq(pw(xi[0], g2, g3), p[0, 2]/12)"
      ]
     },
     "execution_count": 139,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xi_p_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [],
   "source": [
    "U_pw_xi_eqs = [Eq(U_pw_eqs[_k].lhs, U_pw_eqs[_k].rhs.subs(xi_p_eqs[_k].rhs, xi_p_eqs[_k].lhs)) for _k in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{U_{0}}{\\left(z \\right)} = \\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{0},g_{2},g_{3} \\right)}$"
      ],
      "text/plain": [
       "Eq(U0(z), pw(z - z0, g2, g3) - pw(xi[0], g2, g3))"
      ]
     },
     "execution_count": 146,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "U_pw_xi_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{U_{1}}{\\left(z \\right)} = \\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{1},g_{2},g_{3} \\right)}$"
      ],
      "text/plain": [
       "Eq(U1(z), pw(z - z0, g2, g3) - pw(xi[1], g2, g3))"
      ]
     },
     "execution_count": 147,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "U_pw_xi_eqs[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{U_{2}}{\\left(z \\right)} = \\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{2},g_{2},g_{3} \\right)}$"
      ],
      "text/plain": [
       "Eq(U2(z), pw(z - z0, g2, g3) - pw(xi[2], g2, g3))"
      ]
     },
     "execution_count": 148,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "U_pw_xi_eqs[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{U_{3}}{\\left(z \\right)} = \\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{3},g_{2},g_{3} \\right)}$"
      ],
      "text/plain": [
       "Eq(U3(z), pw(z - z0, g2, g3) - pw(xi[3], g2, g3))"
      ]
     },
     "execution_count": 149,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "U_pw_xi_eqs[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [],
   "source": [
    "xi_pwp_sqrd_eqs = [\n",
    "    seperated_U_p_eqs[_i]\n",
    "    .subs(U_pw_xi_eqs[_i].lhs, U_pw_xi_eqs[_i].rhs)\n",
    "    .subs(z,z0+xi[_i])\n",
    "    .doit()\n",
    "    .subs(diff(pw(z,g2,g3),z).subs(z,xi[_i]), pwp(xi[_i],g2,g3))\n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{pwp}^{2}{\\left({\\xi}_{0},g_{2},g_{3} \\right)} = {p}_{0,0}$"
      ],
      "text/plain": [
       "Eq(pwp(xi[0], g2, g3)**2, p[0, 0])"
      ]
     },
     "execution_count": 151,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xi_pwp_sqrd_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 292,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwp_xi_eqs = [Eq(pwp(xi[_i],g2,g3), -I*phase_mod_terms[_i].rhs.subs(Uz[_i],0)) for _i in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 293,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{pwp}{\\left({\\xi}_{0},g_{2},g_{3} \\right)} = - i \\left(K + \\sum_{j=0}^{3} {B}_{0,j} {d}_{j} + \\frac{\\sum_{j=0}^{3} {B}_{0,j}^{2}}{2}\\right)$"
      ],
      "text/plain": [
       "Eq(pwp(xi[0], g2, g3), -I*(K + Sum(B[0, j]*d[j], (j, 0, 3)) + Sum(B[0, j]**2, (j, 0, 3))/2))"
      ]
     },
     "execution_count": 293,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pwp_xi_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 294,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[True, True, True, True]"
      ]
     },
     "execution_count": 294,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pwp_xi_checks = [\n",
    "    (xi_pwp_sqrd_eqs[_i].subs(p_vals[_i][0].lhs, p_vals[_i][0].rhs).rhs - pwp_xi_eqs[_i].rhs**2).simplify() == 0 \n",
    "     for _i in range(Nmodes)\n",
    "]\n",
    "pwp_xi_checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[True, True, True, True, True, True, True, True, True, True, True, True]"
      ]
     },
     "execution_count": 152,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g2_all_equal_checks = [\n",
    "    (g2_eqs[_jj].rhs - g2_eqs[_kk].rhs).subs([\n",
    "        (p_vals[_j][_k].lhs, p_vals[_j][_k].rhs) for _j in range(Nmodes) for _k in range(Nmodes)\n",
    "    ]).expand().doit().subs(d[3], 1 - d[0] - d[1] - d[2]).subs(B_U_subs).expand() == 0\n",
    "    for _jj in range(Nmodes) for _kk in range(Nmodes) if _jj != _kk\n",
    "]\n",
    "g2_all_equal_checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[True, True, True, True, True, True, True, True, True, True, True, True]"
      ]
     },
     "execution_count": 153,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g3_all_equal_checks = [\n",
    "    (g3_eqs[_jj].rhs - g3_eqs[_kk].rhs).subs([\n",
    "    (p_vals[_j][_k].lhs, p_vals[_j][_k].rhs) for _j in range(Nmodes) for _k in range(Nmodes)\n",
    "    ]).expand().doit().subs(d[3], 1 - d[0] - d[1] - d[2]).subs(B_U_subs).expand() == 0\n",
    "    for _jj in range(Nmodes) for _kk in range(Nmodes) if _jj != _kk\n",
    "]\n",
    "g3_all_equal_checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {},
   "outputs": [],
   "source": [
    "g2_avg_eq = Eq(g2, sum([\n",
    "        g2_eqs[_jj].rhs.subs([\n",
    "            (p_vals[_j][_k].lhs, p_vals[_j][_k].rhs) for _j in range(Nmodes) for _k in range(Nmodes)\n",
    "        ])/Nmodes\n",
    "        for _jj in range(Nmodes) for _kk in range(Nmodes) if _jj != _kk\n",
    "    ]).expand().collect(K,factor)\n",
    "  )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Logarithmic Derivatives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 295,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_eq_2_ex1_b = Eq(K_eq_2_ex1.lhs + P_ex1.rhs -  P_ex1.lhs, K_eq_2_ex1.rhs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 296,
   "metadata": {},
   "outputs": [],
   "source": [
    "phase_eqs_ex1 = [\n",
    "    Eq(\n",
    "        diff(phiz[_i],z), \n",
    "        diff(-I*ln(Az[_i]/Acz[_i])/2,z).simplify().subs([\n",
    "            (euler_lagrangian_eqs_ex1[_i].lhs, euler_lagrangian_eqs_ex1[_i].rhs) for _i in range(2*Nmodes)\n",
    "        ])\n",
    "        .simplify()\n",
    "        - (\n",
    "            K_eq_2_ex1_b.rhs/(Az[_i]*Acz[_i]) - K_eq_2_ex1_b.lhs/(Az[_i]*Acz[_i])\n",
    "        ).expand()/2\n",
    "      )\n",
    "    .expand().subs([(Acz[_i], Uz[_i]/Az[_i]) for _i in range(Nmodes)]) \n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 297,
   "metadata": {},
   "outputs": [],
   "source": [
    "P_U_ex1 = P_ex1.subs([(Acz[_j], Uz[_j]/Az[_j]) for _j in range(Nmodes)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 298,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho_eqs = [Eq(rho[_i], d[_i] - Sum(B[_i,j],(j,0,Nmodes - 1))/2 - Rational(1,2)) for _i in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 299,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\rho}_{0} = {d}_{0} - \\frac{\\sum_{j=0}^{3} {B}_{0,j}}{2} - \\frac{1}{2}$"
      ],
      "text/plain": [
       "Eq(rho[0], d[0] - Sum(B[0, j], (j, 0, 3))/2 - 1/2)"
      ]
     },
     "execution_count": 299,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rho_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 591,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sum_{j=0}^{3} {\\rho}_{j} = -1$"
      ],
      "text/plain": [
       "Eq(Sum(rho[j], (j, 0, 3)), -1)"
      ]
     },
     "execution_count": 591,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(\n",
    "    Sum(rho[j],(j,0,Nmodes -1)), \n",
    "    Sum(rho[j],(j,0,Nmodes -1))\n",
    "    .doit()\n",
    "    .subs([r.args for r in rho_eqs])\n",
    "    .doit()\n",
    "    .subs(B_U_subs)\n",
    "    .subs(Sum(d[j],(j,0,Nmodes -1)).doit(), 1)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 300,
   "metadata": {},
   "outputs": [],
   "source": [
    "phase_eqs_ex1_b = [\n",
    "    Eq(phase_eqs_ex1[_i].lhs,\n",
    "        phase_eqs_ex1[_i]\n",
    "        .rhs.subs(P_U_ex1.lhs, P_U_ex1.rhs)\n",
    "        .subs(phase_mod_terms[_i].lhs - K, phase_mod_terms[_i].rhs - K)\n",
    "        .expand()\n",
    "        .collect(Uz[_i])\n",
    "       .subs(Sum(d[j],(j,0,3)),1)\n",
    "       .subs(-I*pwp_xi_eqs[_i].rhs/2, -I*pwp_xi_eqs[_i].lhs/2)\n",
    "       .subs(U_pw_xi_eqs[_i].lhs, U_pw_xi_eqs[_i].rhs)\n",
    "       .subs(rho_eqs[_i].rhs, rho_eqs[_i].lhs)\n",
    "    )\n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 301,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\phi_{0}{\\left(z \\right)} = {\\rho}_{0} + \\frac{i \\operatorname{pwp}{\\left({\\xi}_{0},g_{2},g_{3} \\right)}}{2 \\left(\\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{0},g_{2},g_{3} \\right)}\\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(phi0(z), z), rho[0] + I*pwp(xi[0], g2, g3)/(2*(pw(z - z0, g2, g3) - pw(xi[0], g2, g3))))"
      ]
     },
     "execution_count": 301,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phase_eqs_ex1_b[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 302,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\phi_{1}{\\left(z \\right)} = {\\rho}_{1} + \\frac{i \\operatorname{pwp}{\\left({\\xi}_{1},g_{2},g_{3} \\right)}}{2 \\left(\\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{1},g_{2},g_{3} \\right)}\\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(phi1(z), z), rho[1] + I*pwp(xi[1], g2, g3)/(2*(pw(z - z0, g2, g3) - pw(xi[1], g2, g3))))"
      ]
     },
     "execution_count": 302,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phase_eqs_ex1_b[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 303,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\phi_{2}{\\left(z \\right)} = {\\rho}_{2} + \\frac{i \\operatorname{pwp}{\\left({\\xi}_{2},g_{2},g_{3} \\right)}}{2 \\left(\\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{2},g_{2},g_{3} \\right)}\\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(phi2(z), z), rho[2] + I*pwp(xi[2], g2, g3)/(2*(pw(z - z0, g2, g3) - pw(xi[2], g2, g3))))"
      ]
     },
     "execution_count": 303,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phase_eqs_ex1_b[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 304,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{d}{d z} \\phi_{3}{\\left(z \\right)} = {\\rho}_{3} + \\frac{i \\operatorname{pwp}{\\left({\\xi}_{3},g_{2},g_{3} \\right)}}{2 \\left(\\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{3},g_{2},g_{3} \\right)}\\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(phi3(z), z), rho[3] + I*pwp(xi[3], g2, g3)/(2*(pw(z - z0, g2, g3) - pw(xi[3], g2, g3))))"
      ]
     },
     "execution_count": 304,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phase_eqs_ex1_b[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 305,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_diff_eqns_ex1 = [\n",
    "    Eq(\n",
    "        diff(Az[_i],z)/Az[_i], \n",
    "        (I*diff(phiz[_i],z) + diff(Uz[_i],z)/Uz[_i]/2)\n",
    "        .subs(phase_eqs_ex1_b[_i].lhs, phase_eqs_ex1_b[_i].rhs)\n",
    "        .expand()\n",
    "        .subs(U_pw_xi_eqs[_i].lhs, U_pw_xi_eqs[_i].rhs)\n",
    "        .subs(Derivative(pw(z-z0,g2,g3) - pw(xi[_i],g2,g3),z),pwp(z-z0,g2,g3))\n",
    "        .collect(rho[_i], simplify)\n",
    "    )\n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 306,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)}}{\\operatorname{A_{0}}{\\left(z \\right)}} = i {\\rho}_{0} + \\frac{\\operatorname{pwp}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pwp}{\\left({\\xi}_{0},g_{2},g_{3} \\right)}}{2 \\left(\\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{0},g_{2},g_{3} \\right)}\\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(A0(z), z)/A0(z), I*rho[0] + (pwp(z - z0, g2, g3) - pwp(xi[0], g2, g3))/(2*(pw(z - z0, g2, g3) - pw(xi[0], g2, g3))))"
      ]
     },
     "execution_count": 306,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_eqns_ex1[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 307,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)}}{\\operatorname{A_{1}}{\\left(z \\right)}} = i {\\rho}_{1} + \\frac{\\operatorname{pwp}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pwp}{\\left({\\xi}_{1},g_{2},g_{3} \\right)}}{2 \\left(\\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{1},g_{2},g_{3} \\right)}\\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(A1(z), z)/A1(z), I*rho[1] + (pwp(z - z0, g2, g3) - pwp(xi[1], g2, g3))/(2*(pw(z - z0, g2, g3) - pw(xi[1], g2, g3))))"
      ]
     },
     "execution_count": 307,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_eqns_ex1[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 308,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)}}{\\operatorname{A_{2}}{\\left(z \\right)}} = i {\\rho}_{2} + \\frac{\\operatorname{pwp}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pwp}{\\left({\\xi}_{2},g_{2},g_{3} \\right)}}{2 \\left(\\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{2},g_{2},g_{3} \\right)}\\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(A2(z), z)/A2(z), I*rho[2] + (pwp(z - z0, g2, g3) - pwp(xi[2], g2, g3))/(2*(pw(z - z0, g2, g3) - pw(xi[2], g2, g3))))"
      ]
     },
     "execution_count": 308,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_eqns_ex1[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 309,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)}}{\\operatorname{A_{3}}{\\left(z \\right)}} = i {\\rho}_{3} + \\frac{\\operatorname{pwp}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pwp}{\\left({\\xi}_{3},g_{2},g_{3} \\right)}}{2 \\left(\\operatorname{pw}{\\left(z - z_{0},g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left({\\xi}_{3},g_{2},g_{3} \\right)}\\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(A3(z), z)/A3(z), I*rho[3] + (pwp(z - z0, g2, g3) - pwp(xi[3], g2, g3))/(2*(pw(z - z0, g2, g3) - pw(xi[3], g2, g3))))"
      ]
     },
     "execution_count": 309,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_eqns_ex1[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 310,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\operatorname{pwp}{\\left(x,g_{2},g_{3} \\right)} - \\operatorname{pwp}{\\left(y,g_{2},g_{3} \\right)}}{2 \\left(\\operatorname{pw}{\\left(x,g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left(y,g_{2},g_{3} \\right)}\\right)} = - \\operatorname{zw}{\\left(x,g_{2},g_{3} \\right)} - \\operatorname{zw}{\\left(y,g_{2},g_{3} \\right)} + \\operatorname{zw}{\\left(x + y,g_{2},g_{3} \\right)}$"
      ],
      "text/plain": [
       "Eq((pwp(x, g2, g3) - pwp(y, g2, g3))/(2*(pw(x, g2, g3) - pw(y, g2, g3))), -zw(x, g2, g3) - zw(y, g2, g3) + zw(x + y, g2, g3))"
      ]
     },
     "execution_count": 310,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pw_to_zw_identity = Eq(\n",
    "    (pwp(x,g2,g3) - pwp(y,g2,g3))/(pw(x,g2,g3) - pw(y,g2,g3))/2,\n",
    "    zw(x + y,g2, g3) - zw(x,g2, g3) - zw(y,g2, g3)\n",
    ")\n",
    "pw_to_zw_identity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pw_to_zw_1 = pw_to_zw_identity.subs([(x,xi1),(y,z-z0)])\n",
    "log_diff_A1_zw_xi = log_diff_A1_pw_xi.subs(pw_to_zw_1.lhs, pw_to_zw_1.rhs)\n",
    "log_diff_A1_zw_xi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 323,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_diff_zeta_eqns = [\n",
    "    log_diff_eqns_ex1[_i]\n",
    "    .subs(\n",
    "        pw_to_zw_identity.subs([(x,xi[_i]),(y,z-z0)]).lhs.simplify(), \n",
    "        pw_to_zw_identity.subs([(x,xi[_i]),(y,z-z0)]).rhs\n",
    "    )\n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 324,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{0}}{\\left(z \\right)}}{\\operatorname{A_{0}}{\\left(z \\right)}} = - \\operatorname{zw}{\\left(z - z_{0},g_{2},g_{3} \\right)} + \\operatorname{zw}{\\left(z - z_{0} + {\\xi}_{0},g_{2},g_{3} \\right)} - \\operatorname{zw}{\\left({\\xi}_{0},g_{2},g_{3} \\right)} + i {\\rho}_{0}$"
      ],
      "text/plain": [
       "Eq(Derivative(A0(z), z)/A0(z), -zw(z - z0, g2, g3) + zw(z - z0 + xi[0], g2, g3) - zw(xi[0], g2, g3) + I*rho[0])"
      ]
     },
     "execution_count": 324,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_zeta_eqns[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 325,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{1}}{\\left(z \\right)}}{\\operatorname{A_{1}}{\\left(z \\right)}} = - \\operatorname{zw}{\\left(z - z_{0},g_{2},g_{3} \\right)} + \\operatorname{zw}{\\left(z - z_{0} + {\\xi}_{1},g_{2},g_{3} \\right)} - \\operatorname{zw}{\\left({\\xi}_{1},g_{2},g_{3} \\right)} + i {\\rho}_{1}$"
      ],
      "text/plain": [
       "Eq(Derivative(A1(z), z)/A1(z), -zw(z - z0, g2, g3) + zw(z - z0 + xi[1], g2, g3) - zw(xi[1], g2, g3) + I*rho[1])"
      ]
     },
     "execution_count": 325,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_zeta_eqns[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 326,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{2}}{\\left(z \\right)}}{\\operatorname{A_{2}}{\\left(z \\right)}} = - \\operatorname{zw}{\\left(z - z_{0},g_{2},g_{3} \\right)} + \\operatorname{zw}{\\left(z - z_{0} + {\\xi}_{2},g_{2},g_{3} \\right)} - \\operatorname{zw}{\\left({\\xi}_{2},g_{2},g_{3} \\right)} + i {\\rho}_{2}$"
      ],
      "text/plain": [
       "Eq(Derivative(A2(z), z)/A2(z), -zw(z - z0, g2, g3) + zw(z - z0 + xi[2], g2, g3) - zw(xi[2], g2, g3) + I*rho[2])"
      ]
     },
     "execution_count": 326,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_zeta_eqns[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 327,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{d}{d z} \\operatorname{A_{3}}{\\left(z \\right)}}{\\operatorname{A_{3}}{\\left(z \\right)}} = - \\operatorname{zw}{\\left(z - z_{0},g_{2},g_{3} \\right)} + \\operatorname{zw}{\\left(z - z_{0} + {\\xi}_{3},g_{2},g_{3} \\right)} - \\operatorname{zw}{\\left({\\xi}_{3},g_{2},g_{3} \\right)} + i {\\rho}_{3}$"
      ],
      "text/plain": [
       "Eq(Derivative(A3(z), z)/A3(z), -zw(z - z0, g2, g3) + zw(z - z0 + xi[3], g2, g3) - zw(xi[3], g2, g3) + I*rho[3])"
      ]
     },
     "execution_count": 327,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_zeta_eqns[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 329,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\int \\frac{\\frac{d}{d z} f{\\left(z \\right)}}{f{\\left(z \\right)}}\\, dz = \\log{\\left(f{\\left(z \\right)} \\right)}$"
      ],
      "text/plain": [
       "Eq(Integral(Derivative(f(z), z)/f(z), z), log(f(z)))"
      ]
     },
     "execution_count": 329,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(Integral(Derivative(f(z),z)/f(z)),log(f(z)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'Eq' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-337a7c0bb7fb>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mintegral_zeta_log_sigma\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mEq\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mIntegral\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mzw\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mz\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mg2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mg3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mz\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mlog\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msigma\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mz\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mg2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mg3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mintegral_zeta_log_sigma\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'Eq' is not defined"
     ]
    }
   ],
   "source": [
    "integral_zeta_log_sigma = Eq(Integral(zw(z,g2,g3),z),log(sigma(z,g2,g3)))\n",
    "integral_zeta_log_sigma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 371,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial z} \\log{\\left(\\sigma{\\left(- x + z,g_{2},g_{3} \\right)} \\right)} = \\operatorname{zw}{\\left(- x + z,g_{2},g_{3} \\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(log(sigma(-x + z, g2, g3)), z), zw(-x + z, g2, g3))"
      ]
     },
     "execution_count": 371,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_diff_sigma_zeta = Eq(Derivative(log(sigma(z-x,g2,g3)),z), zw(z-x,g2,g3))\n",
    "log_diff_sigma_zeta_minus = Eq(Derivative(-log(sigma(z-x,g2,g3)),z), -zw(z-x,g2,g3))\n",
    "log_diff_sigma_zeta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 449,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{pw}{\\left(x,g_{2},g_{3} \\right)} - \\operatorname{pw}{\\left(y,g_{2},g_{3} \\right)} = - \\frac{\\sigma{\\left(x - y,g_{2},g_{3} \\right)} \\sigma{\\left(x + y,g_{2},g_{3} \\right)}}{\\sigma^{2}{\\left(x,g_{2},g_{3} \\right)} \\sigma^{2}{\\left(y,g_{2},g_{3} \\right)}}$"
      ],
      "text/plain": [
       "Eq(pw(x, g2, g3) - pw(y, g2, g3), -sigma(x - y, g2, g3)*sigma(x + y, g2, g3)/(sigma(x, g2, g3)**2*sigma(y, g2, g3)**2))"
      ]
     },
     "execution_count": 449,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sigma_p_identity = Eq(\n",
    "    pw(x, g2, g3) - pw(y, g2, g3),\n",
    "    -sigma(x + y, g2, g3) * sigma(x - y, g2, g3) / (sigma(x, g2, g3) ** 2 * sigma(y, g2, g3) ** 2) \n",
    ")\n",
    "sigma_p_identity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 429,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_A_sigma_eqs = [\n",
    "    Eq(\n",
    "        log(Az[_i]), \n",
    "        log(c[_i]) + (I*rho[_i] - zw(xi[_i],g2,g3))*z + log(sigma(xi[_i] + z - z0, g2, g3)) - log(sigma(z - z0, g2, g3))\n",
    "    )\n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 430,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_results_differentiated = [\n",
    "    Eq(\n",
    "        diff(log_A_sigma_eqs[_i].lhs, z), \n",
    "        sum(Derivative(_term, z) for _term in Add.make_args(log_A_sigma_eqs[_i].rhs))\n",
    "    ).subs([\n",
    "        log_diff_sigma_zeta.subs(x,z0-xi[_i]).args,\n",
    "        log_diff_sigma_zeta_minus.subs(x,z0).args \n",
    "    ]).doit()\n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 431,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[True, True, True, True]"
      ]
     },
     "execution_count": 431,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sigma_results_tests = [\n",
    "    (log_diff_zeta_eqns[_i].lhs.subs(*sigma_results_differentiated[_i].args) - log_diff_zeta_eqns[_i].rhs) == 0\n",
    "    for _i in range(Nmodes)\n",
    "]   \n",
    "sigma_results_tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 432,
   "metadata": {},
   "outputs": [],
   "source": [
    "A_sigma_results = [\n",
    "    Eq(\n",
    "        Az[_i], \n",
    "        c[_i]*exp((I*rho[_i] - zw(xi[_i],g2,g3))*z)*sigma(xi[_i] + z - z0, g2, g3)/sigma(z - z0, g2, g3)\n",
    "    )\n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 450,
   "metadata": {},
   "outputs": [],
   "source": [
    "A_sigma_results_conj = [\n",
    "    Eq(U_pw_xi_eqs[_i].lhs/Az[_i], U_pw_xi_eqs[_i].rhs/Az[_i])\n",
    "    .subs([(Uz[_i],Az[_i]*Acz[_i])])\n",
    "    .subs(*sigma_p_identity.subs([(x,z-z0),(y,xi[_i])]).args)\n",
    "    .subs(*A_sigma_results[_i].args)\n",
    "    for _i in range(Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 451,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{A_{0}}{\\left(z \\right)} = \\frac{\\sigma{\\left(z - z_{0} + {\\xi}_{0},g_{2},g_{3} \\right)} e^{z \\left(- \\operatorname{zw}{\\left({\\xi}_{0},g_{2},g_{3} \\right)} + i {\\rho}_{0}\\right)} {c}_{0}}{\\sigma{\\left(z - z_{0},g_{2},g_{3} \\right)}}$"
      ],
      "text/plain": [
       "Eq(A0(z), sigma(z - z0 + xi[0], g2, g3)*exp(z*(-zw(xi[0], g2, g3) + I*rho[0]))*c[0]/sigma(z - z0, g2, g3))"
      ]
     },
     "execution_count": 451,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_sigma_results[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 452,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{A_{1}}{\\left(z \\right)} = \\frac{\\sigma{\\left(z - z_{0} + {\\xi}_{1},g_{2},g_{3} \\right)} e^{z \\left(- \\operatorname{zw}{\\left({\\xi}_{1},g_{2},g_{3} \\right)} + i {\\rho}_{1}\\right)} {c}_{1}}{\\sigma{\\left(z - z_{0},g_{2},g_{3} \\right)}}$"
      ],
      "text/plain": [
       "Eq(A1(z), sigma(z - z0 + xi[1], g2, g3)*exp(z*(-zw(xi[1], g2, g3) + I*rho[1]))*c[1]/sigma(z - z0, g2, g3))"
      ]
     },
     "execution_count": 452,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_sigma_results[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 453,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{A_{2}}{\\left(z \\right)} = \\frac{\\sigma{\\left(z - z_{0} + {\\xi}_{2},g_{2},g_{3} \\right)} e^{z \\left(- \\operatorname{zw}{\\left({\\xi}_{2},g_{2},g_{3} \\right)} + i {\\rho}_{2}\\right)} {c}_{2}}{\\sigma{\\left(z - z_{0},g_{2},g_{3} \\right)}}$"
      ],
      "text/plain": [
       "Eq(A2(z), sigma(z - z0 + xi[2], g2, g3)*exp(z*(-zw(xi[2], g2, g3) + I*rho[2]))*c[2]/sigma(z - z0, g2, g3))"
      ]
     },
     "execution_count": 453,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_sigma_results[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 454,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{A_{3}}{\\left(z \\right)} = \\frac{\\sigma{\\left(z - z_{0} + {\\xi}_{3},g_{2},g_{3} \\right)} e^{z \\left(- \\operatorname{zw}{\\left({\\xi}_{3},g_{2},g_{3} \\right)} + i {\\rho}_{3}\\right)} {c}_{3}}{\\sigma{\\left(z - z_{0},g_{2},g_{3} \\right)}}$"
      ],
      "text/plain": [
       "Eq(A3(z), sigma(z - z0 + xi[3], g2, g3)*exp(z*(-zw(xi[3], g2, g3) + I*rho[3]))*c[3]/sigma(z - z0, g2, g3))"
      ]
     },
     "execution_count": 454,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_sigma_results[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 455,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{Ac_{0}}{\\left(z \\right)} = - \\frac{\\sigma{\\left(z - z_{0} - {\\xi}_{0},g_{2},g_{3} \\right)} e^{- z \\left(- \\operatorname{zw}{\\left({\\xi}_{0},g_{2},g_{3} \\right)} + i {\\rho}_{0}\\right)}}{\\sigma{\\left(z - z_{0},g_{2},g_{3} \\right)} \\sigma^{2}{\\left({\\xi}_{0},g_{2},g_{3} \\right)} {c}_{0}}$"
      ],
      "text/plain": [
       "Eq(Ac0(z), -sigma(z - z0 - xi[0], g2, g3)*exp(-z*(-zw(xi[0], g2, g3) + I*rho[0]))/(sigma(z - z0, g2, g3)*sigma(xi[0], g2, g3)**2*c[0]))"
      ]
     },
     "execution_count": 455,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_sigma_results_conj[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 456,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{Ac_{1}}{\\left(z \\right)} = - \\frac{\\sigma{\\left(z - z_{0} - {\\xi}_{1},g_{2},g_{3} \\right)} e^{- z \\left(- \\operatorname{zw}{\\left({\\xi}_{1},g_{2},g_{3} \\right)} + i {\\rho}_{1}\\right)}}{\\sigma{\\left(z - z_{0},g_{2},g_{3} \\right)} \\sigma^{2}{\\left({\\xi}_{1},g_{2},g_{3} \\right)} {c}_{1}}$"
      ],
      "text/plain": [
       "Eq(Ac1(z), -sigma(z - z0 - xi[1], g2, g3)*exp(-z*(-zw(xi[1], g2, g3) + I*rho[1]))/(sigma(z - z0, g2, g3)*sigma(xi[1], g2, g3)**2*c[1]))"
      ]
     },
     "execution_count": 456,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_sigma_results_conj[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 457,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{Ac_{2}}{\\left(z \\right)} = - \\frac{\\sigma{\\left(z - z_{0} - {\\xi}_{2},g_{2},g_{3} \\right)} e^{- z \\left(- \\operatorname{zw}{\\left({\\xi}_{2},g_{2},g_{3} \\right)} + i {\\rho}_{2}\\right)}}{\\sigma{\\left(z - z_{0},g_{2},g_{3} \\right)} \\sigma^{2}{\\left({\\xi}_{2},g_{2},g_{3} \\right)} {c}_{2}}$"
      ],
      "text/plain": [
       "Eq(Ac2(z), -sigma(z - z0 - xi[2], g2, g3)*exp(-z*(-zw(xi[2], g2, g3) + I*rho[2]))/(sigma(z - z0, g2, g3)*sigma(xi[2], g2, g3)**2*c[2]))"
      ]
     },
     "execution_count": 457,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_sigma_results_conj[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 458,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\operatorname{Ac_{3}}{\\left(z \\right)} = - \\frac{\\sigma{\\left(z - z_{0} - {\\xi}_{3},g_{2},g_{3} \\right)} e^{- z \\left(- \\operatorname{zw}{\\left({\\xi}_{3},g_{2},g_{3} \\right)} + i {\\rho}_{3}\\right)}}{\\sigma{\\left(z - z_{0},g_{2},g_{3} \\right)} \\sigma^{2}{\\left({\\xi}_{3},g_{2},g_{3} \\right)} {c}_{3}}$"
      ],
      "text/plain": [
       "Eq(Ac3(z), -sigma(z - z0 - xi[3], g2, g3)*exp(-z*(-zw(xi[3], g2, g3) + I*rho[3]))/(sigma(z - z0, g2, g3)*sigma(xi[3], g2, g3)**2*c[3]))"
      ]
     },
     "execution_count": 458,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A_sigma_results_conj[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 532,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle h{\\left(z \\right)} = \\prod_{j=0}^{3} \\frac{e^{z \\left(\\operatorname{zw}{\\left({\\xi}_{j},g_{2},g_{3} \\right)} - i {\\rho}_{j}\\right)}}{\\sigma{\\left({\\xi}_{j},g_{2},g_{3} \\right)} {c}_{j}}$"
      ],
      "text/plain": [
       "Eq(h(z), Product(exp(z*(zw(xi[j], g2, g3) - I*rho[j]))/(sigma(xi[j], g2, g3)*c[j]), (j, 0, 3)))"
      ]
     },
     "execution_count": 532,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hz_eq = Eq(h(z), Product(exp(z*(zw(xi[j],g2,g3) - I*rho[j]))/(c[j]*sigma(xi[j],g2,g3)),(j,0,Nmodes - 1)))\n",
    "hz_eq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 544,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {c}_{0} {c}_{1} {c}_{2} {c}_{3} = \\frac{e^{z \\left(\\operatorname{zw}{\\left({\\xi}_{0},g_{2},g_{3} \\right)} + \\operatorname{zw}{\\left({\\xi}_{1},g_{2},g_{3} \\right)} + \\operatorname{zw}{\\left({\\xi}_{2},g_{2},g_{3} \\right)} + \\operatorname{zw}{\\left({\\xi}_{3},g_{2},g_{3} \\right)} - i {\\rho}_{0} - i {\\rho}_{1} - i {\\rho}_{2} - i {\\rho}_{3}\\right)}}{h{\\left(z \\right)} \\sigma{\\left({\\xi}_{0},g_{2},g_{3} \\right)} \\sigma{\\left({\\xi}_{1},g_{2},g_{3} \\right)} \\sigma{\\left({\\xi}_{2},g_{2},g_{3} \\right)} \\sigma{\\left({\\xi}_{3},g_{2},g_{3} \\right)}}$"
      ],
      "text/plain": [
       "Eq(c[0]*c[1]*c[2]*c[3], exp(z*(zw(xi[0], g2, g3) + zw(xi[1], g2, g3) + zw(xi[2], g2, g3) + zw(xi[3], g2, g3) - I*rho[0] - I*rho[1] - I*rho[2] - I*rho[3]))/(h(z)*sigma(xi[0], g2, g3)*sigma(xi[1], g2, g3)*sigma(xi[2], g2, g3)*sigma(xi[3], g2, g3)))"
      ]
     },
     "execution_count": 544,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c_prod_h = Eq(hz_eq.doit().simplify().lhs*Product(c[j],(j,0,Nmodes - 1)).doit()/h(z),\n",
    "                 hz_eq.doit().simplify().rhs*Product(c[j],(j,0,Nmodes - 1)).doit()/h(z)\n",
    "                )\n",
    "c_prod_h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 567,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_eqs = [Eq(delta[_i],zw(xi[_i],g2,g3) + I*d[_i] - I*rho[_i]) for _i in range(Nmodes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 568,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\delta}_{0} = \\operatorname{zw}{\\left({\\xi}_{0},g_{2},g_{3} \\right)} + i {d}_{0} - i {\\rho}_{0}$"
      ],
      "text/plain": [
       "Eq(delta[0], zw(xi[0], g2, g3) + I*d[0] - I*rho[0])"
      ]
     },
     "execution_count": 568,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "delta_eqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 577,
   "metadata": {},
   "outputs": [],
   "source": [
    "xis_plus_minus = [xi[_k] for _k in range(Nmodes)] + [-xi[_k] for _k in range(Nmodes)]\n",
    "f_diff_eqs_a = [\n",
    "    el_eq.subs([\n",
    "        eq.args for eq in A_sigma_results + A_sigma_results_conj\n",
    "    ]).subs([\n",
    "        (sigma(z - z0 + _xi, g2, g3), f(z,_xi)*sigma(z - z0, g2, g3)*sigma(_xi, g2, g3)) \n",
    "        for _xi in xis_plus_minus\n",
    "    ]).doit()\n",
    "    for el_eq in euler_lagrangian_eqs_ex1\n",
    "]\n",
    "f_diff_eqs_b = [\n",
    "    Eq(\n",
    "        diff(f(z, xis_plus_minus[_i]),z), \n",
    "        solve(f_diff_eqs_a[_i].lhs - f_diff_eqs_a[_i].rhs, diff(f(z, xis_plus_minus[_i]),z))[0]\n",
    "        .subs([\n",
    "            (sigma(-xi[_j], g2, g3), -sigma(xi[_j], g2, g3))  for _j in range(Nmodes)\n",
    "        ])\n",
    "        .simplify()\n",
    "        .expand()\n",
    "        .subs(*c_prod_h.args)\n",
    "        .collect(f(z, xis_plus_minus[_i]), simplify)\n",
    "        .subs([eq.args[::-1] for eq in delta_eqs])\n",
    "    )\n",
    "    for _i in range(2*Nmodes)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 578,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial z} f{\\left(z,{\\xi}_{0} \\right)} = - i f{\\left(z,- {\\xi}_{0} \\right)} f^{2}{\\left(z,{\\xi}_{0} \\right)} + f{\\left(z,{\\xi}_{0} \\right)} {\\delta}_{0} - i f{\\left(z,- {\\xi}_{1} \\right)} f{\\left(z,- {\\xi}_{2} \\right)} f{\\left(z,- {\\xi}_{3} \\right)} h{\\left(z \\right)}$"
      ],
      "text/plain": [
       "Eq(Derivative(f(z, xi[0]), z), -I*f(z, -xi[0])*f(z, xi[0])**2 + f(z, xi[0])*delta[0] - I*f(z, -xi[1])*f(z, -xi[2])*f(z, -xi[3])*h(z))"
      ]
     },
     "execution_count": 578,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_diff_eqs_b[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 579,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial z} f{\\left(z,{\\xi}_{1} \\right)} = - i f{\\left(z,- {\\xi}_{0} \\right)} f{\\left(z,- {\\xi}_{2} \\right)} f{\\left(z,- {\\xi}_{3} \\right)} h{\\left(z \\right)} - i f{\\left(z,- {\\xi}_{1} \\right)} f^{2}{\\left(z,{\\xi}_{1} \\right)} + f{\\left(z,{\\xi}_{1} \\right)} {\\delta}_{1}$"
      ],
      "text/plain": [
       "Eq(Derivative(f(z, xi[1]), z), -I*f(z, -xi[0])*f(z, -xi[2])*f(z, -xi[3])*h(z) - I*f(z, -xi[1])*f(z, xi[1])**2 + f(z, xi[1])*delta[1])"
      ]
     },
     "execution_count": 579,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_diff_eqs_b[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 580,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial z} f{\\left(z,{\\xi}_{2} \\right)} = - i f{\\left(z,- {\\xi}_{0} \\right)} f{\\left(z,- {\\xi}_{1} \\right)} f{\\left(z,- {\\xi}_{3} \\right)} h{\\left(z \\right)} - i f{\\left(z,- {\\xi}_{2} \\right)} f^{2}{\\left(z,{\\xi}_{2} \\right)} + f{\\left(z,{\\xi}_{2} \\right)} {\\delta}_{2}$"
      ],
      "text/plain": [
       "Eq(Derivative(f(z, xi[2]), z), -I*f(z, -xi[0])*f(z, -xi[1])*f(z, -xi[3])*h(z) - I*f(z, -xi[2])*f(z, xi[2])**2 + f(z, xi[2])*delta[2])"
      ]
     },
     "execution_count": 580,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_diff_eqs_b[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 581,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial z} f{\\left(z,{\\xi}_{3} \\right)} = - i f{\\left(z,- {\\xi}_{0} \\right)} f{\\left(z,- {\\xi}_{1} \\right)} f{\\left(z,- {\\xi}_{2} \\right)} h{\\left(z \\right)} - i f{\\left(z,- {\\xi}_{3} \\right)} f^{2}{\\left(z,{\\xi}_{3} \\right)} + f{\\left(z,{\\xi}_{3} \\right)} {\\delta}_{3}$"
      ],
      "text/plain": [
       "Eq(Derivative(f(z, xi[3]), z), -I*f(z, -xi[0])*f(z, -xi[1])*f(z, -xi[2])*h(z) - I*f(z, -xi[3])*f(z, xi[3])**2 + f(z, xi[3])*delta[3])"
      ]
     },
     "execution_count": 581,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_diff_eqs_b[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 582,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial z} f{\\left(z,- {\\xi}_{0} \\right)} = i f^{2}{\\left(z,- {\\xi}_{0} \\right)} f{\\left(z,{\\xi}_{0} \\right)} - f{\\left(z,- {\\xi}_{0} \\right)} {\\delta}_{0} + \\frac{i f{\\left(z,{\\xi}_{1} \\right)} f{\\left(z,{\\xi}_{2} \\right)} f{\\left(z,{\\xi}_{3} \\right)}}{h{\\left(z \\right)}}$"
      ],
      "text/plain": [
       "Eq(Derivative(f(z, -xi[0]), z), I*f(z, -xi[0])**2*f(z, xi[0]) - f(z, -xi[0])*delta[0] + I*f(z, xi[1])*f(z, xi[2])*f(z, xi[3])/h(z))"
      ]
     },
     "execution_count": 582,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_diff_eqs_b[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 583,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial z} f{\\left(z,- {\\xi}_{1} \\right)} = \\frac{i f{\\left(z,{\\xi}_{0} \\right)} f{\\left(z,{\\xi}_{2} \\right)} f{\\left(z,{\\xi}_{3} \\right)}}{h{\\left(z \\right)}} + i f^{2}{\\left(z,- {\\xi}_{1} \\right)} f{\\left(z,{\\xi}_{1} \\right)} - f{\\left(z,- {\\xi}_{1} \\right)} {\\delta}_{1}$"
      ],
      "text/plain": [
       "Eq(Derivative(f(z, -xi[1]), z), I*f(z, xi[0])*f(z, xi[2])*f(z, xi[3])/h(z) + I*f(z, -xi[1])**2*f(z, xi[1]) - f(z, -xi[1])*delta[1])"
      ]
     },
     "execution_count": 583,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_diff_eqs_b[5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 584,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial z} f{\\left(z,- {\\xi}_{2} \\right)} = \\frac{i f{\\left(z,{\\xi}_{0} \\right)} f{\\left(z,{\\xi}_{1} \\right)} f{\\left(z,{\\xi}_{3} \\right)}}{h{\\left(z \\right)}} + i f^{2}{\\left(z,- {\\xi}_{2} \\right)} f{\\left(z,{\\xi}_{2} \\right)} - f{\\left(z,- {\\xi}_{2} \\right)} {\\delta}_{2}$"
      ],
      "text/plain": [
       "Eq(Derivative(f(z, -xi[2]), z), I*f(z, xi[0])*f(z, xi[1])*f(z, xi[3])/h(z) + I*f(z, -xi[2])**2*f(z, xi[2]) - f(z, -xi[2])*delta[2])"
      ]
     },
     "execution_count": 584,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_diff_eqs_b[6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 585,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial z} f{\\left(z,- {\\xi}_{3} \\right)} = \\frac{i f{\\left(z,{\\xi}_{0} \\right)} f{\\left(z,{\\xi}_{1} \\right)} f{\\left(z,{\\xi}_{2} \\right)}}{h{\\left(z \\right)}} + i f^{2}{\\left(z,- {\\xi}_{3} \\right)} f{\\left(z,{\\xi}_{3} \\right)} - f{\\left(z,- {\\xi}_{3} \\right)} {\\delta}_{3}$"
      ],
      "text/plain": [
       "Eq(Derivative(f(z, -xi[3]), z), I*f(z, xi[0])*f(z, xi[1])*f(z, xi[2])/h(z) + I*f(z, -xi[3])**2*f(z, xi[3]) - f(z, -xi[3])*delta[3])"
      ]
     },
     "execution_count": 585,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f_diff_eqs_b[7]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 600,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial z} f{\\left(z,- {\\xi}_{0} \\right)} f{\\left(z,{\\xi}_{0} \\right)} = - i f{\\left(z,- {\\xi}_{0} \\right)} f{\\left(z,- {\\xi}_{1} \\right)} f{\\left(z,- {\\xi}_{2} \\right)} f{\\left(z,- {\\xi}_{3} \\right)} h{\\left(z \\right)} + \\frac{i f{\\left(z,{\\xi}_{0} \\right)} f{\\left(z,{\\xi}_{1} \\right)} f{\\left(z,{\\xi}_{2} \\right)} f{\\left(z,{\\xi}_{3} \\right)}}{h{\\left(z \\right)}}$"
      ],
      "text/plain": [
       "Eq(Derivative(f(z, -xi[0])*f(z, xi[0]), z), -I*f(z, -xi[0])*f(z, -xi[1])*f(z, -xi[2])*f(z, -xi[3])*h(z) + I*f(z, xi[0])*f(z, xi[1])*f(z, xi[2])*f(z, xi[3])/h(z))"
      ]
     },
     "execution_count": 600,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(\n",
    "    Derivative(f(z,xi[0])*f(z,-xi[0]),z), \n",
    "    diff(f(z,xi[0])*f(z,-xi[0]),z).subs([f_diff_eqs_b[0].args, f_diff_eqs_b[4].args]).expand()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 601,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial z} \\frac{f{\\left(z,{\\xi}_{0} \\right)}}{f{\\left(z,- {\\xi}_{0} \\right)}} = - 2 i f^{2}{\\left(z,{\\xi}_{0} \\right)} + \\frac{2 f{\\left(z,{\\xi}_{0} \\right)} {\\delta}_{0}}{f{\\left(z,- {\\xi}_{0} \\right)}} - \\frac{i f{\\left(z,- {\\xi}_{1} \\right)} f{\\left(z,- {\\xi}_{2} \\right)} f{\\left(z,- {\\xi}_{3} \\right)} h{\\left(z \\right)}}{f{\\left(z,- {\\xi}_{0} \\right)}} - \\frac{i f{\\left(z,{\\xi}_{0} \\right)} f{\\left(z,{\\xi}_{1} \\right)} f{\\left(z,{\\xi}_{2} \\right)} f{\\left(z,{\\xi}_{3} \\right)}}{f^{2}{\\left(z,- {\\xi}_{0} \\right)} h{\\left(z \\right)}}$"
      ],
      "text/plain": [
       "Eq(Derivative(f(z, xi[0])/f(z, -xi[0]), z), -2*I*f(z, xi[0])**2 + 2*f(z, xi[0])*delta[0]/f(z, -xi[0]) - I*f(z, -xi[1])*f(z, -xi[2])*f(z, -xi[3])*h(z)/f(z, -xi[0]) - I*f(z, xi[0])*f(z, xi[1])*f(z, xi[2])*f(z, xi[3])/(f(z, -xi[0])**2*h(z)))"
      ]
     },
     "execution_count": 601,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Eq(\n",
    "    Derivative(f(z,xi[0])/f(z,-xi[0]),z), \n",
    "    diff(f(z,xi[0])/f(z,-xi[0]),z).subs([f_diff_eqs_b[0].args, f_diff_eqs_b[4].args]).expand()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}