cavity class

import numpy as np
from scipy.signal import argrelextrema

class Device:
    def __init__(self, fc, Q, kin_factor, fe, gamma, C):
        self.Q = Q
        self.fc = fc
        self.wc = 2*np.pi*fc

        self.k = 2*np.pi*self.fc/self.Q # Decay rate
        self.Af_c = self.fc/self.Q # Cavity linewidth

        self.kin = kin_factor*self.k # Waveguide coupling rate

        self.fe = fe
        self.we = 2*np.pi*fe
        self.gamma = gamma

        self.C = C
        self.g = np.sqrt(self.C*self.k*self.gamma/4)

        self.span = 4*self.Af_c

        self.Ace = (fe-fc)/1e9

    def r(self, fp):
        wp = 2*np.pi*fp
        return 1-self.kin/(1j*(wp-self.wc)+self.k/2+self.g**2/(1j*(wp-self.we)+self.gamma/2))

    def R(self, fp):
        return np.abs(self.r(fp))**2

    def phase(self, fp):
        return np.angle(self.r(fp))

    def find_extrema(self):
        fp = np.linspace(self.fc-self.span, self.fc+self.span, 1000000)
        x_axis = (fp-self.fc)/1e9
        R = self.R(fp)

        # Finding local extrema
        # Local maxima
        maxima = argrelextrema(R, np.greater_equal)
        f_max = x_axis[maxima]
        R_max = R[maxima]

        # Local minima
        minima = argrelextrema(R, np.less_equal)
        f_min = x_axis[minima]
        R_min = R[minima]

        # Find extrema closest to emitter frequency
        index_high = np.argmin(np.abs(f_max-self.Ace))
        f_high = f_max[index_high]
        R_high = R_max[index_high]
        index_low = np.argmin(np.abs(f_min-self.Ace))
        f_low = f_min[index_low]
        R_low = R_min[index_low]

        return [f_high, R_high, f_low, R_low]

    def find_overlapping_splitting(self):
        f_high, R_high, f_low, R_low = self.find_extrema()

        res = 100
        Ass = np.linspace(0, 2*(f_low-f_high), res) # Spin splitting (GHz)

        fp = np.linspace(self.fe-2*Ass*1e9, self.fe+2*Ass*1e9, 10000)
        x_axis = (fp-self.fc)/1e9
        wp = 2*np.pi*fp

        diff = np.zeros(np.shape(Ass))
        for i, Ass_i in enumerate(Ass):
            we_i = self.we + 2*np.pi*Ass_i*1e9
            r_i = 1-self.kin/(1j*(wp-self.wc)+self.k/2+self.g**2/(1j*(wp-we_i)+self.gamma/2))
            R_i = np.abs(r_i)**2

            # Finding local extrema
            # Local maxima
            maxima_i = argrelextrema(R_i, np.greater)
            f_max_i = x_axis[maxima_i]
            R_max_i = R_i[maxima_i]

            # Find extrema closest to emitter frequency
            Ace_i = self.Ace + Ass_i
            index_high_i = np.argmin(np.abs(f_max_i-Ace_i))
            f_high_i = f_max_i[index_high_i]
            R_high_i = R_max_i[index_high_i]

            diff[i] = np.abs(f_high_i - f_low)

        index_overlap_high_low = np.argmin(diff)
        Ass_overlap = Ass[index_overlap_high_low]

        return Ass_overlap