---
jupytext:
  encoding: '# -*- coding: utf-8 -*-'
  formats: md:myst,ipynb
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.16.4
kernelspec:
  display_name: Python 3 (ipykernel)
  language: python
  name: python3
---

```{code-cell}
:init_cell: true
:tags: [hide-cell]

import mmf_setup; mmf_setup.nbinit()
import numpy as np
import matplotlib.pyplot as plt
```

(sec:sigmas)=
# Solving for Sigmas: Globally Convergent Newton's Method

To implement the NPSEQ for two components as described in {ref}`sec:NPSEQ` we need to
robustly solve some systems of non-linear equations to determine the values of $\sigma$
in the expansion direction.  Here we do this using a globally convergent Newton's method.

## Asymmetric Traps

If the perpendicular trapping frequencies are different, we need to solve the following
equations with positive coefficients:
\begin{gather*}
  B_x \sigma_x^{2} - \frac{A_x}{\sigma_x^{2}}
  = B_y \sigma_y^{2} - \frac{A_y}{\sigma_y^{2}}
  = f(n_0) = \frac{h(n_0)}{n_0}\\
  n_0 = \frac{n_1}{\pi \sigma_{x}\sigma_{y}},\qquad
  f(n_0) = Q'(n_0) - \frac{Q(n_0)}{n_0} = \frac{h(n_0)}{n_0}.
\end{gather*}
Our general strategy is as follows:
1. Solve both quadratic equations for $\sigma_x$ and $\sigma_y$ so we can obtain a
   single non-linear equation for $n_0 \propto 1/\sigma_x\sigma_y$.
2. Solve this with a globally-convergent Newton iteration.

After some manipulations, we end up with an equation of the following form
\begin{gather*}
  c \frac{q(n_0^2)}{n_0^2} = 
  \left(1 + \sqrt{1 + a_x q(n_0^2)}\right)
  \left(1 + \sqrt{1 + a_y q(n_0^2)}\right),\\
  q(n_0^2) = \frac{1}{f^2(n_0)}, \qquad
  c = \frac{4B_xB_y n_1^2}{\pi^2}, \qquad
  a_x = 4A_{x}B_{x}, \qquad
  a_y = 4A_{y}B_{y}.
\end{gather*}
Motivated by {ref}`sec:PowerLawEoS` where $n_0^2 = Cq^{1/(1-\beta)}$, we demonstrate in
section {ref}`sec:NMq` below a fast, globally convergent Newton method for solving the
following two-parameter equation:
\begin{gather*}
  \tilde{q}^{\beta/(\beta-1)} = \left(1+\sqrt{1+\tilde{a}_{x}\tilde{q}}\right)
                                \left(1+\sqrt{1+\tilde{a}_{\smash{y}}\tilde{q}}\right).
\end{gather*}
We can use this by applying heuristic to estimate parameters $C$ and $\beta$, then define
\begin{gather*}
  \tilde{q} = \left(\frac{c}{C}\right)^{(\beta-1)/\beta} q, \qquad
  \tilde{a}_{i}\tilde{q} = a_i q, \qquad
  \tilde{a}_i = \left(\frac{c}{C}\right)^{(\beta-1)/\beta}a_{i}.
\end{gather*}
We apply the {ref}`sec:NMq` to our original equation, but as a function of the parameter
\begin{gather*}
  \tilde{q} = \left(\frac{c}{C}\right)^{(\beta-1)/\beta}
  \overbrace{\left(\frac{n_0^2}{C}\right)^{1-\beta}}^{\approx q},\qquad
  n_0^2 %= C \left(\left(\frac{C}{c}\right)^{(\beta - 1)/\beta}\tilde{q}\right)^{1/(1-\beta)}
        = C \left(\frac{c}{C}\right)^{1/\beta}\tilde{q}^{1/(1-\beta)},
\end{gather*}
including the appropriate factors through the chain-rule.









this equation.  In the general case, we use the same technique, but apply the chain rule

This can be expressed in terms of two parameters by defining 

We apply Newton's method to this in the form of
\begin{gather*}
  g(n_0^2) = \ln c\frac{q(n_0^2)}{n_0^2}
  - \sum_{i=x,y}\ln\left(1 + \sqrt{1 + a_i q(n_0^2)}\right),\\
  g'(n_0^2) =
  \frac{q'(n_0^2)}{q(n_0^2)}
  - \frac{1}{n_0^2}
  - \frac{q'(n_0^2)}{2}
    \sum_{i={x,y}}\frac{a_i}{1 + a_i q(n_0^2) + \sqrt{1 + a_i q(n_0^2)}},
\end{gather*}
We thus need a way to compute
\begin{gather*}
  \frac{q(n_0^2)}{n_0^2} = \frac{n_0^2}{f^2(n_0)} = \frac{1}{h^2(n_0)},\qquad
  q'(n_0^2) = \frac{-f'(n_0)}{n_0f^3(n_0)},\\
  \frac{q'(n_0^2)}{q(n_0^2)} - \frac{1}{n_0^2}
  = \frac{-f'(n_0)}{n_0f^2(n_0)} - \frac{1}{n_0^2}
  = \frac{-h'(n_0)}{n_0h(n_0)},\\
  f(n_0) = Q'(n_0) - \frac{Q(n_0)}{n_0}, \qquad
  f'(n_0) = Q''(n_0) - \frac{Q'(n_0)}{n_0} + \frac{Q(n_0)}{n_0^2},\\
  h(n_0) = n_0Q'(n_0) - Q(n_0), \qquad h'(n_0) = n_0Q''(n_0).
\end{gather*}

```{code-cell}
:tags: [hide-cell]

from sympy import Function, var, simplify
n = var('n_0')
f = Function('f')(n)
h = Function('h')(n)
q = 1/f**2
dq = simplify(q.diff(n)/2/n)
display(dq)
display(simplify(dq/q-1/n**2))

q = n**2/h**2
dq = simplify(q.diff(n)/2/n)
display(dq)
display(simplify(dq/q-1/n**2))
```

(sec:PowerLawEoS)=
### Power-law Equations of State

If the equation of state is a power-law, then we have some simplifications (with
$\beta=2$ being the GPE):
\begin{gather*}
  \mathcal{E}(n) = \frac{g n^{\beta}}{\beta},\\
  Q(n) = \frac{g n^{\beta}}{\beta^2}, \quad
  Q'(n) = \frac{\mathcal{E}(n)}{n} = \frac{g n^{\beta-1}}{\beta}, \quad
  Q''(n) =  \frac{g (\beta-1)n^{\beta-2}}{\beta},\\
  f(n) = \frac{g (\beta-1)n^{\beta+1}}{\beta^2},\qquad
  f'(n) = \frac{g (\beta^2-1)n^{\beta}}{\beta^2},\\
  h(n) = \frac{g (\beta-1)n^{\beta}}{\beta^2},\qquad
  h'(n) = \frac{g (\beta-1)n^{\beta-1}}{\beta},\\
  q(n^2) = \frac{\beta^4}{g^2 (\beta-1)^2(n^2)^{\beta-1}}, \qquad
  n^2(q) = \left(\frac{g^2 (\beta-1)^2}{\beta^4}q\right)^{-1/(\beta-1)}
         = C q^{1/(1-\beta)}.
\end{gather*}
The equation we need to solve can thus be expressed as:
\begin{gather*}
  g(n_0^2) = \overbrace{\ln c \frac{q(n_0^2)}{n_0^2}}^{\ln \frac{c}{C}q^{\beta/(\beta-1)}}
           - \sum_{i=x,y}\ln\left( 1 + \sqrt{1 + a_i q(n_0^2)}\right),\\
  g(\tilde{q}) = \frac{\beta}{\beta-1}\ln \tilde{q} - \sum_{i=x,y}\ln\left(
    1 + \sqrt{1 + \tilde{a}_i \tilde{q}}
  \right),\\
  \tilde{q} = \left(\frac{c}{C}\right)^{(\beta-1)/\beta} q, \qquad
  \tilde{a}_{i}\tilde{q} = a_i q, \qquad
  \tilde{a}_i = \left(\frac{c}{C}\right)^{(\beta-1)/\beta}a_{i}
\end{gather*}
In {ref}`sec:NMq` below, we find a fast, globally convergent Newton method for solving
this equation.


In the general case, we use the same technique, but apply the chain rule




To analyze the convergence, we can write the Newton iteration in terms of two
parameters:
\begin{gather*}
  \tilde{q}^{\beta/(\beta-1)} = (1+\sqrt{1+\tilde{a}_x\tilde{q}})(1+\sqrt{1+\tilde{a}_y\tilde{q}}).
\end{gather*}

\begin{gather*}
  g(n_0^2) = \overbrace{\ln c \frac{q(n_0^2)}{n_0^2}}^{\ln \frac{c}{C}q^{\beta/(\beta-1)}}
           - \sum_{i=x,y}\ln\left( 1 + \sqrt{1 + a_i q(n_0^2)}\right),\\
  g(\tilde{q}) = \frac{\beta}{\beta-1}\ln \tilde{q} - \sum_{i=x,y}\ln\left(
    1 + \sqrt{1 + \tilde{a}_i \tilde{q}}
  \right),\\
  \tilde{q} = \left(\frac{c}{C}\right)^{(\beta-1)/\beta} q, \qquad
  \tilde{a}_{i}\tilde{q} = a_i q, \qquad
  \tilde{a}_i = \left(\frac{c}{C}\right)^{(\beta-1)/\beta}a_{i}
\end{gather*}
In {ref}`sec:NMq` below, we find a fast, globally convergent Newton method for solving
this equation.  In the general case, we use the same technique, but apply the chain rule
to
\begin{gather*}
  \tilde{q}(n_0^2) = \left(\frac{c}{C}\right)^{(\beta-1)/\beta} q(n_0^2)
\end{gather*}
where the values $C$ and $\beta$ are determined heuristically by evaluating $q(n_0^2)$
and $q'(n_0^2)$ at a few fiducial densities:
\begin{gather*}
  q(n_0^2) \approx \left(\frac{n_0^2}{C}\right)^{1-\beta}, \qquad
  q'(n_0^2) \approx \frac{(1-\beta)}{C}\left(\frac{n_0^2}{C}\right)^{-\beta}\\
  \beta \approx 1-\log_{s} \frac{q(s n_0^2)}{q(n_0^2)}
  \approx -\log_{s} \frac{q'(s n_0^2)}{q'(n_0^2)},\qquad
  C = \frac{n_0^2}{\Bigl(q(n_0^2)\Bigr)^{1/(1-\beta)}}.
\end{gather*}


### Numerical Demonstration

Before we go into too many details, we setup an example.

```{code-cell}
def Q(n, d=0, g=1.0, beta=2.0, **kw):
    """Return Q(n) and its derivatives where  Q'(n) = E/N, the energy per particle."""
    if d == 0:
        return g*n**beta/2/beta
    elif d == 1:
        return g*n**(beta-1)/2
    elif d == 2:
        return g*(beta-1)*n**(beta-2)/2

def f(n, d=0, **kw):
    if d == 0:
        return Q(n, d=1, **kw) - Q(n, d=0, **kw)/n
    elif d == 1:
        return Q(n, d=2, **kw) - Q(n, d=1, **kw)/n + Q(n, d=0, **kw)/n**2

def h(n, d=0, **kw):
    if d == 0:
        return n*Q(n, d=1, **kw) - Q(n, d=0, **kw)
    elif d == 1:
        return n*Q(n, d=2, **kw)

def q(n2, d=0, **kw):
    """Return q(n2) or its derivative."""
    n = np.sqrt(n2)
    if d == 0:
        return 1 / f(n, d=0, **kw)**2
    elif d == 1:
        return -f(n, d=1, **kw) / f(n, d=0, **kw)**3/n

def newton_g(n2, ax, ay, d=0, c=1.0, q=q, **kw):
    _q = q(n2, d=0, **kw)
    _dq = q(n2, d=1, **kw)
    _ax, _ay = 1 + ax*_q, 1 + ay*_q
    _sqrt_ax, _sqrt_ay = np.sqrt(_ax), np.sqrt(_ay)
    _x, _y = 1 + _sqrt_ax, 1 + _sqrt_ay
    if d == 0:
        g = np.log(c*_q/n2) - np.log(_x*_y)
        return g
    elif d == 1:
        dg = _dq/_q - 1/n2 - _dq/2*(ax/(_ax + _sqrt_ax) + ay/(_ay + _sqrt_ay))
        return dg
    raise NotImplementedError(f"{d=}")

def newton_step(n2, ax, ay, beta=2.0, **kw):
    kw.update(ax=ax, ay=ay, beta=beta)
    return n2 - newton_g(n2, **kw)/newton_g(n2, d=1, **kw)
```

```{code-cell}
:tags: [hide-input]

# Some checks
from functools import partial
from mmfutils.math.differentiate import differentiate
rng = np.random.default_rng(seed=2)
kw = dict(zip(['ax', 'ay', 'c', 'beta', 'g'], 1+rng.random(5)))

n = rng.random()
def dQ(n, d=0, **kw):
    return Q(n, d=d+1, **kw)
    
for _f in [Q, dQ, f, h, q, newton_g]:
    _f_ = partial(_f, **kw)
    assert np.allclose(differentiate(_f_, n, h0=0.01), _f_(n, d=1))
```

```{code-cell}
def get_beta_C(q, n0=1.0, s=1.1, **kw):
    """Return estimates for `(beta, C)`."""
    n2 = n0**2
    q0, dq0 = q(n2, **kw), q(n2, d=1, **kw)
    qs0, dqs0 = q(s*n2, **kw), q(s*n2, d=1, **kw)
    betas = [1 - np.log(qs0/q0)/np.log(s), 
             - np.log(dqs0/dq0)/np.log(s)]
    beta = np.mean(betas)
    dbeta = np.std(betas)
    assert (abs(dbeta / beta) < 0.1)
    C = n2 / q0 ** (1 / (1-beta))
    return (beta, C)

def get_q0(ax, ay, beta=None):
    """Return the initial guess `q0`.
    
    Note: these are the tilde version in the notes.
    """
    power = beta/(beta-1)
    rhs = 1.0
    for _a in [ax, ay]:
        lam = _a**2/(1+_a)**2  # Convex interpolation parameter
        rhs *= 1 + np.sqrt(1 + _a)
        power -= 0.5*lam
    q0 = rhs**(1/power)
    return q0

def newton_tilde(qt, ax, ay, c, q=q, **kw):
    """Return `(qt, g)`"""
    kw.update(ax=ax, ay=ay, c=c)
    beta, C = get_beta_C(q, **kw)
    _tmp = C*(c/C)**(1/beta)
    _power = 1/(1-beta)
    n2 = _tmp * qt**_power
    dn2_dqt = _power * n2 / qt
    g = newton_g(n2, **kw)
    dg = newton_g(n2, d=1, **kw) * dn2_dqt
    return qt - g/dg, g

@np.vectorize
def get_qt0(ax, ay, c, q=q, beta=None, C=None, **kw):
    """Return the initial guess `qtilde`."""
    if beta is None or C is None:
        beta, C = get_beta_C(q, ax=ax, ay=ay, **kw)
    _tmp = (c/C)**((beta-1)/beta)
    ax_, ay_ = _tmp * np.array([ax, ay])
    qt0 = get_q0(ax=ax_, ay=ay_, beta=beta)
    return qt0

@np.vectorize
def count(ax, ay, qt, nmax=10, **kw):
    kw.update(ax=ax, ay=ay)
    for n in range(1, nmax+1):
        qt, g = newton_tilde(qt, **kw)
        if abs(g) / max([1.0, abs(qt)]) < 1e-12:
            break
    return n

def plot_region(ax, ay, get_qt0=get_qt0, nmax=10, **kw):
    ax_ = np.linspace(ax[0], ax[1], 101)[:, None]
    ay_ = np.linspace(ay[0], ay[1], 101)[None, :]
    cs = count(ax_, ay_, qt=get_qt0(ax_,ay_, **kw), nmax=nmax, **kw)
    fig, ax = plt.subplots()
    mesh = ax.pcolormesh(ax_.ravel(), ay_.ravel(), cs.T)
    fig.colorbar(mesh, ax=ax, label="steps")
    ax.set(xlabel="$a_x$", ylabel="$a_y$", title=get_qt0.__doc__)

kw.pop('ax')
kw.pop('ay')
kw['beta'] = 2
plot_region(ax=(0.01,100.0), ay=(0.01,100.0), **kw)
```

If we assume that $q(n_0^2)$ is monotonic, then we know that $q(n_0^2) \geq 

(sec:NMq)=
### Globally Convergent Newton's Method
Here we demonstrate a globally convergent Newton method for the simplified equation
\begin{gather*}
  q^{\beta/(\beta-1)} = \bigl(1 + \sqrt{1 + a_x q}\bigr)
                        \bigl(1 + \sqrt{1 + a_y q}\bigr)
\end{gather*}
where $a_{x,y} \geq 0$.
We start with the log of this equation and try Newton's method for
\begin{gather*}
  g(q) = \frac{\beta}{\beta-1} \ln q - \ln\bigl(1 + \sqrt{1 + a_x q}\bigr)
  - \ln\bigl(1 + \sqrt{1 + a_y q}\bigr),\\
  g'(q) = 
  \frac{\beta}{\beta-1} \frac{1}{q} 
  - \frac{a_x/2}{1 + a_x q + \sqrt{1 + a_x q}}
  - \frac{a_y/2}{1 + a_y q + \sqrt{1 + a_y q}},\\
  f(q) = q - \frac{g(q)}{g'(q)}.
\end{gather*}

```{code-cell}
def g(q, ax, ay, beta, d=0):
    _x = 1 + np.sqrt(1+ax*q)
    _y = 1 + np.sqrt(1+ay*q)
    if d == 0:
        g = beta/(beta-1) * np.log(q) - np.log(_x*_y)
        return g
    elif d == 1:
        dg = beta / (beta-1) / q - ax/2/(1+ax*q + np.sqrt(1+ax*q)) - ay/2/(1+ay*q + np.sqrt(1+ay*q))
        return dg
    raise NotImplementedError(f"{d=}")

def f(q, ax, ay, beta):
    kw = dict(ax=ax, ay=ay, beta=beta)
    return q - g(q, **kw)/g(q, d=1, **kw)
```

As an initial guess, we know that $q\geq 1$ so we might start with
\begin{gather*}
  q_0^{\beta/(\beta-1)} = \bigl(1+\sqrt{1+a_x}\bigr)\bigl(1+\sqrt{1+a_y}\bigr).
\end{gather*}

```{code-cell}
from mmfutils.plot import imcontourf

beta = 2.0

def get_q0(ax, ay):
    r"$q_0^{\beta/(\beta-1)} = (1+\sqrt{1+a_x})(1+\sqrt{1+a_y})$"
    global beta
    return ((1+np.sqrt(1+ax))*(1+np.sqrt(1+ay)))**(1-1/beta)
    
@np.vectorize
def count(ax, ay, q0=1.0, nmax=10):
    global beta
    q = q0
    kw = dict(ax=ax, ay=ay, beta=beta)
    for n in range(1, nmax+1):
        q = f(q, **kw)
        g_ = g(q, **kw)
        if abs(g_) / max(1, abs(q0)) < 1e-12:
            break
    return n

def plot_region(ax, ay, get_q0=get_q0, nmax=10):
    ax_ = np.linspace(ax[0], ax[1], 101)[:, None]
    ay_ = np.linspace(ay[0], ay[1], 101)[None, :]
    cs = count(ax_, ay_, q0=get_q0(ax_,ay_), nmax=nmax)
    fig, ax = plt.subplots()
    mesh = ax.pcolormesh(ax_.ravel(), ay_.ravel(), cs.T)
    fig.colorbar(mesh, ax=ax, label="steps")
    ax.set(xlabel="$a_x$", ylabel="$a_y$", title=get_q0.__doc__)
```

```{code-cell}
%time plot_region(ax=(0.01,100.0), ay=(0.01,100.0))
```

Here we see that Newton's method converges quite quickly, but we need a better guess
than $q_0 = 1$ for large values of $a_{x,y}$.  A better guess might be
\begin{gather*}
  q = \sqrt{a_xa_y}^{\beta-1}\\
\end{gather*}

```{code-cell}
def get_q0b(ax, ay):
    r"$q_0=\sqrt{a_xa_y}^{\beta - 1}$"
    global mu
    return np.sqrt(ax*ay)**(beta-1)
    
plot_region(ax=(0.01,100.0), ay=(0.01,100.0), get_q0=get_q0b)
```

:::{margin}
We played a bit with the form for the power $p$.
:::
Here we combine these two methods:
\begin{gather*}
  q^p = \bigl(1+\sqrt{1+a_x}\bigr)\bigl(1+\sqrt{1+a_y}\bigr), \\
  p = \frac{\beta}{\beta-1} - \frac{a_x^2}{2(1+a_x)^2} - \frac{a_y^2}{2(1+a_y)^2}.
\end{gather*}

```{code-cell}
beta = 2.0  # GPE

@np.vectorize
def get_q0c(ax, ay):
    r"combined"
    global mu
    power = beta/(beta-1)
    rhs = 1.0
    for _a in [ax, ay]:
        lam = _a**2/(1+_a)**2  # Convex interpolation parameter
        rhs *= 1 + np.sqrt(1 + _a)
        power -= 0.5*lam
    return rhs**(1/power)

plot_region(ax=(0.01,100.0), ay=(0.01,100.0), get_q0=get_q0c)
```

This appears to uniformly converge in 4 Newton steps for all parameter regions for
$\beta = 2$ (GPE).  We also consider the UFG with $\beta = 5/2$

```{code-cell}
beta = 5/2  # UFG
plot_region(ax=(0.01,100.0), ay=(0.01,100.0), get_q0=get_q0c)
```

This requires one more iteration.







## Old:
:::{margin}
For reference
\begin{gather*}
  A_{x,y} = \frac{\hbar^2}{2m\lambda_{x,y}^2}\\
  B_{x,y} = \frac{m\omega_{x, y}^2}{2},\\
  C = \frac{\alpha}{a\beta^2} \left(\frac{n_1}{\pi}\right)^{\beta-1},
\end{gather*}
where the value of $a$ depends on whether we minimize the energy or chemical potential:
\begin{gather*}
  a = \begin{cases}
    \frac{1}{\beta-1} & \text{(energy)},\\
    \frac{1}{\beta(\beta - 1)} & \text{(chemical potential)},
  \end{cases}
\end{gather*}
For the GPE, $\beta = 2$ and $\alpha = g$.
:::
If the perpendicular trapping frequencies are different, we need to solve the following
equations with positive coefficients:
\begin{gather*}
  B_x \sigma_x^{2} - \frac{A_x}{\sigma_x^{2}}
  = B_y \sigma_y^{2} - \frac{A_y}{\sigma_y^{2}}
  = \underbrace{f\Biggl(
    \overbrace{\frac{n_1}{\pi \sigma_x\sigma_y}}^{n_0}
  \Biggr)}_{Q'(n_0) - \frac{Q(n_0)}{n_0}}
\end{gather*}
Our general strategy is as follows:
1. Solve both quadratic equations for $\sigma_x$ and $\sigma_y$ so we can obtain a
   single non-linear equation for $n_0 \propto 1/\sigma_x\sigma_y$.
2. Solve this with a globally-convergent Newton iteration.

After some manipulations, we end up with an equation of the following form
\begin{gather*}
  c\frac{q(n_0^2)}{n_0^2} = 
  \left(1 + \sqrt{1 + a_x q(n_0^2)}\right)
  \left(1 + \sqrt{1 + a_y q(n_0^2)}\right),\\
  q(n_0^2) = \frac{1}{f^2(n_0)}, \qquad
  c = \frac{4B_xB_y n_1^2}{\pi^2}, \qquad
  a_x = 4A_{x}B_{x}, \qquad
  a_y = 4A_{y}B_{y}.
\end{gather*}
We apply Newton's method to this in the form of
\begin{gather*}
  g(n_0^2) = \ln c + \ln\frac{n_1^2}{n_0^2} + \ln q(n_0^2) 
  - \ln\left(1 + \sqrt{1 + a_x q(n_0^2)}\right)
  - \ln\left(1 + \sqrt{1 + a_y q(n_0^2)}\right),\\
  g'(n_0^2) = -\frac{1}{n_0^2} + \frac{q'(n_0^2)}{q(n_0^2)}
  - \frac{a_x q'(n_0^2)/2}{1 + a_x q(n_0^2) + \sqrt{1 + a_x q(n_0^2)}}
  - \frac{a_y q'(n_0^2)/2}{1 + a_y q(n_0^2) + \sqrt{1 + a_y q(n_0^2)}}.
\end{gather*}
We thus need a way to compute
\begin{gather*}
  q(n_0^2) = \frac{1}{f^2(n_0)}, \qquad
  q'(n_0^2) = \frac{f'(n_0)}{4n_0f^3(n_0)}, \qquad
  f'(n_0) = Q''(n_0) - \frac{Q'(n_0)}{n_0} + \frac{Q(n_0)}{n_0^2}
\end{gather*}

\begin{gather*}
  h(n_0) = n_0f(n_0)\\
  h'(n_0) = n_0f'(n_0) + f(n_0)
  = n_0Q''(n_0)
\end{gather*}








:::::{admonition} Details
:class: dropdown

\begin{gather*}
  \frac{2B\sigma^2}{f} = 1 + \sqrt{1 + \frac{4AB}{f^2}}\\
  \frac{4B_{x}B_{y}(\sigma_x\sigma_y)^2}{f^2(n_0)} 
  = \left(1 + \sqrt{1 + \frac{4A_{x}B_{x}}{f^2(n_0)}}\right)
    \left(1 + \sqrt{1 + \frac{4A_{y}B_{y}}{f^2(n_0)}}\right),\\
  \frac{4B_{x}B_{y}n_1^2}{\pi^2n_0^2f^2(n_0)} 
  = \left(1 + \sqrt{1 + \frac{4A_{x}B_{x}}{f^2(n_0)}}\right)
    \left(1 + \sqrt{1 + \frac{4A_{y}B_{y}}{f^2(n_0)}}\right),\\
  \frac{a}{n_0^2f^2(n_0)} = 
  \left(1 + \sqrt{1 + \frac{a_x}{f^2(n_0)}}\right)
  \left(1 + \sqrt{1 + \frac{a_y}{f^2(n_0)}}\right),\\
  \frac{q(n_0^2)}{n_0^2}
  = 
  \left(1 + \sqrt{1 + a_xq(n_0)}\right)
  \left(1 + \sqrt{1 + a_yq(n_0)}\right),\\
  q(n_0^2) = \frac{4B_{x}B_{y}}{\pi^2f^2(n_0)}, \qquad
  a_x = \frac{\pi^2 A_{x}}{B_{y}}, \qquad
  a_y = \frac{\pi^2 A_{y}}{B_{x}}, \qquad
\end{gather*}
:::::





For power-law equations of state, we have
\begin{gather*}
  Q'(n_0) - \frac{Q(n_0)}{n_0} = \frac{C}{(\sigma_{x}\sigma_{y})^{\beta-1}}.
\end{gather*}

:::{admonition} Symmetric GPE
:class: dropdown

As a check, consider the symmetric GPE ($\beta=2$) with $\sigma_x = \sigma_y = \sigma$,
$A = A_x = A_y$,  and $B_x = B_y = B$:
\begin{gather*}
  B\sigma^4 = A + C, \qquad
  \frac{m\omega_\perp^2}{2}\sigma^4 = \frac{\hbar^2}{2m} + \frac{g n_1}{2^{n}\pi}.
\end{gather*}
:::
To reduce this to a single equation, we define
\begin{gather*}
  \bar{\sigma}^2 = \sigma_x\sigma_y.
\end{gather*}
:::{margin}
Let $f = C/\bar{\sigma}^{2(\beta-1)}$ be fixed, then we have a quadratic equation
\begin{gather*}
  B_{x,y} \sigma_{x,y}^{4} - f\sigma_{x,y}^{2} - A_{x,y} = 0,\\
  \sigma_{x,y}^2 = \frac{f \pm \sqrt{f^2+4A_{x,y}B_{x,y}}}{2B_{x,y}}.
\end{gather*}
:::
Holding this fixed, we can then solve for
\begin{gather*}
  \frac{2B_{x,y}\sigma_{x,y}^2}{C}\bar{\sigma}^{2(\beta-1)} 
  = 1 + \sqrt{1 + \frac{4A_{x,y}B_{x,y}\bar{\sigma}^{4(\beta-1)}}{C^2}}.
\end{gather*}
Multiplying the two equations, we have a single equation for $\bar{\sigma}^2$:
\begin{gather*}
  \frac{4B_{x}B_{y}\bar{\sigma}^{4\beta}}{C^2}
  = \left(1 + \sqrt{1 + \frac{4A_{x}B_{x}\bar{\sigma}^{4(\beta-1)}}{C^2}}\right)
    \left(1 + \sqrt{1 + \frac{4A_{y}B_{y}\bar{\sigma}^{4(\beta-1)}}{C^2}}\right).
\end{gather*}
:::{note}
In the limit of small $C^2$ -- i.e. where the density vanishes -- we have
\begin{gather*}
  \bar{\sigma}^{4} = \sqrt{\frac{A_{x}A_{y}}{B_{x}B_{y}}}
\end{gather*}
:::
:::{admonition} Another approach.
:class: dropdown

Another approach to reducing this from two variables to one is to use the following
relationship between $\sigma_{x}^2$ and $\sigma_y^{2}$:
\begin{gather*}
  \frac{A_{x}}{\sigma_x^{2}} - B_x \sigma_x^{2}
  = \frac{A_{y}}{\sigma_y^{2}} - B_y \sigma_y^{2},\\
  B_x \sigma_x^{4} + \left(\frac{A_{y}}{\sigma_y^{2}} - B_y
  \sigma_y^{2}\right)\sigma_x^2 - A_{x} = 0.
\end{gather*}
This is simply a quadratic equation that can be solved reducing the problem to a single
variable, but lead to complicated equations.
:::

This can be simplified by introducing the variable $q$ and two positive parameters
$a_{x}$ and $a_{y}$
\begin{gather*}
  q^{\beta/(\beta-1)} = \bigl(1 + \sqrt{1 + a_x q}\bigr)
                        \bigl(1 + \sqrt{1 + a_y q}\bigr)
\end{gather*}
where
\begin{gather*}
  q^{\beta/(\beta-1)} = \frac{4B_{x}B_{y}\bar{\sigma}^{4\beta}}{C^2},\\
  \bar{\sigma}^{4(\beta-1)} = q\left(\frac{C^2}{4B_{x}B_{y}}\right)^{(\beta-1)/\beta},\\
  a_{x,y} = \frac{4A_{x,y}B_{x,y}}{C^2}\left(\frac{C^2}{4B_{x}B_{y}}\right)^{(\beta-1)/\beta}
          = \frac{A_{x,y}B_{x,y}}{B_{x}B_{y}}
          \left(\frac{4B_{x}B_{y}}{C^2}\right)^{1/\beta}.
\end{gather*}
Once we find $q$, we invert to find $\bar{\sigma}$ and $f = C/\bar{\sigma}^{2(\beta-1)}$
and then solve for $\sigma_x$ and $\sigma_y$:
\begin{gather*}
  \bar{\sigma}^{4\beta} = \frac{C^2}{4B_{x}B_{y}}q^{\beta/(\beta-1)}, \qquad
  f = \frac{C}{\bar{\sigma}^{2(\beta-1)}}, \qquad
  \sigma_{x,y}^2 = \frac{f + \sqrt{f^2+4A_{x,y}B_{x,y}}}{2B_{x,y}}.
\end{gather*}

### Old Globally Convergent Newton's Method
Here we demonstrate a globally convergent Newton method for the simplified equation
\begin{gather*}
  q^{\beta/(\beta-1)} = \bigl(1 + \sqrt{1 + a_x q}\bigr)
                        \bigl(1 + \sqrt{1 + a_y q}\bigr)
\end{gather*}
where
\begin{gather*}
  q^{\beta/(\beta-1)} = \frac{4B_{x}B_{y}\bar{\sigma}^{4\beta}}{C^2},\\
  \bar{\sigma}^{4(\beta-1)} = q\left(\frac{C^2}{4B_{x}B_{y}}\right)^{(\beta-1)/\beta},\\
  a_{x,y} = \frac{4A_{x,y}B_{x,y}}{C^2}\left(\frac{C^2}{4B_{x}B_{y}}\right)^{(\beta-1)/\beta}
          = \frac{A_{x,y}B_{x,y}}{B_{x}B_{y}}
          \left(\frac{4B_{x}B_{y}}{C^2}\right)^{1/\beta}.
\end{gather*}
Once we find $q$, we invert to find $\bar{\sigma}$ and $f = C/\bar{\sigma}^{2(\beta-1)}$
and then solve for $\sigma_x$ and $\sigma_y$:
\begin{gather*}
  \bar{\sigma}^{4\beta} = \frac{C^2}{4B_{x}B_{y}}q^{\beta/(\beta-1)}, \qquad
  f = \frac{C}{\bar{\sigma}^{2(\beta-1)}}, \qquad
  \sigma_{x,y}^2 = \frac{f + \sqrt{f^2+4A_{x,y}B_{x,y}}}{2B_{x,y}}.
\end{gather*}


We start with the log of this equation and try Newton's method for
\begin{gather*}
  g(q) = \frac{\beta}{\beta-1} \ln q - \ln\bigl(1 + \sqrt{1 + a_x q}\bigr)
  - \ln\bigl(1 + \sqrt{1 + a_y q}\bigr),\\
  g'(q) = 
  \frac{\beta}{\beta-1} \frac{1}{q} 
  - \frac{a_x/2}{1 + a_x q + \sqrt{1 + a_x q}}
  - \frac{a_y/2}{1 + a_y q + \sqrt{1 + a_y q}},\\
  f(q) = q - \frac{g(q)}{g'(q)}.
\end{gather*}

```{code-cell}
def g(q, ax, ay, beta, d=0):
    _x = 1 + np.sqrt(1+ax*q)
    _y = 1 + np.sqrt(1+ay*q)
    if d == 0:
        g = beta/(beta-1) * np.log(q) - np.log(_x*_y)
        return g
    elif d == 1:
        dg = beta / (beta-1) / q - ax/2/(1+ax*q + np.sqrt(1+ax*q)) - ay/2/(1+ay*q + np.sqrt(1+ay*q))
        return dg
    raise NotImplementedError(f"{d=}")

def f(q, ax, ay, beta):
    kw = dict(ax=ax, ay=ay, beta=beta)
    return q - g(q, **kw)/g(q, d=1, **kw)
```

As an initial guess, we know that $q\geq 1$ so we might start with
\begin{gather*}
  q_0^{\beta/(\beta-1)} = \bigl(1+\sqrt{1+a_x}\bigr)\bigl(1+\sqrt{1+a_y}\bigr).
\end{gather*}

```{code-cell}
from mmfutils.plot import imcontourf

beta = 2.0

def get_q0(ax, ay):
    r"$q_0^{\beta/(\beta-1)} = (1+\sqrt{1+a_x})(1+\sqrt{1+a_y})$"
    global beta
    return ((1+np.sqrt(1+ax))*(1+np.sqrt(1+ay)))**(1-1/beta)
    
@np.vectorize
def count(ax, ay, q0=1.0, nmax=10):
    global beta
    q = q0
    kw = dict(ax=ax, ay=ay, beta=beta)
    for n in range(1, nmax+1):
        q = f(q, **kw)
        g_ = g(q, **kw)
        if abs(g_) / max(1, abs(q0)) < 1e-12:
            break
    return n

def plot_region(ax, ay, get_q0=get_q0, nmax=10):
    ax_ = np.linspace(ax[0], ax[1], 101)[:, None]
    ay_ = np.linspace(ay[0], ay[1], 101)[None, :]
    cs = count(ax_, ay_, q0=get_q0(ax_,ay_), nmax=nmax)
    fig, ax = plt.subplots()
    mesh = ax.pcolormesh(ax_.ravel(), ay_.ravel(), cs.T)
    fig.colorbar(mesh, ax=ax, label="steps")
    ax.set(xlabel="$a_x$", ylabel="$a_y$", title=get_q0.__doc__)
```

```{code-cell}
%time plot_region(ax=(0.01,100.0), ay=(0.01,100.0))
```

Here we see that Newton's method converges quite quickly, but we need a better guess
than $q_0 = 1$ for large values of $a_{x,y}$.  A better guess might be
\begin{gather*}
  q = \sqrt{a_xa_y}^{\beta-1}\\
\end{gather*}

```{code-cell}
def get_q0b(ax, ay):
    r"$q_0=\sqrt{a_xa_y}^{\beta - 1}$"
    global mu
    return np.sqrt(ax*ay)**(beta-1)
    
plot_region(ax=(0.01,100.0), ay=(0.01,100.0), get_q0=get_q0b)
```

:::{margin}
We played a bit with the form for the power $p$.
:::
Here we combine these two methods:
\begin{gather*}
  q^p = \bigl(1+\sqrt{1+a_x}\bigr)\bigl(1+\sqrt{1+a_y}\bigr), \\
  p = \frac{\beta}{\beta-1} - \frac{a_x^2}{2(1+a_x)^2} - \frac{a_y^2}{2(1+a_y)^2}.
\end{gather*}

```{code-cell}
beta = 2.0  # GPE

@np.vectorize
def get_q0c(ax, ay):
    r"combined"
    global mu
    power = beta/(beta-1)
    rhs = 1.0
    for _a in [ax, ay]:
        lam = _a**2/(1+_a)**2  # Convex interpolation parameter
        rhs *= 1 + np.sqrt(1 + _a)
        power -= 0.5*lam
    return rhs**(1/power)

plot_region(ax=(0.01,100.0), ay=(0.01,100.0), get_q0=get_q0c)
```

This appears to uniformly converge in 4 Newton steps for all parameter regions for
$\beta = 2$ (GPE).  We also consider the UFG with $\beta = 5/2$

```{code-cell}
beta = 5/2  # UFG
plot_region(ax=(0.01,100.0), ay=(0.01,100.0), get_q0=get_q0c)
```

This requires one more iteration.

## Two Components

For the two component systems, we need to solve

\begin{gather*}
  \sigma_a^4 = B_{a} + \frac{C_{a}\sigma_a^4}{(\sigma_a^2 + \sigma_b^2)^2} = B_{a} +
  \frac{C_{a}}{\left(1 + \frac{\sigma_b^2}{\sigma_a^2}\right)^2}, \\
  B_{a} = \frac{\hbar^2}{m^2\omega_\perp^2} + \frac{g_{aa}\abs{\psi_a}^2}{2\pi
  m\omega_\perp^2}, \qquad
  C_{a} = \frac{2 g_{ab}\abs{\psi_b}^2}{\pi m\omega_\perp^2}
\end{gather*}

and $a \leftrightarrow b$. These can be rearranged into a single equation for the
fraction $x$:

\begin{gather*}
 x = \frac{\sigma_a^2}{\sigma_b^2}, \qquad
 a = \frac{B_a}{B_b}, \qquad
 b = \frac{C_a - C_b}{B_b}, \\
 x^4 + 2x^3 + x^2 = a(1+x)^2 + b x^2 = 0, \qquad
 (1+x)^2(x^2-a) = b x^2.
\end{gather*}

Here $a\approx 1$ and for physical systems $\abs{b} < 1$, but for testing, we might have
larger values.  We can take $b>0$ by flipping $a\leftrightarrow b$  $C_a < C_b$.

The problem is now reduced to finding a good strategy for finding the largest positive
root of this polynomial.  Our approach is to use Newton's method with an appropriate
initial guess to ensure rapid convergence.  A good way to do this, is to make a function
that counts the number of iterations required to converge.

```{code-cell}
x = np.linspace(0, 3, 100)
a = 1.0
b = 10.0
plt.plot(x, a + b * x ** 2 / (1 + x) ** 2 - x ** 2)
```

```{code-cell}
x = np.linspace(0, np.sqrt(11), 100)
a = 0.01
b = 0.1
plt.plot(x, 1.0 / (a / x ** 2 + b / (1 + x) ** 2) - 1)
```

```{code-cell}
x = np.linspace(0, 2, 100)
a = 0.01
b = 3.0
plt.plot(x, np.sqrt(a * (1 + x) ** 2 + b * x ** 2) - (1 + x) * x)
```

```{code-cell}
x = np.linspace(0, 2, 100)
a = 0.01
b = 3.0
plt.plot(x, (1 + x) ** 2)
plt.plot(x, a * (1 + x) ** 2 / x ** 2 + b)
```

```{code-cell}
%pylab inline --no-import-all
from mmfutils.plot import imcontourf


def step(x, a, b):
    P = [1.0, 2.0, 1 - a - b, -2 * a, -a]
    dP = np.polyder(P)
    f = np.polyval(P, x)
    return x - f / np.polyval(dP, x), f


def step1(x, a, b):
    f = x ** 2 - a - b * x ** 2 / (1 + x) ** 2
    df = 2.0 * x - 2.0 * b * x / (1 + x) ** 2 + 2.0 * b * x ** 2 / (1 + x) ** 3
    return x - np.divide(f, df), f


def step2(x, a, b):
    f = np.sqrt(a * (1 + x) ** 2 + b * x ** 2) - (1 + x) * x
    df = (a * (1 + x) + b * x) / np.sqrt(a * (1 + x) ** 2 + b * x ** 2) - 1 - 2 * x
    return x - np.divide(f, df), f


def step3(x, a, b):
    """This is very good except near a=0, b<1."""
    a = np.asarray(a)
    f = np.divide(1, (a / x ** 2 + b / (1 + x) ** 2)) - 1
    df = 2.0 * (a / x ** 3 + b / (1 + x) ** 3) * (1 + f) ** 2
    return x - np.divide(f, df), f


@np.vectorize
def count(a, b, x0=None, nmax=10):
    if x0 is None:
        x0 = np.sqrt(a + b)
    for n in range(1, nmax + 1):
        x0, f0 = step(x0, a, b)
        if abs(f0) / max(1, abs(x0) ** 4) < 1e-12:
            break
    return n


def plot_region(a, b, f=lambda a, b: np.sqrt(a + b), nmax=10):
    a_ = np.linspace(a[0], a[1], 101)[:, None]
    b_ = np.linspace(b[0], b[1], 101)[None, :]
    x0 = f(a_, b_)
    cs = count(a_, b_, x0=x0, nmax=nmax)
    imcontourf(a_, b_, cs)
    plt.colorbar(label="steps")
    plt.xlabel("a")
    plt.ylabel("b")
```

We can now plot over a region of reasonable parameter values.  Here we start with an initial guess $x_0 = 1.0$:

```{code-cell}
plot_region(a=(0, 5), b=(0, 20), f=lambda a, b: 1.0, nmax=100)
plt.plot(a_, 4 * (1 - a_), "r", scalex=False, scaley=False)
```

We see that this is really quite poor except for a small region along the line $b=4(1-a)$ which is the region near where $x=1$ is a solution.  *(This can be found by substituting $x=1$ into the original relationship.)*

To do better, we consider the factored form:

\begin{gather*}
  (1+x)^2(x^2-a) = b x^2.
\end{gather*}

If $b=0$, this has the solution $x|_{b=0} = \sqrt{a}$.  Likewise, if $a=0$, this has the solution $x|_{a=0} = \sqrt{b} - 1$.  This suggests a couple of what turn out to be an excellent guesses:

\begin{gather*}
  x_0 = \sqrt{a+b}, \qquad
  x_0 = \sqrt{a+\max(0, \sqrt{b}-1)^2}.
\end{gather*}

```{code-cell}
plot_region(a=(0, 5), b=(0, 20), f=lambda a, b: np.sqrt(a + b), nmax=10)
```

The only region this does not work well in is for small $a<0.2$ and small $b$.

```{code-cell}
plot_region(
    a=(0, 5),
    b=(0, 20),
    nmax=10,
    f=lambda a, b: np.sqrt(a + np.maximum(0, np.sqrt(b) - 1) ** 2),
)
```

```{code-cell}
plot_region(
    a=(0, 0.1),
    b=(0, 2),
    nmax=10,
    f=lambda a, b: np.sqrt(a + np.maximum(0, np.sqrt(b) - 1) ** 2),
)
```

The last approach only fails for very small $a$ near $b = 1 + \epsilon$.  To improve the guess, we start with

\begin{gather*}
  x^3(x + 2) = \epsilon x^2 + a(1+x)^2.
\end{gather*}

Now, noting that the solution is small, $b=1$ gives $2x_0^3 \approx a$, and inserting this, we obtain the perturbative solution:

\begin{gather*}
  x \approx x_0 + \frac{\epsilon}{6} = \left(\frac{a}{2}\right)^{1/3} + \frac{b-1}{6}.
\end{gather*}

```{code-cell}
plot_region(
    a=(0, 0.01),
    b=(0, 2),
    nmax=10,
    f=lambda a, b: (a / 2) ** (1.0 / 3.0) + (np.sqrt(b) - 1) / 3.0,
)
```

```{code-cell}
plot_region(
    a=(0, 0.01),
    b=(0, 2),
    nmax=10,
    f=lambda a, b: (a / 2) ** (1.0 / 3.0) + (b - 1) / 6.0,
)
```

Here is our final solution:

```{code-cell}
plot_region(
    a=(0, 0.01),
    b=(0, 2),
    nmax=10,
    f=lambda a, b: np.where(
        np.logical_and(abs(b - 1) < 0.2, a < 0.2),
        (a / 2) ** (1.0 / 3.0) + (b - 1) / 6.0,
        np.sqrt(a + np.maximum(0, np.sqrt(b) - 1) ** 2),
    ),
)
```

There are still some issues for extremely small $a$.

## Reverse Approach

Consider now the reverse approach of specifying the solution $x$ and looking at which
values of $a$ and $b$ will give that solution.

\begin{gather*}
  (1+x)^2(x^2-a) = b x^2, \qquad
  b = \frac{(1+x)^2(x^2-a)}{x^2}, \qquad
  a = x^2 - b \frac{x^2}{(1+x)^2}.
\end{gather*}

For large $x \gg 1$, the equation simplifies $x \approx \sqrt{a+b}$ which explains the
excellent convergence of this initial guess for large values of $a$ and $b$.

For small $x \ll 1$ we have $x \approx \sqrt{a/(1-b)}$.

Finally for $x = 1 + \epsilon$ we have $x \approx (b/2-4)/(b - 6 + 2a)$.

```{code-cell}
plot_region(
    a=(0, 2),
    b=(0, 2),
    nmax=7,
    f=lambda a, b: np.sqrt(a + np.maximum(0, np.sqrt(b) - 1) ** 2),
)
```

```{code-cell}
plot_region(a=(0, 2), b=(0, 2), nmax=7, f=lambda a, b: np.sqrt(a / (1 - b)))
```

```{code-cell}
plot_region(
    a=(0, 2), b=(0, 2), nmax=7, f=lambda a, b: (b / 2.0 - 4.0) / (b - 6.0 + 2.0 * a)
)
```

```{code-cell}
a_ = np.linspace(0, 1, 101)[:, None]
b_ = np.linspace(0, 1, 101)[None, :]

x0 = np.sqrt(a_ / (1 - b_))
cs = count(a_, b_, x0=x0, nmax=10)
imcontourf(a_, b_, cs)
plt.colorbar(label="steps")
plt.xlabel("a")
plt.ylabel("b")
cs = plt.contour(
    a_.ravel(), b_.ravel(), (x0.T) ** 2, levels=[0, 0.1, 0.2, 0.3, 0.4], c="k"
)
plt.clabel(cs, fontsize=9, inline=1)
```

```{code-cell}
a_ = np.linspace(0, 5, 101)[:, None]
b_ = np.linspace(0, 3, 101)[None, :]

x0 = (b_ / 2.0 - 4.0) / (b_ - 6 + 2 * a_)
cs = count(a_, b_, x0=x0, nmax=10)
cs[np.where(x0 < 0)] = 0
x0[np.where(x0 < 0)] = 0
imcontourf(a_, b_, cs)
plt.colorbar(label="steps")
plt.xlabel("a")
plt.ylabel("b")
cs = plt.contour(
    a_.ravel(), b_.ravel(), x0, 100
)  # levels=[0, 0.1, 0.2, 0.3, 0.4], c='k')
plt.clabel(cs, fontsize=9, inline=1)
```

```{code-cell}
a_ = np.linspace(0, 1, 101)[:, None]
b_ = np.linspace(0, 3, 101)[None, :]

x0 = np.sqrt(a_ + b_)
cs = count(a_, b_, x0=x0, nmax=10)
# cs[np.where(x0 < 0)] = 0
# x0[np.where(x0 < 0)] = 0
imcontourf(a_, b_, cs)
plt.colorbar(label="steps")
plt.xlabel("a")
plt.ylabel("b")
cs = plt.contour(
    a_.ravel(), b_.ravel(), x0, 10
)  # levels=[0, 0.1, 0.2, 0.3, 0.4], c='k')
plt.clabel(cs, fontsize=9, inline=1)
```

```{code-cell}
a_ = np.linspace(0, 0.3, 101)[:, None]
b_ = np.linspace(0.9, 1.2, 101)[None, :]
x0 = np.where(
    a_ + b_ > 1.0,
    np.sqrt(a_ + np.maximum((np.sqrt(b_) - 1) ** 2, 0)),
    np.where(a_ < 5 * (1 - b_), np.sqrt(a_ / (1 - b_)), np.nan),
)
cs = count(a_, b_, x0=x0, nmax=10)
# cs[np.where(x0 < 0)] = 0
# x0[np.where(x0 < 0)] = 0
imcontourf(a_, b_, cs)
plt.colorbar(label="steps")
plt.xlabel("a")
plt.ylabel("b")
```

```{code-cell}
a_, b_
```

```{code-cell}
a = 0.00006
b = 0.2
x0 = np.sqrt(a / (1 - b))
print count(a, b, x0=x0, nmax=100)
for n in xrange(10):
    x0, f0 = step(x0, a=a, b=b)
    print (x0, f0)
```

```{code-cell}
x1 = (b_ / 2.0 - 4) / (b_ - 6 + 2 * a_)
x2 = np.sqrt(a_ + b_)

np.where(x0 < np.minimum(x1, 1))

cs = count(a_, b_, x0=1.0, nmax=100)
imcontourf(a_, b_, cs)
plt.colorbar(label="steps")
plt.xlabel("a")
plt.ylabel("b")
plt.plot(a_, 4 * (1 - a_), "r", scalex=False, scaley=False)

np.where(a / (1 - b) < 0.5)
```

# Old Approach

+++

Our initial approach was to formulate the following system, and apply Newton's method, but this failed to converge in some cases.

\begin{gather*}
  \vec{F}(x) = \begin{pmatrix}
    (1 - C_ax_c) x_a^2 - B_a\\
    (1 - C_bx_c) x_b^2 - B_b\\
    x_c (x_a + x_b)^2 - 1
  \end{pmatrix}, \qquad
  \mat{J} = \begin{pmatrix}
    2(1 - C_ax_c) x_a &  0 & -C_ax_a^2\\
    0 & 2(1 - C_bx_c) x_b & -C_bx_b^2\\
    2x_c (x_a + x_b) & 2x_c (x_a + x_b) & (x_a + x_b)^2
  \end{pmatrix}
  \\
  x_a = \sigma_a^2, \qquad
  x_c = \frac{1}{(\sigma_a^2+\sigma_b^2)^2}, \qquad
\end{gather*}

+++

For physical parameters, $\alpha \approx 1$ while $\abs{\beta} < 1$.  If we let $\alpha = 1 - \epsilon$, then we have

\begin{gather*}
  (1+x)^2(x^2-1) = (\beta-\epsilon)x^2 - 2\epsilon x - \epsilon\\
  x_0 = 1 + \frac{\beta - 4(1-\alpha)}{8-2\beta+4(1-\alpha)}
\end{gather*}


We at looking for the largest root, so if $\alpha + \beta - 1 \gg 0$, then we have an approximate solution:

\begin{gather*}
  x_0 = \sqrt{\alpha + \beta - 1}
\end{gather*}

which is a good initial point, converging in at most 8 iterations for $\alpha > 0.25$ and $\beta > 1-\alpha/2$.  If $\alpha$ is small, then

\begin{gather*}
  x_0 = \sqrt{\beta} - 1
\end{gather*}

+++

\begin{gather*}
  x^4 + 2x^3 - 2x +1 = (x+1)^2(x^2-1)
\end{gather*}

```{code-cell}
import gpe.tube2

reload(gpe.tube2)
from gpe.tube2 import Sigma2s

Cs = [[1.0], [2.0]]
Bs = np.array([[0.6749164517058756], [0.1749387337039167]])
Cs = np.array([[0], [2.0001124853078269]])

np.random.seed(0)
Bs = np.random.random((2, 1))
Cs = np.random.random((2, 1))
ss = Sigma2s(Bs=Bs, Cs=Cs)
sa, sb = ss.solve3([[0.82], [2.0], [1 / (2.82) ** 2]])[:2]
alpha = Bs[0] / Bs[1]
beta = (Cs[0] - Cs[1]) / Bs[1]
# x = a/b
# assert np.allclose(x**4 + 2*x**3 + x**2 - alpha*(1+x)**2 - beta*x**2, 0)
print alpha, beta
# print(sa, sb, sa/sb)
ss.solve(), sa, sb
```

```{code-cell}
imcontourf(a_, b_, cs)
plt.plot(a_, 1 - a_ / 2)
plt.colorbar()
```

```{code-cell}
count(x0=9.0, alpha=0.8, beta=100)
```

```{code-cell}
alpha = 1.0
beta = 0.8
P = [1.0, 2.0, 1 - alpha - beta, -2 * alpha, -alpha]
xs = np.linspace(0, 2.0, 100)
plt.axhline(0)
plt.plot(xs, np.polyval(P, xs))
```

```{code-cell}
xs = np.linspace(1, 100.0, 100)
plt.plot(xs, count(x0=xs, alpha=100.0, beta=100.0))
```

```{code-cell}
alpha = 1.0
beta = 100.0
x0 = 10.0
for n in xrange(100):
    x0, f0 = step(x0, alpha=alpha, beta=beta)
    print x0, f0
```