Minimization

import mmf_setup;mmf_setup.nbinit()

This cell adds /home/docs/checkouts/readthedocs.org/user_builds/gpe/checkouts/latest/src to your path, and contains some definitions for equations and some CSS for styling the notebook. If things look a bit strange, please try the following:

  • Choose "Trust Notebook" from the "File" menu.
  • Re-execute this cell.
  • Reload the notebook.

Minimization#

Motivation#

To start a simulation, one needs an initial state, and a good approximation is often the ground state in some external potential. I.e., an experiment might cool the trapped gas in an external potential ending with a fixed number of particles \(N\).

For the GPE with energy density \(\mathcal{E}_{GPE}(n) = gn^2/2\), one can use the Thomas-Fermi approximation

\[\begin{gather*} n_{TF}(\vect{r}) = \frac{\mu - V(\vect{r})}{g}. \end{gather*}\]

This corresponds to minimizing the energy without any kinetic energy at fixed total particle number:

\[\begin{gather*} \min_{n(x)}E_{TF}[n] = \int \d^{d}{\vect{r}}\; \Bigl(\mathcal{E}(n) + n(\vect{r})V(\vect{r})\Bigr), \qquad \int \d^{d}{\vect{r}}\;n(\vect{r}) = N,\\ \frac{\delta}{\delta n(x)} \bigl(E_{TF}[n] -\mu N[n]\bigr) = 0,\qquad \mathcal{E}'(n_{TF}) = \mu - V(\vect{r}). \end{gather*}\]

This neglects the kinetic energy and will not be a good approximation at the edges of the cloud where the density vanishes. More generally, the full GPE follows from minimizing the following energy \(E\)

\[\begin{gather*} E[\psi] = \int \d^{d}{\vect{r}} \left( \frac{\abs{-\I\hbar \vect\nabla \psi(\vect{r})}^2}{2m} + \mathcal{E}(n) + n(\vect{r})V(\vect{r}) \right), \qquad n(\vect{r}) = \psi^\dagger(\vect{r})\psi(\vect{r}). \end{gather*}\]

Performing a functional minimization defines the single-particle Hamiltonian \(\op{H}\) whose action is

\[\begin{gather*} \op{H}\psi(\vect{r}) = \pdiff{E[\psi]}{\psi^\dagger(\vect{r})} = \mu \psi(\vect{r}). \end{gather*}\]

Numerical Techniques#

A reliable strategy for minimizing a function \(E(x)\) is gradient descent: simply proceed in the downhill direction \(-\nabla E(x)\) opposite to the gradient. This can be formally implemented by solving the ODE:

\[\begin{gather*} \dot{x} = - \nabla E(x), \end{gather*}\]

with a simplistic approach \(\delta x = - \delta t \nabla E\) with an appropriately chosen step size \(\delta t\). Since we don’t care about actually solving the ODE, we can choose large steps \(\delta t\), monitoring the energy \(E(x)\), and reducing the step size if the energy does not decrease. With suitable termination criterion, this method is guaranteed to find a local minimum (or to diverge) if the function \(E(x)\) is sufficiently smooth. Alternatively, one can use an adaptive IVP solver to determine the step size, but this can be overkill.

Constraints can be implemented by ensuring that the step \(\delta x\) is orthogonal to the constraint, or a corrections can be made. To see how this works formally, consider the following problem of minimizing \(E\) subject to fixed \(N\) (which w.l.g. we take to be zero here):

\[\begin{gather*} \min_{\ket{x}} E, \qquad N = 0,\\ \mat{H}\ket{x} = \pdiff{E}{\bra{x}}, \qquad \mat{C}\ket{x} = \pdiff{N}{\bra{x}},\qquad (\mat{H} - \mu \mat{C})\ket{x} = 0. \end{gather*}\]

The formulation as a gradient descent

\[\begin{gather*} \ket{\dot{x}} = -\pdiff{E}{\bra{x}} = -\mat{H}\ket{x} \end{gather*}\]

follows from

\[\begin{gather*} \dot{E} = \bra{\dot{x}}\pdiff{E}{\bra{x}} = \braket{\dot{x}|\mat{H}|x} = -\braket{x|\mat{H}^\dagger\mat{H}|x} \leq 0. \end{gather*}\]

Evolution is guaranteed to proceed downhill unless the gradient vanishes.

The constraint can be enforced by projecting onto the null-space of the constraint:

\[\begin{gather*} \ket{\dot{x}} = \mat{P}_{\perp}\mat{H}\ket{x} = \mat{P}_{\perp}\mat{H}\ket{x} \end{gather*}\]

subtracting any component of the proposed step that has a non-zero projection along the gradient of the constraint:

\[\begin{gather*} \ket{\dot{x}} = \mat{H}\ket{x} - \ket{\alpha}. \end{gather*}\]

The appropriate correction follows from requiring

\[\begin{gather*} \dot{N} = 0 = \braket{\dot{x}|\mat{C}|x} = \braket{x|\mat{H}^\dagger\mat{C}|x} + \braket{\alpha|\mat{C}|x},\\ \braket{x|\mat{C}^\dagger|\alpha} = -\braket{x|\mat{C}^\dagger\mat{H}|x}. \end{gather*}\]
\[\begin{gather*} \braket{x|\mat{C}^\dagger \end{gather*}\]

\(\ket{x} \rightarrow \op{P}_{N}\ket{x}\) the step into the null-space of the constraint:

\[\begin{gather*} \ket{\dot{x}} = -\mat{P}_{N}\mat{H}\ket{x}, \qquad \mat{C}\op{P}_{N}\ket{x} = 0. \end{gather*}\]

This is equivalent to subtracting the component appropriate component

Consider trying to find the minimum of a function \(E(x)\). One approach is gradient descent:

\[\begin{gather*} \dot{x} = -E'(x): \end{gather*}\]

I.e., move downhill towards the minimum. Closely related is the arrested Newton flow algorithm, which attempts to minimize the energy by following the trajectory that a particle of unit mass \(m=1\) would follow in a potential \(V(x) = E(x)\)

\[\begin{gather*} \ddot{x} = -V'(x) = -E'(x) \end{gather*}\]

starting from rest \(\dot{x} = 0\). Of course, due to inertia, as the velocity rises, the particle might overshoot the minimum and start to rise \(\dot{V} > 0\). If this happens one sets the kinetic energy to zero \(\dot{x} = 0\) and restart from this point.

Here we play with this a bit, but I don’t see why it is better than gradient descent, especially if we add a bit of step control.

Numerical Minimization#

Numerically, integrals are performed by summing over the lattice points, and include a factor of the metric \(\d{\vect{r}}^d\) which we call metric in the code and must be explicitly included in the minimizer. Thus, we can formulate the problem as a minimization of the objective function

\[\begin{gather*} f(\vect{x}) = E[\psi], \qquad N[\vect{x}] = N_0, \qquad \vect{x} = P(\psi) \end{gather*}\]

where \(\vect{x} = P(\psi)\) is a real parametrization of \(\psi\) (i.e. the real and imaginary components in some order. We use the underlying memory layout to efficiently do this conversion). The minimizer requires the gradient, so we have:

\[\begin{gather*} \nabla f(\vect{x}) = P\left(\mat{H}\psi\right)(2\d{\vect{r}}^d). \end{gather*}\]

The last factor of 2*metric accounts for the complex-to-real conversion and the metric implicit in the summation of the integral.

Fixed Particle Number#

If the minimization routine explicitly deals with constraints, we can use the previous formalism. However, if we need to manually fix the particle number, we can instead write the objective function as

\[\begin{gather*} \tilde{f}(\vect{x}) = \tilde{E}[\psi] = E[s\psi] + \frac{E_0}{4}\left(\frac{1}{s^2} - 1\right)^2, \qquad s = \sqrt{\frac{N}{\braket{\psi|\psi}}}\\ \pdiff{\tilde{E}}{\psi^\dagger} = s^*(\op{H} - \mu)\ket{\tilde{\psi}} + \frac{E_0}{2s\braket{\tilde{\psi}|\tilde{\psi}}}\left(\frac{1}{s^2} - 1\right)\ket{\tilde{\psi}}, \qquad \mu = \frac{\braket{\tilde{\psi}|\op{H}|\tilde{\psi}}}{\braket{\tilde{\psi}|\tilde{\psi}}}. \end{gather*}\]
\[\begin{gather*} \pdiff{}{\psi_i^\dagger}\frac{E_0}{4}\left(\frac{1}{s^2} - 1\right)^2 = \frac{-E_0}{s^3}\left(\frac{1}{s^2} - 1\right)\pdiff{s}{\psi_i^\dagger} = \frac{E_0}{2s^2}\left(\frac{1}{s^2} - 1\right)\frac{\psi_i}{\braket{\psi|\psi}} = \frac{E_0}{2s}\left(\frac{1}{s^2} - 1\right)\frac{\tilde{\psi}_i}{\braket{\tilde{\psi}|\tilde{\psi}}}\\ = \frac{E_0}{2sN}\left(\frac{1}{s^2} - 1\right)\tilde{\psi}_i \end{gather*}\]

Detailed derivation

Here is the detailed derivation:

\[\begin{gather*} \pdiff{}{\psi^\dagger_i} s = -s\frac{\psi_i}{2\braket{\psi|\psi}}, \qquad \pdiff{}{\psi_i} s = -s\frac{\psi_i^\dagger}{2\braket{\psi|\psi}},\\ \end{gather*}\]

Thus, let \(\tilde{\psi} = s\psi\), then we have

(21)#\[\begin{align} \pdiff{\tilde{E}}{\psi^\dagger_i} &= \pdiff{E[\tilde{\psi}]}{\tilde{\psi}^\dagger_j} \pdiff{\tilde{\psi}^\dagger_j}{\psi^\dagger_i} + \pdiff{E[\tilde{\psi}]}{\tilde{\psi}_j} \pdiff{\tilde{\psi}_j}{\psi^\dagger_i} \\ &= [\op{H}\tilde{\psi}]_j \left( s\delta_{ij} - \psi^\dagger_j s\frac{\psi_i}{2\braket{\psi|\psi}} \right) + [\tilde{\psi}^\dagger\op{H}]_j \left( - \psi_j s\frac{\psi_i}{2\braket{\psi|\psi}} \right)\\ &= s[\op{H}\tilde{\psi}]_i - s\Re\frac{\braket{\psi|\op{H}|\tilde{\psi}}}{\braket{\psi|\psi}} \psi_i = s[\op{H}\tilde{\psi}]_i - s\Re\frac{\braket{\tilde{\psi}|\op{H}|\tilde{\psi}}}{\braket{\tilde{\psi}|\tilde{\psi}}} \tilde{\psi}_i\\ &= s[\op{H}-\mu]\tilde{\psi}. \end{align}\]

Fixed Phase#

For initialization vorticity and turbulence, it can be useful to minimize the state with fixed phase \(\phi\). In this case, we minimize with respect to \(\sqrt{n}\) which is real:

\[ \tilde{f}(\sqrt{n}) = E[\overbrace{\sqrt{n}\phi}^{\psi}], \qquad \pdiff{\tilde{f}}{\sqrt{n}} = \pdiff{E}{\psi^\dagger}[\sqrt{n}\phi]\pdiff{\sqrt{n}\phi^\dagger}{\sqrt{n}} = \phi^*\mat{H}\psi. \]

One should really take the real part of this since.

%pylab inline --no-import-all
from importlib import reload
from gpe import bec
u = bec.u
from gpe.imports import *
from gpe import minimize;reload(minimize)
from gpe.minimize import MinimizeStateFixedPhase

class State(bec.StateHOMixin, bec.StateBase):
    def __init__(self, phase=None, **kw):
        if phase is not None:
            phase = np.exp(1j*np.angle(phase))
        self.phase = phase
        super().__init__(**kw)

s0 = State(Nxyz=(64, 64), Lxyz=(10., 10.), ws=(200.0*u.Hz,)*2, mu=0.01)
x, y = s0.xyz
phase = np.exp(1j*np.angle(x+1j*y))

%pylab is deprecated, use %matplotlib inline and import the required libraries.
Populating the interactive namespace from numpy and matplotlib
[I 03:37:56 root] Patching zope.interface.document.asReStructuredText to format code
[I 03:37:56 numexpr.utils] NumExpr defaulting to 2 threads.

KeyboardInterrupt

m = MinimizeStateFixedPhase(s0, fix_N=False, phase=phase)
m.check()
def callback(state):
    plt.clf()
    state.plot()
    display(plt.gcf())
    clear_output(wait=True)
    plt.close('all')
s = m.minimize(gtol=1e-12, ftol=1e-12, disp=1)#callback=callback)
m = MinimizeStateFixedPhase(s, fix_N=True, phase=phase)
s = m.minimize(gtol=1e-12, ftol=1e-12, disp=1)#callback=callback)
m = MinimizeStateFixedPhase(s, fix_N=True, phase=phase)
s = m.minimize(gtol=1e-12, ftol=1e-12, disp=1)#callback=callback)
m = MinimizeStateFixedPhase(s, fix_N=True, phase=phase)
s = m.minimize(gtol=1e-12, ftol=1e-12, disp=1)#callback=callback)
m = MinimizeStateFixedPhase(s, fix_N=True, phase=phase)
s = m.minimize(gtol=1e-12, ftol=1e-12, disp=1)#callback=callback)
m = MinimizeStateFixedPhase(s, fix_N=True, phase=phase)
s = m.minimize(gtol=1e-12, ftol=1e-12, disp=1)#callback=callback)
s.plot()

s.minimize_results

s1 = s0.copy()
m = MinimizeStateFixedPhase(s1, fix_N=True, phase=phase)
m.check()
s1 = m.minimize(callback=callback, psi_tol=1e-16, E_tol=1e-16)
s1.plot()

m = MinimizeStateFixedPhase(s1, fix_N=True, phase=phase)
m.check()
s1 = m.minimize()
s1.plot()

s.cooling_phase = 1.0j
ev = EvolverABM(s1, dt=0.1*s.t_scale)
NoInterrupt.unregister()
with NoInterrupt() as interrupted:
    while not interrupted:
        ev.evolve(100)
        plt.clf()
        ev.y.plot()
        display(plt.gcf())
        clear_output(wait=True)

Harmonic Oscillator#

As a simple test, we consider a simple harmonic oscillator setting \(V(x) = m\omega^2x^2/2\) and \(g=0\) which has a ground state wavefunction:

\[ \psi_0(x) = \frac{1}{\pi^{1/4}\sqrt{a}} e^{-x^2/2a^2}, \qquad a = \sqrt{\frac{\hbar}{m\omega}}. \]

To ensure we have a sufficiently large box, we choose \(L\) so that the density falls by a factor of \(10^{-15}\) at the boundary. This will ensure that we have no long-range (IR) errors.

\[ n(L/2) = e^{-L^2/4a^2} \approx 2^{-52} \approx 2\times 10^{-16} \approx e^{-36},\qquad L \approx 12.0 a. \]

The code for this (exact_solutions.HarmonicOscillator) also allows for interacting systems \(g \neq 0\), adjusting the potential so that the ground state is still Gaussian. We also test against this exact solution.

%pylab inline --no-import-all
from gpe import bec;reload(bec)
from gpe.bec import u
dim = 1
w = 1.0
a = np.sqrt(u.hbar/u.m/w)
N = 32
L = 12*a
E0 = dim * u.hbar * w / 2.0

def get_analytic_state(N=N, L=L):
    """Return the analytic ground state."""
    s = bec.State(Nxyz=(N,)*dim, Lxyz=(L,)*dim, ws=(w,)*dim, g=0.0, N=1.0)
    r2 = sum(_x**2 for _x in s.xyz)
    s[...] = np.sqrt(np.exp(-r2/a**2)/np.sqrt(np.pi)/a)
    return s

Ns = 2**np.arange(1,8)
Es = np.array([get_analytic_state(_N).get_energy() for _N in Ns])
Es_ = np.array([get_analytic_state(_N, L=L/1.1).get_energy() for _N in Ns])
plt.semilogy(Ns, abs(Es - E0), '-+', label='L/a = {}'.format(L/a))
plt.semilogy(Ns, abs(Es_ - E0), 'r--+', label='L/a = {}'.format(L/a/1.2))
plt.xlabel('N')
plt.ylabel('err (E)')
plt.legend(loc='best')

The errors here are due to an insufficient number of lattice points \(N\) or too large of a lattice spacing \(\d{x}\) (ultraviolet or UV errors). Notice that the error decreases exponentially as a function of the lattice size, reaching machine precision at \(N=32\):

\[ \delta E \propto a^{-N} \propto (\d{x})^{N}. \]

This exponential convergence is a special property of Fourier methods for analytic functions. Finite-difference methods exhibit at best a power-law scaling \(\delta E \propto N^{a}\) (though higher-order methods can have fairly large exponents). The other curve shows a slightly smaller box, demonstrating that \(L=12a\) is required for full convergence (infrared or IR errors).

Exercise: Explain UV Errors#

I explained above why \(L=12a\) was required to reduce the errors to machine precision. Explain why \(N=32\) is sufficient for this purpose.

Numerical Example#

Having established that we can achieve machine precision in the energy, let’s see how well the minimization routine can do.

%pylab inline --no-import-all
from gpe import exact_solutions; reload(exact_solutions)
from gpe import minimize;reload(minimize);from gpe.minimize import MinimizeState

def check_minimization(State, **kw):
    s = State(**kw)  # Assumes that this is an analytic solution
    E0 = s.get_energy()
    s0 = s.copy()
    s0[...] = 1
    s0.normalize()
   
    tols = 10.0**np.arange(-1, -17, -1)
    dEs_f = []
    dEs_g = []
    dns_f = []
    dns_g = []
    for _t in tols:
        _sf = MinimizeState(s0.copy(), fix_N=True).minimize(
            use_scipy=True, ftol=_t, gtol=1e-32)
        _sg = MinimizeState(s0.copy(), fix_N=True).minimize(
            use_scipy=True, ftol=1e-32, gtol=_t)
        dEs_f.append(abs(_sf.get_energy() - E0))
        dEs_g.append(abs(_sg.get_energy() - E0))
        dns_f.append(abs(abs(_sf[...])**2 - abs(s[...])**2).max())
        dns_g.append(abs(abs(_sg[...])**2 - abs(s[...])**2).max())    

    plt.figure(figsize=(10,5))
    plt.subplot(121)
    plt.loglog(tols, dEs_f, 'r+-', label='ftol (E)')
    plt.loglog(tols, dEs_g, 'g+-', label='gtol (E)')
    plt.loglog(tols, dns_f, 'r+--', label='ftol (n)')
    plt.loglog(tols, dns_g, 'g+--', label='gtol (n)')

    plt.xlabel('tol')
    plt.ylabel('err')
    plt.legend(loc='best')    

    plt.subplot(122)
    plt.plot(s.xyz[0], abs(s[...])**2)
    plt.xlabel('x')
    plt.ylabel('n')    
    plt.twinx()
    plt.plot(s.xyz[0], s.get_Vext(), '--r')    
    plt.xlabel('x')
    plt.ylabel('V')
    
check_minimization(State=exact_solutions.HarmonicOscillator, g=0.0);
plt.title("g=0");
plt.figure();
check_minimization(State=exact_solutions.HarmonicOscillator, g=2.0);
plt.title("g=2");

Here we note several things. First - the energy can be computed to machine precision, which requires ftol\(\sim\epsilon\) and gtol\(\sim \sqrt{\epsilon}\), however, the state itself can only be computed to an absolute accuracy of \(\sqrt{\epsilon}\) because at this point, the energy no-longer changes. This is a common problem with minimization proceedures. About the minimium, the function is quadratic \(f(x) \approx f_0 + a (x-x_0)^2\) so changes in \(x \sim \delta\) will give rise to changes in \(f(x_0+\delta) \approx f_0 + a\delta^2\).

s = exact_solutions.HarmonicOscillator(g=2.0)
E0 = s.get_energy()
n0 = s.get_density()
s[...] = 1.0
s.normalize()
m = Minimize(s.copy())
s1 = m.minimize(E_tol=0.1, psi_tol=1e-1, disp=1)
print s1.get_energy() - E0, abs(s1.get_density() - n0).max()

Root-Finding#

How can we improve the precision of the solution? Instead of solving the minimization problem, we can apply a root-finding algorithm to the evolution equation. Thus, we try to find solutions to the non-linear equations:

\[ \left(\op{H}[\psi] - \mu[\psi]\right)\psi = 0, \qquad \mu = \frac{\braket{\psi|\op{H}|\psi}}{\braket{\psi|\psi}}. \]

The general strategy here for finding a root of a vector-valued function \(\vect{F}(\vect{x})\) is to use Newton’s method via the iterations:

\[ \vect{x} \rightarrow \vect{x} - \mat{J}^{-1}\cdot \vect{F} \]

where \([\mat{J}]_{ij} = \partial F_i/\partial x_j\) is the Jacobian. However, for large problems, even forming \(\mat{J}\) is generally prohibitive given the size of the wavefunction, let along computing the inverse. Instead, one can use a memory-limited quasi-Newton method like Broyden’s method which approximates \(\mat{J}\) or \(\mat{B} = \mat{J}^{-1}\) with a matrix of limited rank:

\[ \mat{B}_n = \mat{B}_0 + \sum_{i=1}^{n}\ket{a_i}\bra{b_i}. \]

This only requires storing \(2n\) vectors. One starts with a guess \(\mat{B}_0\) for the inverse of the Jacobian (often taken as the identity \(\mat{B}_0 \propto \mat{1}\), and then as one iterates towards the solution, one updates \(\mat{B}_n\) to ensure that the secant condition is satisfied:

\[ \mat{B}_n\cdot(\vect{x}_{n+1} - \vect{x}_n) = \vect{F}_{n+1} - \vect{F}_{n}. \]

In one-dimension, this defines the secant method, but in higher dimensions the choice is ambiguous. Broyden’s approach is to satisfy this condition while minimizing the change in \(\mat{B}\), typically minimizing the \(L_2\) norm of \(\mat{B}\) (the so-called bad Broyden’s so-method) of minimizing the \(L_2\) norm of \(\mat{J} = \mat{B}^{-1}\) (the so-called goot Broyden’s method). Two simple implementations of these are provided in SciPy (though they have limitations and we ultimately would prefer to implement our own.)

In order for Broyden’s method to work well, the initial guess \(\mat{B}_0\) should be a good approximation of the inverse Jacobian of the function \(\vect{F}(\vect{x})\). One way to ensure that this is roughly true for a constant \(\mat{B}_0\) is to find an iteration

\[ \vect{x} \mapsto \vect{G}(\vect{x}) \]

that converges to a fixed-point at the solution \(\vect{x}^* = \vect{G}(\vect{x}^*)\). Sometimes one can coax such an iteration to have better convergence properties by applying the following iteration called “linear mixing”:

\[ \vect{x} \rightarrow\vect{G}_\alpha(\vect{x}) = \alpha\vect{G}(\vect{x}) + (1-\alpha)\vect{x}. \]

Once such a convergent approach is found, then a good candidate for Broyden’s method is

\[ \vect{F}(\vect{x}) = \vect{x} - \vect{G}(\vect{x}). \]

where we start with \(\mat{B}_0 = \alpha\mat{1}\). This will give the convergent linear mixing step as the first step:

\[ \vect{x} \mapsto \vect{x} - \mat{B}_0\cdot\vect{F}(\vect{x}) = \alpha\vect{G}(\vect{x}) + (1-\alpha)\vect{x}. \]

Note: Constructing \(\vect{F}(x) = \vect{x} - \vect{G}_\alpha(\vect{x}) = \alpha[\vect{x} - \vect{G}(\vect{x})]\) from the convergent linear mixing iteration is equivalent to scaling the objective function.

A good choice for the iteration \(\vect{G}(\vect{x})\) when trying to minimize \(f(\vect{x})\) is to step in the direction of steepest descent:

\[ \vect{G}(\vect{x}) = \vect{x} - a\vect{\nabla}f(\vect{x}), \qquad \vect{F}(\vect{x}) = a\vect{\nabla}f(\vect{x}). \]

which is guaranteed to descend if the step-size \(a\) is small enough. In our case this follows from \(\vect{\nabla}f(\vect{x}) = (\op{H}-\mu)\psi\):

\[ \vect{G}(\psi) = \psi - a\frac{\op{H}[\psi] - \mu}{\mu}\psi, \qquad \mu = \frac{\braket{\psi|\op{H}|\psi}}{\braket{\psi|\psi}}. \]
\[ \psi \rightarrow \vect{G}(\psi) = \psi - a \frac{(H[\psi] - \mu)\psi}{\mu} = (1-a)\psi - a\frac{H[\psi]\psi}{\mu}. \]

Here we have rendered the factor \(a\) dimensionless by using the chemical potential, but one could use any reasonable energy scale. (Note: we want to conserve particle number, so we also normalize the state with each iteration, so the actual function \(\vect{G}(\vect{x})\) has an additional scale factor).

%pylab inline --no-import-all
from IPython.display import clear_output
from scipy.optimize import root

from gpe import minimize;reload(minimize);from gpe.minimize import Minimize
from gpe.exact_solutions import HarmonicOscillator
s0 = HarmonicOscillator()
E0 = s0.get_energy()
n0 = s0.get_density()

s = s0.copy()
s[...] = 1
s.normalize()
s.init()

def G(s):
    Hy = s.get_Hy(subtract_mu=True)
    s = s - Hy/s._mu
    s.normalize()
    return s

def F(s, alpha=0.1):
    return alpha*(s - G(s))

alpha = 0.1
s = s0.copy(); s[...] = 1; s.normalize(); s.init()
for n in range(20):
    s = alpha*G(s) + (1-alpha)*s
    plt.clf();s.plot();display(plt.gcf());clear_output(wait=True)

Here we demonstrate simple convergence of the iteration. Note that this requires quite a small \(\alpha\) though it does converge to a few places.

ns = np.arange(100)
alphas = [0.1, 0.05]
errs = []
for alpha in alphas:
    _errs = []
    s = s0.copy()
    s[...] = 1
    s.normalize()
    s.init()
    for n in ns:
        s = alpha*G(s) + (1-alpha)*s
        s.normalize()
        _errs.append((s.get_energy() - E0, abs(s.get_density() - n0).max()))
    errs.append(np.asarray(_errs))
for alpha, err in zip(alphas, errs):
    l = plt.semilogy(ns, abs(err[:,0]), '-', label='E, alpha={}'.format(alpha))[0]
    plt.semilogy(ns, abs(err[:,1]), '--', label='n', c=l.get_c())
plt.legend(loc='best')
plt.ylabel('err')
plt.xlabel('iteration')

Broyden Algorithm#

We thus have an algorithm that effectively minimizes the energy, but quite inefficiently, taking almost 700 function evaluations in order to achieve machine precision compared with the 59 evaluations needed by the L-BFGS-M minimizer. We can now improve this by using the Broyden algorithm. Here will use the same objective function \(f(\vect{x})\) and its gradient \(\vect{\nabla} f(\vect{x})\) considered above.

As shown above, the following linear mixing converges, although quite slowly:

\[ \vect{x} \mapsto \alpha\vect{G}(\vect{x}) + (1-\alpha)\vect{x}, \qquad \vect{G}(\vect{x}) = -\vect{\nabla}f(\vect{x}). \]

This iteration can be induced in two ways:

  • By using \(\vect{F}(\vect{x}) = \vect{x} - \vect{G}(\vect{x})\) as the objective function for the root finder with an initial inverse Jacobian of \(\mat{B}_0 = \alpha \mat{1}\).

  • By using \(\vect{F}(\vect{x}) = \alpha[\vect{x} - \vect{G}(\vect{x})]\)

%pylab inline --no-import-all
from IPython.display import clear_output
from scipy.optimize import root

import mmfutils.solve.broyden; reload(mmfutils.solve.broyden)
from mmfutils.solve import broyden

from gpe import minimize;reload(minimize);from gpe.minimize import MinimizeState
from gpe.exact_solutions import HarmonicOscillator
s0 = HarmonicOscillator()
E0 = s0.get_energy()
n0 = s0.get_density()

from scipy.optimize import line_search

def f(x):
    s = s0.empty()
    s[...] = x.view(dtype=s.dtype).reshape(s.shape)
    scale_factor = np.sqrt(s0._N/s.get_N())
    s *= scale_factor
    return s.get_energy()

def df(x):
    s = s0.empty()
    s[...] = x.view(dtype=s.dtype).reshape(s.shape)
    scale_factor = np.sqrt(s0._N/s.get_N())
    s *= scale_factor
    return (scale_factor * 2 * s.metric * 
            s.get_Hy(subtract_mu=True)[...].view(dtype=float).ravel())

def G(x):
    return x - df(x)
    
def F(x):
    return df(x)

n = 0
ns = []
errs = []
s = s0.copy()
s[...] = 1
s.normalize()
s.init()

# Start with descent direction
x0 = s[...].copy().view(dtype=float).ravel()
F0 = F(x0)
dx0 = -F0
alpha, fc, gc, new_fval, old_fval, new_slop = line_search(f, df, xk=x0, pk=dx0)
print alpha
x1 = x0 + alpha*dx0
#Jacobian, Nsteps1, Nsteps2 = broyden.Jacobian, 120, 150
Jacobian, Nsteps1, Nsteps2 = broyden.JacobianBFGS, 64, 80
B = Jacobian(alpha=alpha)
B.update(x=x0, f=F0)
dx0 = -B.solve(F0)
assert np.allclose(x0 + dx0, x1)
s[...] = x1.copy().view(dtype=s.dtype).reshape(s.shape)
s.normalize()
x1 = s[...].copy().view(dtype=float).ravel()
F1 = F(x1)
n += fc; ns.append(n)
errs.append((s.get_energy() - E0, abs(s.get_density() - n0).max()))

while n < Nsteps1:
    B.update(x=x1, f=F1)
    x0, F0 = x1, F1
    dx0 = -B.solve(F0)
    alpha, fc, gc, new_fval, old_fval, new_slop = line_search(f, df, xk=x0, pk=dx0)
    if alpha is None:
        assert np.allclose(F(x0), F0)
        dx0 = -F0
        alpha, fc, gc, new_fval, old_fval, new_slop = line_search(f, df, xk=x0, pk=dx0)
    x1 = x0 + alpha*dx0
    s[...] = x1.view(dtype=s.dtype).reshape(s.shape); s.normalize()
    x1 = s[...].copy().view(dtype=float).ravel()
    F1 = F(x1)
    n += fc; ns.append(n)
    errs.append((s.get_energy() - E0, abs(s.get_density() - n0).max()))
    
errs_ = np.asarray(errs)
plt.semilogy(ns, abs(errs_[:,0]), '-', label='E')
plt.semilogy(ns, abs(errs_[:,1]), '--', label='n')
plt.legend(loc='best')
plt.ylabel('err')
plt.xlabel('Function evaluations')

Now we polish off the densities using pure Broyden steps. Here we use a line-search algorithm from scipy version 0.19.

from scipy.optimize.nonlin import _nonlin_line_search
while n < Nsteps2:
    B.update(x=x1, f=F1)
    x0, F0 = x1, F1
    dx0 = -B.solve(F0)
    alpha, x1, F1, F1_norm = _nonlin_line_search(F, x0, F0, dx0, search_type='wolfe')
    s[...] = x1.view(dtype=s.dtype).reshape(s.shape); s.normalize()
    x1 = s[...].copy().view(dtype=float).ravel()
    F1 = F(x1)
    errs.append((s.get_energy() - E0, abs(s.get_density() - n0).max()))
    n += 1
    ns.append(n)

errs_ = np.asarray(errs)
plt.semilogy(ns, abs(errs_[:,0]), '-', label='E')
plt.semilogy(ns, abs(errs_[:,1]), '--', label='n')
plt.legend(loc='best')
plt.ylabel('err')
plt.xlabel('Function evaluations')

Here is the pure L-BFGS algorithm applied directly to the non-linear problem:

from scipy.optimize.nonlin import _nonlin_line_search

n = 0
ns = []
errs = []
s = s0.copy()
s[...] = 1
s.normalize()
s.init()

x1 = s[...].copy().view(dtype=float).ravel()
F1 = F(x1)
B = broyden.JacobianBFGS()

while n < Nsteps2:
    B.update(x=x1, f=F1)
    x0, F0 = x1, F1
    dx0 = -B.solve(F0)
    alpha, x1, F1, F1_norm = _nonlin_line_search(F, x0, F0, dx0, search_type='wolfe')
    s[...] = x1.view(dtype=s.dtype).reshape(s.shape); s.normalize()
    x1 = s[...].copy().view(dtype=float).ravel()
    F1 = F(x1)
    errs.append((s.get_energy() - E0, abs(s.get_density() - n0).max()))
    n += 1
    ns.append(n)

errs_ = np.asarray(errs)
plt.semilogy(ns, abs(errs_[:,0]), '-', label='E')
plt.semilogy(ns, abs(errs_[:,1]), '--', label='n')
plt.legend(loc='best')
plt.ylabel('err')
plt.xlabel('Function evaluations')

Initial Jacobian#

Note: By playing a bit here, we find that it is crucial for convergence of the Broyden iterations that the initial inverse Jacobian approximation \(\mat{B}_0 = \mat{J}_0^{-1} = \alpha\mat{1}\) be appropriately chosen. We find in practice that determining the factor by performing a line-search in the descent direction provides a good approximation.

Here we build up the Broyden approximation to the Jacobian using the minimize routine:

%pylab inline --no-import-all
from IPython.display import clear_output
from scipy.optimize import root

from gpe import minimize;reload(minimize);from gpe.minimize import MinimizeState
from gpe.exact_solutions import HarmonicOscillator
s0 = HarmonicOscillator()
E0 = s0.get_energy()
n0 = s0.get_density()

ns = [0]
errs = []
def callback(s):
    ns.append(ns[-1]+1)
    errs.append((s.get_energy() - E0, 
                  abs(s.get_density() - n0).max()))
s = s0.copy()
s[...] = 1
s.normalize()
s.init()
minimizer = MinimizeState(s.copy(), fix_N=True)
state = minimizer.minimize(f_tol=1e-32, x_tol=1e-32, _log=True, polish=True, 
                           broyden_alpha=0.2,
                           callback=callback)
state.minimize_res.nfev
ns = ns[1:]
errs_ = np.asarray(errs)
plt.semilogy(ns, abs(errs_[:,0]), '-', label='E')
plt.semilogy(ns, abs(errs_[:,1]), '--', label='n')
plt.semilogy([abs(_c[1] - E0) for _c in minimizer._calls], label='E (per function call)')
plt.legend(loc='best')
plt.ylabel('err')
plt.xlabel('Iterations evaluations')    

x0, F0, B, J, df = state.broyden_data
_ns = list(ns)
_errs = list(errs)

while _ns[-1] < 60:
    x1 = x0 - J.solve(F0)
    s[...] = x1.view(dtype=s.dtype).reshape(s.shape); s.normalize()
    x1 = s[...].copy().view(dtype=float).ravel()
    F1 = df(x1)
    _ns.append(_ns[-1]+1)
    _errs.append((s.get_energy() - E0, abs(s.get_density() - n0).max()))
    J.update(x1, F1)
    x0, F0 = x1, F1
    
errs_ = np.asarray(_errs)
plt.semilogy(_ns, abs(errs_[:,0]), '-', label='E')
plt.semilogy(_ns, abs(errs_[:,1]), '--', label='n')
plt.legend(loc='best')
plt.ylabel('err')
plt.xlabel('Function evaluations')    

x0, F0, B, J, df = state.broyden_data
_ns = list(ns)
_errs = list(errs)

while _ns[-1] < 60:
    x1 = x0 - (B*F0)
    s[...] = x1.view(dtype=s.dtype).reshape(s.shape); s.normalize()
    x1 = s[...].copy().view(dtype=float).ravel()
    F1 = df(x1)
    _ns.append(_ns[-1]+1)
    _errs.append((s.get_energy() - E0, abs(s.get_density() - n0).max()))
    B.update_broyden(dx=x1-x0, df=F1-F0)
    x0, F0 = x1, F1
    
errs_ = np.asarray(_errs)
plt.semilogy(_ns, abs(errs_[:,0]), '-', label='E')
plt.semilogy(_ns, abs(errs_[:,1]), '--', label='n')
plt.legend(loc='best')
plt.ylabel('err')
plt.xlabel('Function evaluations')    

_s = s.copy()
def _f(c):
    _s[...] = c[0].view(dtype=complex).reshape(_s.shape)
    return _s.get_N()
plt.plot(map(_f, minimizer._calls))

from mmfutils.solve.broyden import DyadicSum
n = 0
ns = []
errs = []
s = s0.copy()
s[...] = 1
s.normalize()
s.init()

# Start with descent direction
x0 = s[...].copy().view(dtype=float).ravel()
F0 = dx0 = -df(x0)
alpha, fc, gc, new_fval, old_fval, new_slop = line_search(f, df, xk=x0, pk=F0)
x1 = x0 + alpha*F0
B = DyadicSum(alpha=-alpha)
assert np.allclose(x0 - B*F0, x1)
s[...] = x1.copy().view(dtype=s.dtype).reshape(s.shape)
s.normalize()
x1 = s[...].copy().view(dtype=float).ravel()
F1 = -df(x1)
n += fc; ns.append(n)
errs.append((s.get_energy() - E0, abs(s.get_density() - n0).max()))

while n < 120:
    B.update_broyden(dx=x1-x0, df=F1-F0)
    x0, F0 = x1, F1
    dx0 = -(B*F0)
    alpha, fc, gc, new_fval, old_fval, new_slop = line_search(f, df, xk=x0, pk=dx0)
    if alpha is None:
        assert np.allclose(-df(x0), F0)
        dx0 = F0
        alpha, fc, gc, new_fval, old_fval, new_slop = line_search(f, df, xk=x0, pk=dx0)
    x1 = x0 + alpha*dx0
    s[...] = x1.view(dtype=s.dtype).reshape(s.shape); s.normalize()
    x1 = s[...].copy().view(dtype=float).ravel()
    F1 = -df(x1)
    n += fc; ns.append(n)
    errs.append((s.get_energy() - E0, abs(s.get_density() - n0).max()))
    
errs_ = np.asarray(errs)
plt.semilogy(ns, abs(errs_[:,0]), '-', label='E')
plt.semilogy(ns, abs(errs_[:,1]), '--', label='n')
plt.legend(loc='best')
plt.ylabel('err')
plt.xlabel('Function evaluations')

A Harder Problem#

Here is a more realistic problem that has been posing some difficulties.

%pylab inline --no-import-all
from gpe import bec2; reload(bec2)
from gpe import bec; reload(bec)
from gpe import soc_soliton; reload(soc_soliton)
from gpe import minimize;reload(minimize)
from gpe.soc_soliton import Experiment
from gpe.soc_soliton import u, Experiment
e = Experiment()
s0 = e.get_state(Nxyz=(2**12/8,),
                 Lxyz=(1000*u.micron,), 
                 cooling_phase=1+0.000j,
                 constraint='N')
s0[0,...] = 1.0
s0[1,...] = np.sqrt(0.5)

m = minimize.MinimizeState(s0, fix_N=True)
assert m.check()
m.minimize(psi_tol=1e-7, tries=10, disp=1)
assert m.check()
s1 = m.minimize(psi_tol=1e-15, tries=1000, disp=1)
s1.plot()
print(s1.get_Ns())

E0 = s1.get_energy()

Here we trace the performance of both solvers:

from gpe import minimize;reload(minimize)
e = Experiment()
s0 = e.get_state(Nxyz=(2**12/8,),
                 Lxyz=(1000*u.micron,), 
                 cooling_phase=1+0.000j,
                 constraint='N')
s0[0,...] = 1.0
s0[1,...] = np.sqrt(0.5)

m0 = minimize.MinimizeState(s0, fix_N=True)
m0.minimize(psi_tol=1e-32, E_tol=1e-32, disp=2, use_scipy=True)
Es0 = np.array(m0._Es)
m0.minimize_results

from gpe import minimize;reload(minimize)
m1 = minimize.MinimizeState(s0, fix_N=True)
m1.minimize(psi_tol=1e-32, E_tol=1e-32, tries=10, disp=2)
Es1 = np.array(m1._Es0)

There seems to be a problem with the computation of the derivatives:

plt.semilogy(Es0-E0, '+')
plt.semilogy(Es1-E0, ':')
plt.semilogy(np.arange(len(Es1), len(Es1)+len(m1._Es)), m1._Es-E0, '--')
#len(Es1)

plt.semilogy(Es0-E0, '+')
plt.semilogy(Es1-E0, '.')
#len(Es1)

plt.semilogy(Es0-E0, '+')
plt.semilogy(Es1-E0, '.')
len(Es1)

reload(minimize)
m0 = minimize.MinimizeState(s0, fix_N=True)
np.random.seed(11)
x0 = np.random.random(m0.x.shape) - 0.5
m0.x = x0
s0 = m0.state
m0 = minimize.MinimizeState(s0, fix_N=True)
print(m0.state.get_N(), m0.state.get_energy())
#m0.f_df(m0.x)[0]
#m0.check()
x1 = np.random.random(m0.x.shape) - 0.5
s1 = m0.unpack(x1)
s1.normalize()
x2 = m0.pack(s1)
m0.check(x2)

It appears like the line search is bailing too early. We have instrumented to code to store the local state at the failing line_search call:

locals().update(m1._locals)
c1 = 1e-4
c2 = 0.1
ls_res = line_search(f, df, xk=x0, pk=dx0, old_fval=f0, old_old_fval=f00, c1=1e-4, c2=0.99)
alphas = np.linspace(0,1,100)
def phi(alpha):
    return f(x0 + alpha*dx0)
def dphi(alpha):
    return dx0.dot(df(x0 + alpha*dx0))
Es = map(phi, alphas)
plt.plot(alphas, Es)
plt.plot(alphas, phi(0) + dphi(0)*alphas)
plt.plot(alphas, f(x0) + c1*dx0.dot(df(x0))*alphas)

a0 = 0.6
plt.plot(alphas, phi(a0) + dphi(a0)*(alphas - a0), 'g-')
print dphi(a0), (phi(0.61) - phi(0.6))/0.01

x1 = x0 + a0*dx0
h = 0.0001
(f(x1 + h*dx0) - f(x1 - h*dx0))/2/h, dx0.dot(df(x1))

#plt.plot(alphas, f(x0(abs(derphi_a1) <= -c2*derphi0):
plt.axvline(ls_res[0], c='y')
#plt.twinx()
#plt.plot(alphas, [dx0.dot(df(x0+_a*dx0)) for _a in alphas])
#plt.axhline(c2*dx0.dot(df(x0)), c='y')

m1._alphas[-10:]

The line search is failing, but the function looks fine. Let’s explicitly check the wolf conditions.

line_search?

Old Notes#

Here are some notes about material that is no-longer relevant or important.

Minimizing the Residual Norm#

As we shall see below, we will also need to be able to minimize an auxilliary funtion:

\[ g(x) = (\psi\mat{H})^\dagger (\mat{H}\psi). \]

This requires a little care since \(\mat{H}[n]\) itself depends non-linearly on the \(\psi\) through the density \(n\). Let us consider a slightly more general non-linear energy density:

\[ \frac{g}{2}n(x)^2 = \mathcal{E}(n) \]

The only portion of \(\op{H}\) that depends on \(n\) is \([H]_{ij} = \cdots + \mathcal{E}'(n_j)\delta_{ij}\), so the derivative is:

\[ \pdiff{[H]_{ij}}{n_k} = g\delta_{ij}\delta_{jk}. \]

The resulting derivative contains the piece

\[ \pdiff{g}{\psi^\dagger_k} = [\mat{H}^\dagger\mat{H}\psi]_k + \sum_{ij}\psi^\dagger_i\mathcal{E}''(n_j)\psi_j\delta_{jk}\delta_{ij}[\mat{H}\psi]_j + \sum_{ij}[\psi^\dagger\mat{H}^\dagger]_i\mathcal{E}''(n_j)\psi_j\delta_{jk}\delta_{ij}\psi_j = [\mat{H}^\dagger\mat{H}\psi]_k + \psi^\dagger_k\mathcal{E}''(n_k)\psi_k[\mat{H}\psi]_k + [\psi^\dagger\mat{H}^\dagger]_k\mathcal{E}''(n_k)\psi_k\psi_k \]
\[ \pdiff{g}{\psi^\dagger_k} = [\mat{H}^\dagger\mat{H}\psi]_k + n_k\mathcal{E}''(n_k)\psi_k \]