import EOBRun_module as EOB
import matplotlib.pyplot as plt
import numpy as np 


def plot_teob_wf():

    fig, ax = plt.subplots()
    for lam in [400., 500., 800., 1000.]:
        pars = {
            # System parametes, assuming aligned spins        
            'M'                  : 2.7,       # Total mass
            'q'                  : 1.,     # Mass ratio m1/m2 > 1
            'chi1'               : 0.2,     # Z component of chi_1
            'chi2'               : 0.1,     # Z component of chi_2
            'LambdaAl2'          : lam,    # Quadrupolar tidal parameter of body 1 (A)
            'LambdaBl2'          : lam,    # Quadrupolar tidal parameter of body 2 (B)
            # Initial conditions and output time grid
            'domain'             : 1,      # Frequency domain, EOBSPA
            'srate_interp'       : 4096.,  # srate at which to interpolate. Default = 4096.
            'use_geometric_units': "no",   # output quantities in geometric units. Default = 1
            'initial_frequency'  : 10.0,   # in Hz if use_geometric_units = 0, else in geometric units
            'interp_uniform_grid': "no",   # interpolate mode by mode on a uniform grid. Default = 0 (no interpolation)
            'df'                 : 1./256,
            # Modes
            'use_mode_lm'        : [1],    # List of modes to use/output through EOBRunPy

            # Output parameters (Python)
            'arg_out'            : "no",  # Output hlm/hflm. Default = 0
        }
        f, rhp, ihp, _, _ = EOB.EOBRunPy(pars)
        hp   = rhp -1j*ihp
        A    = abs(hp)
        ax.loglog(f, A, label=lam)
    
    ax.set_xlabel('f [Hz]')
    ax.set_ylabel(r'$\tilde{A}$')
    # plt.legend()
    ax.set_xlim(600, 2500)
    ax.set_ylim(1e-24,1e-22)
    plt.legend()
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    plot_teob_wf()

Updated: