I have a system of coupled PDEs which I solve numerically. I wish to optimise a particular objective function using gradient ascent, where the system of PDEs feature as constraints. I can't get the method to work so any help in checking that I've not misunderstood anything would be appreciated. Sorry that the problem is quite verbose but I've tried to write down every step I've carried out.
Here are the PDEs,
\begin{align} &\begin{aligned} \partial_{\tau} \mathcal{G}(z, \tau) = -i \Omega^{*}_{c}(\tau) P_{ge}(z, \tau) + c.c. \end{aligned}\label{G}\\ &\begin{aligned} \partial_{\tau} \xi(z, \tau) =i \Omega^{*}_{c}(\tau) P_{ge}(z, \tau) + c.c. \end{aligned}\label{Xi}\\ &\begin{aligned} \partial_{\tau} P_{ge}(z, \tau) = -i \Delta_{c} P_{ge}(z, \tau) -i \Omega_{c}(\tau) (\mathcal{G}(z, \tau) - \xi(z, \tau)) \end{aligned}\label{Pge}\\ &\begin{aligned} \partial_{\tau} S_{gs}(z, \tau) = -i(\Delta_{c} + \Delta_{t})S_{gs}(z, \tau) + i \Omega_{c}(\tau) P_{es}(z, \tau) - \sqrt{d} P_{ge}(z, \tau) \mathcal{E}_{t}(z, \tau) \end{aligned}\label{Sgs}\\ &\begin{aligned} \partial_{\tau} P_{es}(z, \tau) = -i \Delta_{t} P_{es}(z, \tau) - \sqrt{d} \xi(z, \tau) \mathcal{E}_{t}(z, \tau) + i \Omega^{*}_{c}(\tau) S_{gs}(z, \tau) \end{aligned}\label{Pes}\\ &\begin{aligned} \partial_{z}\mathcal{E}_{t}(z, \tau) = -\sqrt{d} P_{es}(z, \tau) \end{aligned}\label{Et} \end{align} where $\mathcal{G}$ and $\xi$ are strictly real, and c.c stands for the complex conjugate.
Initial boundary conditions, \begin{align} &\begin{aligned} \partial_{\tau} \mathcal{G}(z, 0) = 1 \end{aligned}\label{Gb}\\ &\begin{aligned} \partial_{\tau} \xi(z, 0) = 0 \end{aligned}\label{Xib}\\ &\begin{aligned} \partial_{\tau} P_{ge}(z, 0) = 0 \end{aligned}\label{Pgeb}\\ &\begin{aligned} \partial_{\tau} S_{gs}(z, 0) = 0 \end{aligned}\label{Sgsb}\\ &\begin{aligned} \partial_{\tau} P_{es}(z, 0) = 0 \end{aligned}\label{Pesb}\\ &\begin{aligned} \partial_{z}\mathcal{E}_{t}(0, \tau) = E_{in}(\tau) \end{aligned}\label{Etb} \end{align}
Here is the objective funcition to maximise,
\begin{equation} \begin{split} J & = \int_{0}^{1}dz S_{gs}(z, T) S^{*}_{gs}(z, T)\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{\mathcal{G}}(z, \tau)[-\dot{\mathcal{G}} - i \Omega^{*}_{c} P_{ge} + i \Omega_{c} P^{*}_{ge}] + \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{\xi}(z, \tau)[-\dot{\xi} + i \Omega^{*}_{c} P_{ge} - i \Omega_{c} P^{*}_{ge}]\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{P}^{*}_{ge}(z, \tau)[-\dot{P_{ge}} - i\Delta_{c} P_{ge} - i \Omega_{c} (\mathcal{G} - \xi)] + c.c.\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{S}^{*}_{gs}(z, \tau)[-\dot{S_{gs}} + i \Omega_{c} P_{es} + \sqrt{d} P_{ge} \mathcal{E}_{t}] + c.c.\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{P}^{*}_{es}(z, \tau)[-\dot{P_{es}} - i\Delta_{t} P_{es} + \sqrt{d} \xi \mathcal{E}_{t} + i \Omega^{*}_{c} S_{gs}] + c.c.\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{\mathcal{E}}^{*}_{t}(z, \tau)[-\partial_{z}\mathcal{E}_{t} - \sqrt{d}P_{es}] + c.c. \end{split} \end{equation} where Lagrange multiplier functions have been prepended to each of the constraints and denoted by hats. To deal with the functions being complex, we have treated each function and its complex conjugate as independent. Note that this leads to the below (slightly modified) gradient ascent update.
I want to maximise this function with respect to $\Omega(\tau)$, \begin{align} &\begin{aligned} \Omega \leftarrow \Omega + \lambda \frac{\partial J}{\partial \Omega^{*}} \end{aligned}\\ &\begin{aligned} \frac{\partial J}{\partial \Omega^{*}} = i \int_{0}^{1}dz\int_{0}^{T}d\tau\ (\hat{\xi} -\hat{\mathcal{G}}) P_{ge} - \hat{P}_{ge}(\xi - \mathcal{G}) + \hat{P}^{*}_{es} S_{gs} - \hat{S}_{gs} P^{*}_{es} \end{aligned} \end{align} where $\lambda$ represents the step size of the gradient ascent.
The Lagrange multiplier functions feature in the gradient, so we use the fact that J must be stationary with respect to variations in $\mathcal{G}$, $\xi$, $P_{ge}$, $S_{gs}$, $P_{es}$ and $\mathcal{E}_{t}$ in order to derive equations of motion and boundary conditions for these functions,
\begin{equation} \begin{split} \delta J & = \int_{0}^{1}dz S_{gs}(z, T) \delta S^{*}_{gs}(z, T) + \delta S_{gs}(z, T) S^{*}_{gs}(z, T)\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{\mathcal{G}}[-\delta\dot{\mathcal{G}} - i \Omega^{*}_{c} \delta P_{ge} + i \Omega_{c} \delta P^{*}_{ge}]\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{\xi}[-\delta\dot{\xi} + i \Omega^{*}_{c} \delta P_{ge} - i \Omega_{c} \delta P^{*}_{ge}]\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{P}^{*}_{ge}[-\delta\dot{P_{ge}} -i\Delta_{c}\delta P_{ge} - i \Omega_{c} (\delta\mathcal{G} - \delta\xi)] + c.c.\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{S}^{*}_{gs}[-\delta\dot{S_{gs}} + i \Omega_{c} \delta P_{es} + \sqrt{d}( \delta P_{ge} \mathcal{E}_{t} + P_{ge} \delta\mathcal{E}_{t})] + c.c.\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{P}^{*}_{es}[-\delta\dot{P_{es}} -i\Delta_{t}\delta P_{es} + \sqrt{d} (\delta \xi \mathcal{E}_{t} + \xi \delta\mathcal{E}_{t}) + i \Omega^{*}_{c} \delta S_{gs}] + c.c.\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \hat{\mathcal{E}}^{*}_{t}[-\partial_{z}\delta\mathcal{E}_{t} - \sqrt{d}\delta P_{es}] + c.c. \end{split} \end{equation}
Next we integrate by parts and use that $\delta S_{gs}(z, 0) = \delta P_{ge}(z, 0) = \delta P_{es}(z, 0) = \delta \mathcal{G}(z, 0) = \delta \xi(z, 0) = \delta \mathcal{E}_{t}(0, \tau) = 0 $ as initial conditions are fixed, \begin{equation} \begin{split} \delta J & = \int_{0}^{1}dz S_{gs}(z, T) \delta S^{*}_{gs}(z, T) + \delta S_{gs}(z, T) S^{*}_{gs}(z, T)\\ &- \int_{0}^{1}dz \hat{\mathcal{G}}(z, T)\delta\mathcal{G}(z, T) + \int_{0}^{1}dz\int_{0}^{T}d\tau \dot{\hat{\mathcal{G}}}\delta\mathcal{G} + \hat{\mathcal{G}}[- i \Omega^{*}_{c} \delta P_{ge} + i \Omega_{c} \delta P^{*}_{ge}]\\ &- \int_{0}^{1}dz \hat{\xi}(z, T)\delta\xi(z, T)+ \int_{0}^{1}dz\int_{0}^{T}d\tau \dot{\hat{\xi}}\delta\xi + \hat{\xi}[i \Omega^{*}_{c} \delta P_{ge} - i \Omega_{c} \delta P^{*}_{ge}]\\ &- \int_{0}^{1}dz \hat{P^{*}_{ge}}(z, T)\delta P_{ge}(z, T) + \int_{0}^{1}dz\int_{0}^{T}d\tau \dot{\hat{P}}^{*}_{ge}\delta P_{ge} + \hat{P^{*}_{ge}}[-i\Delta_{c}\delta P_{ge} - i \Omega_{c} (\delta\mathcal{G} - \delta\xi) + c.c.\\ &- \int_{0}^{1}dz \hat{S^{*}_{gs}}(z, T)\delta S_{gs}(z, T)\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \dot{\hat{S}}^{*}_{gs}\delta S_{gs} + \hat{S^{*}_{gs}}[i \Omega_{c} \delta P_{es} + \sqrt{d}( \delta P_{ge} \mathcal{E}_{t} + P_{ge} \delta\mathcal{E}_{t})] + c.c.\\ &- \int_{0}^{1}dz \hat{P^{*}_{es}}(z, T)\delta P_{es}(z, T)\\ &+ \int_{0}^{1}dz\int_{0}^{T}d\tau \dot{\hat{P}}^{*}_{es}\delta P_{es} + \hat{P^{*}_{es}}[-i\Delta_{t}\delta P_{es} + \sqrt{d} (\delta \xi \mathcal{E}_{t} + \xi \delta\mathcal{E}_{t}) + i \Omega^{*}_{c} \delta S_{gs}] + c.c.\\ &- \int_{0}^{1}d\tau \hat{\mathcal{E}}_{t}(L, \tau)\delta\mathcal{E}_{t}(L, \tau) + \int_{0}^{1}dz\int_{0}^{T}d\tau \ \partial_{z}{\hat{\mathcal{E}}}^{*}_{t}\delta \mathcal{E}_{t} + \hat{\mathcal{E}}^{*}_{t}[- (\sqrt{d})^{*}\delta P_{es}] + c.c. \end{split} \end{equation}
Optimum requires that $\delta J = 0$ for any variation and so we can collect terms multiplying $\delta X$ (where $X$ represents some variable) and set the results to zero. This results in the following equations of motion for the Lagrange multipier functions,
\begin{align} &\begin{aligned} \partial_{\tau} \hat{\mathcal{G}} = -i \Omega^{*}_{c} \hat{P}_{ge} + i \Omega_{c} \hat{P}^{*}_{ge} \end{aligned}\label{Gadj}\\ &\begin{aligned} \partial_{\tau} \hat{\xi} = - \sqrt{d}(\mathcal{E}_{t} \hat{P}^{*}_{es} + \mathcal{E}^{*}_{t} \hat{P}_{es}) - i \Omega_{c} \hat{P}^{*}_{ge} + i \Omega^{*}_{c} \hat{P}_{ge} \end{aligned}\label{Xiadj}\\ &\begin{aligned} \partial_{\tau} \hat{P}_{ge} = - i \Delta_{c} \hat{P}_{ge} - \sqrt{d} \mathcal{E}^{*}_{t} \hat{S}_{gs} -i\Omega_{c}(\hat{\xi} - \hat{\mathcal{G}}) \end{aligned}\label{Pgeadj}\\ &\begin{aligned} \partial_{\tau} \hat{S}_{gs} = i \Omega_{c} \hat{P}_{es} \end{aligned}\label{Sgsadj}\\ &\begin{aligned} \partial_{\tau} \hat{P}_{es} =i \Delta_{t} \hat{P}_{es} + \sqrt{d} \hat{\mathcal{E}}_{t} \end{aligned}\label{Pesadj}\\ &\begin{aligned} \partial_{z} \hat{\mathcal{E}}_{t} = -\sqrt{d} (\xi\hat{P}_{es} + P^{*}_{ge}\hat{S}_{gs})\, \end{aligned}\label{Eadj} \end{align}
and the following boundary conditions, \begin{align} &\begin{aligned} \hat{\mathcal{G}}(z, T) = 0 \end{aligned}\\ &\begin{aligned} \hat{\xi}(z, T) = 0 \end{aligned}\\ &\begin{aligned} \hat{P}_{ge}(z, T) = 0 \end{aligned}\\ &\begin{aligned} \hat{S}_{gs}(z, T) = S_{gs}(z, T) \end{aligned}\\ &\begin{aligned} \hat{P}_{es}(z, T) = 0 \end{aligned}\\ &\begin{aligned} \hat{\mathcal{E}}_{t}(L, \tau) = 0\, \end{aligned} \end{align}
However, solving the initial equations for some $\Omega(\tau)$, then solving the Lagrange multipier equations, calculating the gradient and updating $\Omega(\tau)$, does not cause the objective function to increase.
Any help would be greatly appreciated.