Solving for Sigmas: Globally Convergent Newton’s Method

Hide code cell content

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

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.

Solving for Sigmas: Globally Convergent Newton’s Method#

To implement the NPSEQ for two components as described in Non-Polynomial Schrödinger Equation (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 Power-law Equations of State where \(n_0^2 = Cq^{1/(1-\beta)}\), we demonstrate in section Globally Convergent Newton’s Method 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 Globally Convergent Newton’s Method 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*}\]

Hide code cell content

 
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))
\[\displaystyle - \frac{\frac{d}{d n_{0}} f{\left(n_{0} \right)}}{n_{0} f^{3}{\left(n_{0} \right)}}\]
\[\displaystyle - \frac{\frac{d}{d n_{0}} f{\left(n_{0} \right)}}{n_{0} f{\left(n_{0} \right)}} - \frac{1}{n_{0}^{2}}\]
\[\displaystyle \frac{- n_{0} \frac{d}{d n_{0}} h{\left(n_{0} \right)} + h{\left(n_{0} \right)}}{h^{3}{\left(n_{0} \right)}}\]
\[\displaystyle - \frac{\frac{d}{d n_{0}} h{\left(n_{0} \right)}}{n_{0} h{\left(n_{0} \right)}}\]

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 Globally Convergent Newton’s Method 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 Globally Convergent Newton’s Method 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.

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)

Hide code cell source

 
# 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))

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)

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[5], line 71
     67 
     68 kw.pop('ax')
     69 kw.pop('ay')
     70 kw['beta'] = 2
---> 71 plot_region(ax=(0.01,100.0), ay=(0.01,100.0), **kw)

Cell In[5], line 62, in plot_region(ax, ay, get_qt0, nmax, **kw)
     59 def plot_region(ax, ay, get_qt0=get_qt0, nmax=10, **kw):
     60     ax_ = np.linspace(ax[0], ax[1], 101)[:, None]
     61     ay_ = np.linspace(ay[0], ay[1], 101)[None, :]
---> 62     cs = count(ax_, ay_, qt=get_qt0(ax_,ay_, **kw), nmax=nmax, **kw)
     63     fig, ax = plt.subplots()
     64     mesh = ax.pcolormesh(ax_.ravel(), ay_.ravel(), cs.T)
     65     fig.colorbar(mesh, ax=ax, label="steps")

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/numpy/lib/_function_base_impl.py:2518, in vectorize.__call__(self, *args, **kwargs)
   2515     self._init_stage_2(*args, **kwargs)
   2516     return self
-> 2518 return self._call_as_normal(*args, **kwargs)

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/numpy/lib/_function_base_impl.py:2511, in vectorize._call_as_normal(self, *args, **kwargs)
   2508     vargs = [args[_i] for _i in inds]
   2509     vargs.extend([kwargs[_n] for _n in names])
-> 2511 return self._vectorize_call(func=func, args=vargs)

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/numpy/lib/_function_base_impl.py:2599, in vectorize._vectorize_call(self, func, args)
   2597 # gh-29196: `dtype=object` should eventually be removed
   2598 args = [asanyarray(a, dtype=object) for a in args]
-> 2599 outputs = ufunc(*args, out=...)
   2601 if ufunc.nout == 1:
   2602     res = asanyarray(outputs, dtype=otypes[0])

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/numpy/lib/_function_base_impl.py:2506, in vectorize._call_as_normal.<locals>.func(*vargs)
   2504     the_args[_i] = vargs[_n]
   2505 kwargs.update(zip(names, vargs[len(inds):]))
-> 2506 return self.pyfunc(*the_args, **kwargs)

Cell In[5], line 54, in count(ax, ay, qt, nmax, **kw)
     50 @np.vectorize
     51 def count(ax, ay, qt, nmax=10, **kw):
     52     kw.update(ax=ax, ay=ay)
     53     for n in range(1, nmax+1):
---> 54         qt, g = newton_tilde(qt, **kw)
     55         if abs(g) / max([1.0, abs(qt)]) < 1e-12:
     56             break
     57     return n

Cell In[5], line 36, in newton_tilde(qt, ax, ay, c, q, **kw)
     32     _tmp = C*(c/C)**(1/beta)
     33     _power = 1/(1-beta)
     34     n2 = _tmp * qt**_power
     35     dn2_dqt = _power * n2 / qt
---> 36     g = newton_g(n2, **kw)
     37     dg = newton_g(n2, d=1, **kw) * dn2_dqt
     38     return qt - g/dg, g

Cell In[3], line 32, in newton_g(n2, ax, ay, d, c, q, **kw)
     30 def newton_g(n2, ax, ay, d=0, c=1.0, q=q, **kw):
     31     _q = q(n2, d=0, **kw)
---> 32     _dq = q(n2, d=1, **kw)
     33     _ax, _ay = 1 + ax*_q, 1 + ay*_q
     34     _sqrt_ax, _sqrt_ay = np.sqrt(_ax), np.sqrt(_ay)
     35     _x, _y = 1 + _sqrt_ax, 1 + _sqrt_ay

Cell In[3], line 28, in q(n2, d, **kw)
     24     n = np.sqrt(n2)
     25     if d == 0:
     26         return 1 / f(n, d=0, **kw)**2
     27     elif d == 1:
---> 28         return -f(n, d=1, **kw) / f(n, d=0, **kw)**3/n

Cell In[3], line 14, in f(n, d, **kw)
     10 def f(n, d=0, **kw):
     11     if d == 0:
     12         return Q(n, d=1, **kw) - Q(n, d=0, **kw)/n
     13     elif d == 1:
---> 14         return Q(n, d=2, **kw) - Q(n, d=1, **kw)/n + Q(n, d=0, **kw)/n**2

Cell In[3], line 1, in Q(n, d, g, beta, **kw)
----> 1 def Q(n, d=0, g=1.0, beta=2.0, **kw):
      2     """Return Q(n) and its derivatives where  Q'(n) = E/N, the energy per particle."""
      3     if d == 0:
      4         return g*n**beta/2/beta

KeyboardInterrupt:

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

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*}\]
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*}\]
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__)

%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*}\]
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)

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*}\]
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\)

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:#

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*}\]

Details

\[\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*}\]

Symmetric GPE

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*}\]

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*}\]

Another approach.

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*}\]
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*}\]
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__)

%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*}\]
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)

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*}\]
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\)

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.

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

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)

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)

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)

%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\):

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*}\]
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\).

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),
)

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*}\]
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,
)

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:

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)\).

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),
)

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

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)
)

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)

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)

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)

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")

a_, b_

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)

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*}\]
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

imcontourf(a_, b_, cs)
plt.plot(a_, 1 - a_ / 2)
plt.colorbar()

count(x0=9.0, alpha=0.8, beta=100)

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))

xs = np.linspace(1, 100.0, 100)
plt.plot(xs, count(x0=xs, alpha=100.0, beta=100.0))

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