<jupyter_start><jupyter_text>The root-finding function for case D is:$$ f(\alpha, \xi, \chi, \beta, \lambda) = \alpha + \frac{1}{2}\left(\lambda - \beta\sqrt{\lambda^2 + \chi^2 + 4(1+\chi)\sin^2(\alpha) + 2\lambda\sin(2\alpha)}\right) - \xi. $$The goal is to find $\alpha_0$ such that $f(\alpha = \alpha_0) = 0$. Some numerical root finders ( e.g. Newton ) expect $f'(\alpha)$ for accuracy. With chain rule we have:$$ f'(\alpha) = 1 - \frac{\beta}{4}\left( \lambda^2 + \chi^2 + 4(1+\chi)\sin^2(\alpha) + 2\lambda\sin(2\alpha)\right)^{-1/2} \times (8(1+\chi)\sin(\alpha)\cos(\alpha) + 4\lambda\cos(2\alpha)) $$ With simplification:$$ f'(\alpha) = 1 - \beta \left( \frac{2(1+\chi)\sin(\alpha)\cos(\alpha) + \lambda\cos(2\alpha)}{ \sqrt{ \lambda^2 + \chi^2 + 4(1+\chi)\sin^2(\alpha) + 2\lambda\sin(2\alpha) }}  \right) $$<jupyter_code># Note that ff can NOT be np-vectorized since it will be called by the jitted "newton" ot "brentq".

@njit
def ff(a, z, x, beta, lamb):
    return a + 1/2 * (lamb - beta* sqrt(lamb**2 + x**2 + 4*(1+x)*sin(a)**2 + 2*lamb*sin(2*a)) - z)

@njit
def ff_prime(a, z, x, beta, lamb):
    return 1 - beta * (2*(1+x)*sin(a)*cos(a) + lamb*cos(2*a)) / sqrt(lamb**2 + x**2 + 4*(1+x)*sin(a)**2 + 2*lamb*sin(2*a))

# Newton needs the first derivate and a good initial guess ( otherwise may not converge )  
print(newton(ff, 0, ff_prime, args=(1e-5, 1e-5, beta, lamb))[0])

@vectorize([float64(float64, float64, float64, float64, float64)], target='parallel')
def newton_vec_ff(a0, z, x, beta, lamb):
    return newton(ff, a0, ff_prime, args=(z, x, beta, lamb))[0]

print( newton_vec_ff(0, 1e-5, 1e-5, beta, lamb) )
print( newton_vec_ff(0, 1.1e-5, 1e-5, beta, lamb))
print( newton_vec_ff(0, 2.1e-5, 1e-5, beta, lamb))<jupyter_output>0.02548575935978343
0.02548575935978343<jupyter_text>Since Newton is hard to use ( have to guess initial $\alpha$ smartly ), let's try the Brent's method<jupyter_code># Brent's method needs a "guess interval" 
print( brentq(ff, -0.01, 0.1, args=(1e-5, 1e-5, beta, lamb))[0])

@vectorize([float64(float64, float64, float64, float64)], target='parallel')
def brentq_vec_ff(z, x, beta, lamb):
    #return brentq(ff, -0.01, 0.1, args=(z, x, beta, lamb))[0]
    return brentq(ff, -1, 1, args=(z, x, beta, lamb))[0]

print(brentq_vec_ff( 1e-5, 1e-5, beta, lamb))
print(brentq_vec_ff( 1.1e-5, 1e-5, beta, lamb))
print(brentq_vec_ff( 2.1e-5, 1e-5, beta, lamb))<jupyter_output>0.025485759359775418
0.025485759359768323
0.02640272292353774
0.03354960543650097<jupyter_text>Timing test<jupyter_code># Define a simple 2D grid 

zvec = np.linspace(-6,6,501)*1e-5
xvec = np.linspace(-100,100,501)*1e-5

zm,xm = np.meshgrid(zvec, xvec, indexing="ij")<jupyter_output><empty_output><jupyter_text>Old method (with np.vectorize and no numba) is SLOW<jupyter_code># OLD method ( root-finding without numba )

# This function cannot be directly numba-jitted or numba-vectorized
# because "scipy.optimize.root_scalar" is not compatible. 

@np.vectorize
def alpha_exact_case_D(z, x, beta, lamb):
  
    f = lambda a: a + 1/2 * (lamb - beta*sqrt(lamb**2 + x**2 + 4*(1+x)*sin(a)**2 + 2*lamb*sin(2*a)) - z)
    
    res = scipy.optimize.root_scalar(f, bracket=(-1,1))
    
    return res.root

t1 = time.time()
alpha_exact_case_D(zm, xm, beta, lamb)
t2 = time.time()
print('Without Numba it takes', t2-t1, 'seconds')<jupyter_output>Without Numba it takes 36.07687425613403 seconds<jupyter_text>With numba-jitted root-finder (from quantecon), it's FAST<jupyter_code>%%timeit
# Timing with Numba
GG = brentq_vec_ff(zm, xm, beta, lamb)
GG = brentq_vec_ff(zm, xm, beta, lamb)
from mpl_toolkits import mplot3d

fig, ax = plt.subplots(figsize=(8,8))
ax = plt.axes(projection='3d')
#ax.contour3D(Z,X, psi_s_grid , 1000, cmap='inferno')
ax.plot_surface(zm*1e5, xm*1e5, GG, cmap='inferno')
#ax.set_xlabel('$z/2/\rho$ (1e-3)')
plt.xlabel(r'$\xi (10^{-5})$')
plt.ylabel(r'$\chi (10^{-5})$')
ax.set_zlabel(r'$\alpha (rad)$')
plt.title(r'$\alpha(\xi,\chi)$ for case D')

ax.azim = 240
plt.plot(zvec,GG[:,250], label = '$\chi$ = 0')
plt.plot(zvec,GG[:,0],label = '$\chi$ = -0.001')
plt.plot(zvec,GG[:,-1], label = '$\chi$ = 0.001')

plt.xlabel(r'$\xi$', fontsize=16)
plt.ylabel(r'$\alpha$ (rad)', fontsize=16)
plt.title(r'$\alpha$ v.s $\xi$ for case D',  fontsize=20)
plt.legend()<jupyter_output><empty_output><jupyter_text>Below are scratch<jupyter_code>import numpy as np
import scipy
from numba import vectorize, float64, njit
#@vectorize([float64(float64, float64, float64, float64)])

@np.vectorize
def alpha_exact_case_D(z, x, beta, lamb):
    """
    Exact alpha calculation using numerical root finding.

    """
    #beta = np.sqrt(beta2)
    f = lambda a: a + 1/2 * (lamb - beta*np.sqrt(lamb**2 + x**2 + 4*(1+x)*np.sin(a)**2 + 2*lamb*np.sin(2*a)) - z)
    
    res = scipy.optimize.root_scalar(f, bracket=(-1,1))
    
    return res.root


#@vectorize([float64(float64, float64, float64, float64)])
@np.vectorize
def Es_case_D(z, x, beta, lamb):
    """
    Eq.(?) from Ref[2] slide #21 with no constant factor e*beta**2/rho**2.
    Note that 'x' here corresponds to 'chi = x/rho', 
    and 'z' here corresponds to 'xi = z/2/rho' in the paper. 
    """
  
    if z == 0 and x == 0:
        return 0
    
    #beta2 = beta**2
    alp = alpha_exact_case_D(z, x, beta, lamb)
    sin2a = np.sin(2*alp)
    cos2a = np.cos(2*alp) 

    kap = np.sqrt(lamb**2 + x**2 + 4*(1+x)*np.sin(alp)**2 + 2*lamb*np.sin(2*alp)) # kappa for case D
    
    N1 = cos2a - (1+x)
    N2 = lamb*cos2a + (1+x)*sin2a - beta*kap
    D = kap - beta*(lamb*cos2a + (1+x)*sin2a)
    
    return N1*N2/D**3
import numpy as np
zvec = np.linspace(-100,100,101)*1e-6
xvec = np.linspace(-100,100,101)*1e-6
beta = 0.999999
lamb = 0.3
zm,xm = np.meshgrid(zvec, xvec, indexing="ij")
alpha_exact_case_D(zm,xm, beta, lamb)
Es_case_D(zm,xm, beta, lamb)
f = lambda a: a + 1/2 * (lamb - beta*np.sqrt(lamb**2 + x**2 + 4*(1+x)*np.sin(a)**2 + 2*lamb*np.sin(2*a)) - z)
    
    res = scipy.optimize.root_scalar(f, bracket=(-1,1))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19953557803947586, 'fraction_numerical': 0.07546856858517167, 'mean_word_length': 4.2663755458515285, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '60858', 'license': 'Apache-2.0', 'n_tokens_mistral': 2910, 'n_tokens_neox': 2640, 'n_words': 574, 'path': 'benchmarking/one_bend_transient/caseD_vectorize.ipynb', 'repo_name': 'ChristopherMayes/PyCSR2D', 'size': '6029'}