Solving multiple solution equation in Python

Question

I need to get the solutions of the following trigonometric equation:

I know there exists either M or M-1 solutions depending on the value of $\Delta$. Does anyone know of any module or algorithm that I can use in Python for this task?

I've tried this way but it is very sensitive to the tolerance and it doesn't work:

def f(k, N, d):
    return np.tan(k*(N)) - d*np.sin(k)/(1+d*np.cos(k))
    

k = np.linspace(0, np.pi, 10000+1, endpoint=False)[1:] 

def ksolutions(k,N,d):
    solutions=[]
    tol=4*1e-4
    for i in k:
        if abs(f(i,N,d))<tol:
            solutions.append(i)
    
    print(solutions) 

Answer 1

I know there exists either M or M-1 solutions depending on the value of delta

No. There are more.

Because this is a scalar function and the analytic first and second partial derivatives are straightforward to calculate, I suggest using Halley's method. Unfortunately, it does not magically find all roots; but it is very fast at finding some root. There are various ways to perform multi-root iteration; the simplest is to have a preset locus of search starts (similar to what you did, but lower-resolution). Care needs to be taken to eliminate duplicates.

import functools

import matplotlib.pyplot as plt
import numpy as np
from numpy import sin, cos, tan
from scipy.optimize import root_scalar, check_grad


def f(x: float, M: int, Δ: float) -> float:
    return (1 + Δ*cos(x))*tan(M*x) + Δ*sin(x)


def dfdx(x: float, M: int, Δ: float) -> float:
    return (
        M*(1 + Δ*cos(x))/cos(M*x)**2
        + Δ*(cos(x) - sin(x)*tan(M*x))
    )


def d2fdx2(x: float, M: int, Δ: float) -> float:
    sec = 2*M/cos(M*x)**2
    tanM = tan(M*x)
    cosx = cos(x)
    return (
        M*sec*tanM*(1 + Δ*cosx)
        - Δ*(
            cosx*tanM + sin(x)*(sec + 1)
        )
    )


def solve(M: int, Δ: float, start_factor: int = 20, tol: float = 1e-6):
    starts = np.linspace(start=0, stop=np.pi, num=start_factor*M)
    roots = []
    approx_roots = set()
    args = {'M': M, 'Δ': Δ}
    fp = functools.partial(f, **args)
    df = functools.partial(dfdx, **args)
    df2 = functools.partial(d2fdx2, **args)

    err1 = check_grad(fp, df, 0.5)
    err2 = check_grad(df, df2, 0.5)
    assert err1 < 1e-4
    assert err2 < 1e-3

    for x0 in starts:
        sol = root_scalar(f=fp, fprime=df, fprime2=df2, x0=x0, method='halley')
        if sol.converged:
            x = sol.root
            y = fp(x)
            if abs(y) < tol and -tol < x < np.pi + tol:
                approx_root = round(x, 4)
                if approx_root not in approx_roots:
                    approx_roots.add(approx_root)
                    roots.append(x)

    return np.array(roots)


def demo() -> None:
    xhi = np.linspace(start=0, stop=np.pi, num=1_001)

    for M in (1, 4):
        for Δ in (0.5, 3.1):
            roots = solve(M=M, Δ=Δ)
            print(f'M={M} Δ={Δ} roots={roots}')
            print('f =', f(roots, M, Δ))
            print()

            fig, ax = plt.subplots()
            ax.set_title(f'M={M} Δ={Δ}')
            ax.set_xlabel('x')
            ax.set_ylabel('f')
            y = f(xhi, M, Δ)
            mask = (y > -10) & (y < 10)
            ax.plot(xhi, np.where(mask, y, np.nan))
            ax.scatter(roots, np.zeros_like(roots))

    plt.show()


if __name__ == '__main__':
    demo()

M=1 Δ=0.5 roots=[0.         3.14159264]
f = [0. 0.]

M=1 Δ=3.1 roots=[0.         1.73279428 3.14159265]
f = [ 0.00000000e+00 -3.10862447e-15  6.36816336e-16]

M=4 Δ=0.5 roots=[0.         1.4607308  0.72612363 2.22665935 3.14159265]
f = [ 0.00000000e+00 -2.22044605e-16 -1.11022302e-16 -2.22044605e-16
 -1.83697020e-16]

M=4 Δ=3.1 roots=[0.         0.65956474 1.3127522  3.14159265 2.56082681 1.95055524]
f = [ 0.00000000e+00 -8.88178420e-16  3.10862447e-15  1.40834382e-15
  1.55431223e-15  1.77635684e-15]