def ampli(dpt1, f, rp=[], layers=[1500, 1000, 5540, 3200, 2500], theta = radians(15.71)):
    """
    Compute amplification coefficient for P and S waves. 

    Args:
        dpt1 (ndarray): bathymetry grid in meters
        f (ndarray): frequency vector in Hz
        rp (ndarray, optional): ray parameter matrix to integrate over. 
        layers (list, optional): layers properties [Vp_w, rho_w, Vp_c, Vs_c, rho].
        theta (float, optional): limit angle of. Defaults to radians(15.71).

    Returns:
        cP (ndarray): Amplification coefficient for P waves.
        cS (ndarray): Amplification coefficient for S waves.
        bathy_ampli_P (ndarray): Amplification coefficient for P waves divided by theta.
        bathy_ampli_S (ndarray): Amplification coefficient for S waves divided by theta.

    """
    # compute amplification for Tp and rp following G14
    (cP, cS) = bathy(dpt1, f, rp, layers)
    if np.size(f)==1:
        cP = abs(cP.reshape(dpt1.shape, order='F').copy())
        cS = abs(cS.reshape(dpt1.shape, order='F').copy())
        bathy_ampli_P = cP/theta_lim
        bathy_ampli_S = cS/theta_lim
        return cP, cS, bathy_ampli_P, bathy_ampli_S
    else:
        new_cP = np.empty((np.size(f), dpt1.shape[0], dpt1.shape[1]))
        new_cS = np.empty((np.size(f), dpt1.shape[0], dpt1.shape[1]))
        for i in range(np.size(f)):
            new_cP[i, :] = abs(cP[i,:].reshape(dpt1.shape, order='F').copy())
            new_cS[i, :] = abs(cS[i,:].reshape(dpt1.shape, order='F').copy())
        bathy_ampli_P = new_cP/theta
        bathy_ampli_S = new_cS/theta
        return new_cP, new_cS, bathy_ampli_P, bathy_ampli_S

def bathy(z, f, p=[], m= [1500, 1000, 55400, 3200, 2500]):
    """Bathymetry secondary microseismic excitation coefficients (for P, S amplitude).
    Based on LI Lei, ll.ynyf@gmail.com modified by Pierre Boue 23/11/2020.

    Examples:
        >>> z = np.linspace(0, 25000, 5001)
        >>> f = 1/8
        >>> cP, cS = bathy(z, f, [], [])
        >>> x = f*z/1500
        >>> plt.figure()
        >>> plt.plot(x, abs(cP[0]), x, abs(cS[0]))
        >>> plt.show()

    Args:
        z (np.ndarray): thickness of water layer in meters, if z in m|km, v should be in m/s|km/s. All rhos must keep the same units.
        f (np.ndarray): seismic frequency in Hz
        p (np.ndarray, optional): slowness, if not specified return p values integral to 1/vp_crust
        m (list, optional): [vp_water, rho_water, vp_crust, vs_crust, rho_crust]

    Returns:
        cP (np.ndarray): excitation coefficients of P waves in shape of [[p,]f, z]
        cS (np.ndarray): excitation coefficients of S waves in shape of [[p,]f, z]

    """
    ##
    (rhow, rhoc) = (m[1], m[4])  # in kg/m³
    (vpw, vpc, vsc) = (m[0], m[2], m[3])  # in m/s
    if len(m) > 0:
        if (len(m) >= 1 and m[0] > 0):
            vpw = m[0]
        if (len(m) >= 2 and m[1] > 0):
            rhow = m[1]
        if (len(m) >= 3 and m[2] > 0):
            vpc = m[2]
        if (len(m) >= 4 and m[3] > 0):
            vsc = m[3]
        if (len(m) >= 5 and m[4] > 0):
            rhoc = m[4]

    elif np.max(z) < 50:
        (vpw, rhow, vpc, vsc, rhoc) = (1.5, 1, 5.54, 3.2, 2.5)
        print('[rho_w, rho_c] = [1, 2.5] g/cm³')
        print('[vpw, vpc, vsc] = [1.5, 5.54, 3.2] km/s')
    if len(p)==0:
        p = np.linspace(0, 0.995, 200)/vpc
        m = np.array([vpw, rhow, vpc, vsc, rhoc])
        a = np.arcsin(vpw*p)
        cP = np.zeros((np.size(f), np.size(z)), dtype="complex_")
        cS = np.zeros((np.size(f), np.size(z)), dtype="complex_")
        for i in range(np.size(f)):
            if np.size(f) == 1:
                c_P, c_S = bathy(z, f, p, m)
            else:
                c_P, c_S = bathy(z, f[i], p, m)
            cP[i, :] = cP[i, :] + np.trapz(abs(c_P)**2, x=a, axis=0)
            cS[i, :] = cS[i, :] + np.trapz(abs(c_S)**2, x=a, axis=0)
        (cP, cS) = (csqrt(cP), csqrt(cS))
        return cP, cS

    elif (p != np.linspace(0, 0.995, 200)/vpc).all():
        c_P = np.empty((np.size(p), np.size(f), np.size(z)), dtype="complex_")
        c_P[:] = np.nan
        c_S = np.empty((np.size(p), np.size(f), np.size(z)), dtype="complex_")
        c_S[:] = np.nan
        cP = np.zeros((np.size(f), np.size(z)), dtype="complex_")
        cS = np.zeros((np.size(f), np.size(z)), dtype="complex_")
        Rpp, Tpp, Tps = subfcn_liquid_solid(p, [vpw, rhow], [vpc, vsc, rhoc])
        qw = csqrt(1/vpw**2 - p**2)
        if np.size(f) == 1:
            phi = 4*np.pi*f*(z.flatten('F')).T
            for i in range(np.size(p)):
                C = 1*np.reciprocal(1 + Rpp[i]*np.exp(1j*phi*qw[i]))
                c_P[i, :, :] = np.dot(Tpp[i], C) 
                c_S[i, :, :] = np.dot(Tps[i], C)  
        else:
            for i in range(np.size(f)):
                c_Pi, c_Si = bathy(z, f[i], p, m)
                c_P[:, i, :] = c_Pi
                c_S[:, i, :] = c_Si

        for i in range((np.size(f))):
                a = np.arcsin(vpw*p)
                cP[i, :] = cP[i, :] + np.trapz(abs(c_P[:,i,:])**2, x=a, axis=0)
                cS[i, :] = cS[i, :] + np.trapz(abs(c_S[:, i, :])**2, x=a, axis=0)
        (cP, cS) = (np.squeeze(cP), np.squeeze(cS))
    Rpp, Tpp, Tps = subfcn_liquid_solid(p, [vpw, rhow], [vpc, vsc, rhoc])
    qw = csqrt(1/vpw**2 - p**2)
    if np.size(f) == 1:
        phi = 4*np.pi*f*(z.flatten('F')).T
    else:
        phi = 4*np.pi*np.dot(f.flatten('F'), (z.flatten('F')).T)
    cP = np.empty((np.size(p), np.size(f), np.size(z)), dtype="complex_")
    cP[:] = np.nan
    cS = cP.copy()
    for i in range(np.size(p)):
        C = 1*np.reciprocal(1 + Rpp[i]*np.exp(1j*phi*qw[i]))
        cP[i, :, :] = np.dot(Tpp[i], C)  
        cS[i, :, :] = np.dot(Tps[i], C)
    (cP, cS) = (np.squeeze(cP), np.squeeze(cS))
    return cP, cS

def download_ww3_local(YEAR, MONTH, ftp_path_to_files="ftp://ftp.ifremer.fr/ifremer/dataref/ww3/GLOBMULTI_ERA5_GLOBCUR_01/GLOB-30M/2020/FIELD_NC/", ww3_local_path= '../../data/ww3/', prefix = "WW3-GLOB-30M"):
    """Download WW3 files for a given year and month from the specified FTP path to a local directory.

    Args:
        YEAR (int): The year for which the files should be downloaded.
        MONTH (list): The month or months for which the files should be downloaded.
        ftp_path_to_files (str, optional): The FTP path to the WW3 files.
        ww3_local_path (str, optional): The local directory where the files should be saved.
        prefix (str, optional): The prefix for the WW3 files.

    """
    workdir = os.getcwd()
    # create directory if it does not exist
    try: 
        os.mkdir(ww3_local_path) 
    except OSError as error: 
        print(error)
    os.chdir(ww3_local_path)

    # download files
    if len(MONTH) == 0:
        print("-----------------------------------------------------------------\n")
        print("Downloading WW3 files for year %d\n"%(YEAR))
        print("-----------------------------------------------------------------\n")
        MONTH = np.arange(1, 13)
    for m in MONTH:
        print("Downloading can take some time...\n")
        file_p2l = ftp_path_to_files + "%s_%d%02d_p2l.nc"%(prefix, YEAR, m) # p2l file
        check_file_p2l = ww3_local_path + "%s_%d%02d_p2l.nc"%(prefix, YEAR, m) # p2l file

        if os.path.exists(check_file_p2l):
            print("-----------------------------------------------------------------\n")
            print(check_file_p2l + " already downloaded\n")
            print("-----------------------------------------------------------------\n")
        else:
            os.system("wget -nv -c %s"%(file_p2l))
            print("-----------------------------------------------------------------\n")
            print(file_p2l + " downloaded\n")
            print("-----------------------------------------------------------------\n")

    print( "WW3 files downloaded in %s"%(ww3_local_path))
    os.chdir(workdir)
    print("current directory : ", os.getcwd())

def loop_ww3_sources(paths, dpt1, zlon, zlat, wave_type='P', date_vec=[2020, [], [], []], extent=[-180, 180, -90, 90],parameters= [1/12, 1/2], c_file = "../../data/cP.nc", prefix = "WW3-GLOB-30M", **kwargs):
    """Compute the equivalent vertical force on the seafloor for a given wave type (P or S), given a path to the ww3 p2l file, the bathymetry, the wave type, the date vector and the spatial extent.
    Saves in netcdf format the equivalent vertical force for each frequency if save argument True.
    Plots in PNG source maps of P/S waves at given intervals depending on plot variables.

    Args:
        paths (list): [file_bathy, ww3_local_path]: paths of additional files bathymetry, ww3 p2l file
        dpt1 (xarray.ndarray): bathymetry grid in m (depth) with dimensions lon x lat
        zlon (xarray.ndarray): longitude of bathymetry file (°)
        zlat (xarray.ndarray): latitude of bathymetry file (°)
        wave_type (str, optional): P or S waves.
        date_vec (list, optional): date vector [year, month, day, hour], with hour in [0, 3, 6, 9, 12, 15, 18, 21].
        extent (list, optional): spatial extent format [lon_min, lon_max, lat_min, lat_max].
        parameters (list, optional): parameters minimum frequency, maximum frequency.
        c_file (str, optional): path to amplification coefficient file.
        prefix (str, optional) : prefix of ww3 p2l file.
        plot (bool, optional): default: True
        plot_hourly (bool, optional): plot maps every 3-hours default: False
        plot_daily (bool, optional): plot maps every day, default: False
        plot_monthly (bool, optional): plot maps every month, default : True
        plot_yearly (bool, optional): plot map for the year average, default: False
        save (bool, optional): save 3-hourly matrix, default: False

    """

    file_bathy = paths[0]
    ww3_local_path = paths[1]

    # Constants
    radius = 6.371*1e6 # radius of the earth in meters
    lg10 = log(10) # log of 10
    #
    f1 = parameters[0]
    f2 = parameters[1]
    ## Initialize variables
    if 'plot_type' in kwargs:
        plot_type = kwargs['plot_type']
        plot = True
        if plot_type == 'hourly':
            plot_hourly = True
            F_hourly =  np.zeros(dpt1.shape)
        else:
            plot_hourly = False
        if plot_type == 'daily':
            plot_daily = True
            F_daily =  np.zeros(dpt1.shape)
        else:
            plot_daily = False
        if plot_type == 'monthly':
            plot_monthly = True
            F_monthly =  np.zeros(dpt1.shape)
        else:
            plot_monthly = False
        if plot_type == 'yearly':
            plot_yearly = True
            F_yearly =  np.zeros(dpt1.shape)
        else:
            plot_yearly = False
    else:
        plot = True

    if 'save' in kwargs:
        save = kwargs['save']
    else:
        save = False

    if 'vmin' in kwargs:
        vmin = kwargs['vmin']
    else:
        vmin = 0

    if 'vmax' in kwargs:
        vmax = kwargs['vmax']
    else:
        vmax = 1e10

    ## Adapt latitude and longitude to values in parameters file
    lon_min = extent[0]
    lon_max = extent[1]
    lat_min = extent[2]
    lat_max = extent[3]

    dpt1 = dpt1.sel(latitude = slice(lat_min, lat_max), longitude = slice(lon_min, lon_max))
    zlat = zlat.sel(latitude = slice(lat_min, lat_max))
    zlon = zlon.sel(longitude = slice(lon_min, lon_max))

    ## Open Amplification Coefficient
    # check for refined bathymetry
    res_bathy = abs(zlon[1] - zlon[0])
    if res_bathy == 0.5:
        ds_ampli = xr.open_dataset('../../data/c%s.nc'%(wave_type)).astype('float64')
        amplification_coeff = ds_ampli['c%s'%wave_type]
        refined = False
    else:
        print("Refined bathymetry grid \n PLEASE RUN amplification_coefficients.ipynb before running this script")
        ds_ampli = xr.open_dataset(c_file).astype('float64')
        amplification_coeff = ds_ampli['c%s'%wave_type]        
        refined = True

    ## Surface Element
    msin = np.array([np.sin(np.pi/2 - np.radians(zlat))]).T
    ones = np.ones((1, len(zlon)))
    res_mod = radians(abs(zlat[1] - zlat[0]))
    dA = radius**2*res_mod**2*np.dot(msin,ones)

    ## Loop over dates
    YEAR = date_vec[0]
    MONTH = date_vec[1]
    DAY = date_vec[2]
    HOUR = date_vec[3]

    for iyear in np.array([YEAR]):
        if isinstance(MONTH, int):
            MONTH = np.array([MONTH])
        elif not len(MONTH):
            MONTH = np.arange(1, 13)
        else:
            MONTH = np.array(MONTH)
        for imonth in MONTH:
            TOTAL_month = np.zeros(dpt1.shape)  # Initiate monthly source of Rayleigh wave matrix
            daymax = monthrange(iyear,imonth)[1]
            filename_p2l = '%s/%s_%d%02d_p2l.nc'%(ww3_local_path, prefix, iyear, imonth)
            print("File WW3 ", filename_p2l)
            try:
                day = np.array(DAY)
                if day[0] > day[-1]:
                    index = np.squeeze(np.argwhere(day==daymax))
                    if imonth == MONTH[0]:
                        day = day[:index+1]
                    elif imonth == MONTH[-1]:
                        index = np.squeeze(np.argwhere(day==monthrange(iyear, imonth-1)[1]))
                        day = day[index+1:]
                    else:
                        day = np.arange(1,daymax+1)
            except:
                try:
                    day = int(np.squeeze(day))
                except:
                    day= np.arange(1,(monthrange(iyear,imonth)[1])+1)
            for iday in day:
                if isinstance(HOUR, int):
                    HOUR = np.array([HOUR])
                elif not len(HOUR):
                    HOUR = np.arange(0,24,3)
                else:
                    HOUR = np.array(HOUR)
                for ih in HOUR:

                    ## Open F_p3D 
                    (lati, longi, freq_ocean, p2l, unit1) = read_p2l(filename_p2l, [iyear, imonth, iday, ih], [lon_min, lon_max], [lat_min, lat_max])
                    nf = len(freq_ocean)  # number of frequencies 
                    xfr = np.exp(np.log(freq_ocean[-1]/freq_ocean[0])/(nf-1))  # determines the xfr geometric progression factor
                    df = freq_ocean*0.5*(xfr-1/xfr)  # frequency interval in wave model times 2
                    freq_seismic = 2*freq_ocean  # ocean to seismic waves freq
                    ## Check units of the model, depends on version
                    if unit1 == 'log10(Pa2 m2 s+1E-12':
                        p2l = np.exp(lg10*p2l)  - (1e-12-1e-16)
                    elif unit1 == 'log10(m4s+0.01':
                        p2l = np.exp(lg10*p2l) - 0.009999
                    elif unit1 == 'log10(Pa2 m2 s+1E-12)':
                        p2l = np.exp(lg10*p2l)  - (1e-12-1e-16)

                    ## Integral over a frequency band

                    if f1 < f2:
                        index_freq = np.logical_and(f1 <= freq_seismic, freq_seismic <= f2)
                        df = df[index_freq]
                        freq_seismic = freq_seismic[index_freq]
                        n_freq = len(df)
                        Fp = p2l.sel(frequency = freq_ocean[index_freq], latitude = slice(lat_min, lat_max), longitude = slice(lon_min, lon_max))
                        Fp = Fp.where(np.isfinite(Fp))
                        Fp = Fp.drop('frequency')
                        Fp.coords['frequency'] = freq_seismic
                        res_p2l = abs(Fp.longitude[1]-Fp.longitude[0])
                        if res_bathy != res_p2l:
                            # interpolate Fp
                            Fp = Fp.interp(longitude=zlon, latitude=zlat, method='linear')
                        amplification_coeff = amplification_coeff.sel(frequency = freq_seismic, method='nearest', tolerance=0.01)
                        amplification_coeff = amplification_coeff.sel(latitude = slice(lat_min, lat_max), longitude = slice(lon_min, lon_max))
                        amplification_coeff = amplification_coeff.reindex_like(Fp, method='nearest', tolerance=0.01)
                        F_f = Fp*amplification_coeff**2
                        F = F_f.copy()
                        for ifq, fq in enumerate(freq_seismic):
                            F_f[ifq, :, :] *= dA
                            F[ifq, :, :] = F_f[ifq, :, :]*df[ifq]
                        F = 2*np.pi*np.sqrt(F.sum(dim = 'frequency'))
                    ## Single frequency
                    elif f1 == f2:
                        index_freq = np.squeeze(np.argmin(abs(freq-f1)))
                        print('unique frequency ', f1)
                        return

                    ## Exception in parametrization of frequencies
                    else:
                        print('two frequencies with f2 < f1 were given')
                        return

                    ## Save F to file
                    if save == True:
                        path_out = './F/'
                        if not os.path.exists(path_out):
                            print("make directory %s"%path_out)
                            os.makedirs(path_out)

                        # Create netCDF 3-hourly file
                        ncfile = Dataset(path_out+"F_%d%02d%02d%02d.nc"%(iyear, imonth, iday, ih), mode='w',format='NETCDF4_CLASSIC')
                        lat_dim = ncfile.createDimension('latitude', len(zlat))  # latitude axis
                        lon_dim = ncfile.createDimension('longitude', len(zlon))  # longitude axis
                        time_dim = ncfile.createDimension('time', daymax*8)  # unlimited axis (can be appended to).
                        freq_dim = ncfile.createDimension('frequency', n_freq)
                        ncfile.title='Equivalent Vertical Force on %d-%02d-%02d-%02d'%(iyear, imonth, iday, ih)
                        ncfile.subtitle='Equivalent Force every 3 hours for the secondary microseismic peak'
                        lat = ncfile.createVariable('latitude', np.float64, ('latitude',))
                        lat.units = 'degrees_north'
                        lat.long_name = 'latitude'
                        lon = ncfile.createVariable('longitude', np.float64, ('longitude',))
                        lon.units = 'degrees_east'
                        lon.long_name = 'longitude'
                        time = ncfile.createVariable('time', np.float64, ('time',))
                        time.units = 'hours since 1990-01-01'
                        time.long_name = 'time'
                        freq_nc = ncfile.createVariable('frequency', np.float64, ('frequency',))
                        freq_nc.units = 'Hz'
                        freq_nc.long_name = 'frequency'
                        F_freq = ncfile.createVariable('F_f', np.float64, ('frequency','latitude', 'longitude'))
                        F_freq.units = 'N.s^{1/2}'
                        F_freq.long_name = 'Equivalent Vertical Force spectrum'
                        lat[:] = zlat
                        lon[:] = zlon
                        freq_nc[:] = freq_seismic
                        F_freq[:] = F_f
                        time[:] = date2num(datetime(iyear, imonth, iday, ih), units='hours since 1990-01-01', calendar='standard')
                        ncfile.close()

                    ## Plot F
                    if plot_hourly == True:
                        plt.close('all') 
                        F_plot = xr.DataArray(F, 
                                                coords={'latitude': zlat,'longitude': zlon}, 
                                                dims=["latitude", "longitude"],
                                                name = 'Frequency %.3f-%.3f Hz.%d-%02d-%02dT%02d\n %s waves.\n'%(f1, f2, iyear, imonth, iday, ih, wave_type))
                        fig = plt.figure(figsize=(9,6))
                        fig.suptitle('Frequency %.3f-%.3f Hz.%d-%02d-%02dT%02d\n %s waves.\n'%(f1, f2, iyear, imonth, iday, ih, wave_type))
                        ax = plt.axes(projection=ccrs.Robinson())
                        ax.coastlines()
                        gl = ax.gridlines()
                        gl.xformatter = LONGITUDE_FORMATTER
                        gl.yformatter = LATITUDE_FORMATTER
                        ax.add_feature(cartopy.feature.LAND, zorder=100, edgecolor='k', facecolor='linen')
                        if wave_type == 'P':
                            F_plot.plot(ax=ax, transform=ccrs.PlateCarree(),  cbar_kwargs={'label':'F (N)', 'orientation': 'horizontal'}, vmin=vmin, vmax=vmax)
                        else:
                            F_plot.plot(ax=ax, transform=ccrs.PlateCarree(),  cbar_kwargs={'label':'F (N)', 'orientation': 'horizontal'}, vmin=vmin, vmax=vmax)
                        plt.savefig('F_%s_%d%02d%02dT%02d.png'%(wave_type, iyear, imonth, iday, ih), dpi = 300, bbox_inches='tight')

                    ## Sum F
                    if plot_daily == True:
                        F_daily += F
                    if plot_monthly == True:
                        F_monthly += F
                    if plot_yearly == True:
                        F_yearly += F

                if plot_daily == True:
                    plt.close('all') 
                    F_plot = xr.DataArray(F_daily, 
                        coords={'latitude': zlat,'longitude': zlon}, 
                        dims=["latitude", "longitude"],
                        name = 'Frequency %.3f-%.3f Hz.%d-%02d-%02d\n %s waves.\n'%(f1, f2, iyear, imonth, iday, wave_type))
                    fig = plt.figure(figsize=(9,6))
                    fig.suptitle('Frequency %.3f-%.3f Hz.%d-%02d-%02d.\n %s waves.\n'%(f1, f2, iyear, imonth, iday, wave_type))
                    ax = plt.axes(projection=ccrs.Robinson())
                    ax.coastlines()
                    gl = ax.gridlines()
                    gl.xformatter = LONGITUDE_FORMATTER
                    gl.yformatter = LATITUDE_FORMATTER
                    ax.add_feature(cartopy.feature.LAND, zorder=100, edgecolor='k', facecolor='linen')
                    if wave_type == 'P':
                        F_plot.plot(ax=ax, transform=ccrs.PlateCarree(),  cbar_kwargs={'label':'F (N)', 'orientation': 'horizontal'}, vmin=vmin, vmax=vmax),
                    else:
                        F_plot.plot(ax=ax, transform=ccrs.PlateCarree(),  cbar_kwargs={'label':'F (N)', 'orientation': 'horizontal'}, vmin=vmin, vmax=vmax)
                    plt.savefig('F_%s_%d%02d%02d.png'%(wave_type, iyear, imonth, iday), dpi = 300, bbox_inches='tight')
                    F_daily = np.zeros((dpt1.shape))

            if plot_monthly == True:
                plt.close('all')
                F_plot = xr.DataArray(F_monthly, 
                    coords={'latitude': zlat,'longitude': zlon}, 
                    dims=["latitude", "longitude"],
                    name = 'Frequency %.3f-%.3f Hz.%d-%02d\n %s waves.\n'%(f1, f2, iyear, imonth, wave_type))
                fig = plt.figure(figsize=(9,6))
                fig.suptitle('Frequency %.3f-%.3f Hz.%d-%02d\n %s waves.\n'%(f1, f2, iyear, imonth, wave_type))
                ax = plt.axes(projection=ccrs.Robinson())
                ax.coastlines()
                gl = ax.gridlines()
                gl.xformatter = LONGITUDE_FORMATTER
                gl.yformatter = LATITUDE_FORMATTER
                ax.add_feature(cartopy.feature.LAND, zorder=100, edgecolor='k', facecolor='linen')
                if wave_type == 'P':
                    F_plot.plot(ax=ax, transform=ccrs.PlateCarree(),  cbar_kwargs={'label':'F (N)', 'orientation': 'horizontal'}, vmin=vmin, vmax=vmax)
                else:
                    F_plot.plot(ax=ax, transform=ccrs.PlateCarree(),  cbar_kwargs={'label':'F (N)', 'orientation': 'horizontal'}, vmin=vmin, vmax=vmax)
                plt.savefig('F_%s_%d%02d.png'%(wave_type, iyear, imonth), dpi = 300, bbox_inches='tight')
                F_monthly = np.zeros((dpt1.shape))


        if plot_yearly == True:
            plt.close('all')
            F_plot = xr.DataArray(F_yearly,
                                coords={'latitude': zlat,'longitude': zlon}, 
                                dims=["latitude", "longitude"],
                                name = 'Frequency %.3f-%.3f Hz.%d\n %s waves.\n'%(f1, f2, iyear, wave_type))
            fig = plt.figure(figsize=(9,6))
            fig.suptitle('Frequency %.3f-%.3f Hz.%d\n %s waves.\n'%(f1, f2, iyear, wave_type))
            ax = plt.axes(projection=ccrs.Robinson())
            ax.coastlines()
            gl = ax.gridlines()
            gl.xformatter = LONGITUDE_FORMATTER
            gl.yformatter = LATITUDE_FORMATTER
            ax.add_feature(cartopy.feature.LAND, zorder=100, edgecolor='k', facecolor='linen')
            if wave_type == 'P':
                F_plot.plot(ax=ax, transform=ccrs.PlateCarree(),  cbar_kwargs={'label':'F (N)', 'orientation': 'horizontal'}, vmin=vmin, vmax=vmax)
            else:
                F_plot.plot(ax=ax, transform=ccrs.PlateCarree(),  cbar_kwargs={'label':'F (N)', 'orientation': 'horizontal'}, vmin=vmin, vmax=vmax)
            plt.savefig('F_%s_%d.png'%(wave_type, iyear), dpi = 300, bbox_inches='tight')
            F_daily = np.zeros((dpt1.shape))
            F_yearly = np.zeros((dpt1.shape))
        plt.close('all')
    print('%s source maps done!'%wave_type)

def open_bathy(file_bathy = '../../data/WW3-GLOB-30M_202002_p2l.nc', refined_bathymetry=False, extent=[-180, 180, -90, 90]):
    """Open bathymetry file and optionally refine bathymetry using ETOPOv2 dataset. 

    Args:
        file_bathy (str): Path to the bathymetry file.
        refined_bathymetry (bool, optional): Whether to use the refined ETOPOv2 dataset. Defaults to False.
        extent (list, optional): The geographical extent of the bathymetry data in the format [lon_min, lon_max, lat_min, lat_max].

    Returns:
        dpt1_mask (xarray.DataArray): Masked bathymetry data.
        zlon (xarray.DataArray): Longitude coordinates.
        zlat (xarray.DataArray): Latitude coordinates.
    """
    [lon_min, lon_max, lat_min, lat_max] = extent
    ds = xr.open_mfdataset(file_bathy, combine='by_coords')
    dpt1 = ds['dpt'].squeeze(dim = 'time', drop=True)
    dpt1 = dpt1.sel(latitude = slice(lat_min, lat_max), longitude = slice(lon_min, lon_max))
    if refined_bathymetry or file_bathy == '../../data/ETOPO_2022_v1_60s_N90W180_bed.nc':
        # load refined bathymetry ETOPOv2
        file_bathy = '../../data/ETOPO_2022_v1_60s_N90W180_bed.nc'
        try:
            ds = xr.open_mfdataset(file_bathy, combine='by_coords')
            ds  = ds.rename({'lon':'longitude', 'lat': 'latitude'})
            z = ds['z']
            z *= -1 # ETOPOv2 to Depth
            z = z.where(z>0, other=np.nan)
            dpt1 = z.sel(latitude = slice(lat_min, lat_max), longitude = slice(lon_min, lon_max))         
        except:
            print("Refined bathymetry ETOPOv2 not found. \nYou can download it from:\n https://www.ngdc.noaa.gov/thredds/catalog/global/ETOPO2022/60s/60s_bed_elev_netcdf/catalog.html?dataset=globalDatasetScan/ETOPO2022/60s/60s_bed_elev_netcdf/ETOPO_2022_v1_60s_N90W180_bed.nc\nSave in ../data/")
            return None, None, None
    ## Mask nan values    
    dpt1_mask = dpt1.where(np.isfinite(dpt1))
    zlon = dpt1_mask.longitude
    zlat = dpt1_mask.latitude
    return dpt1_mask, zlon, zlat

def subfcn_liquid_solid(p, mi, mt):
    """Calculate the reflection and transmission coefficients for P and S waves between two media.
    Author: LI Lei, ll.ynyf@gmail.com modified by Pierre Boue 23/11/2020

    Args:
        p (float or array-like): The wave slowness.
        mi (tuple): The properties of the incident medium, consisting of the vertical P-wave velocity (vp1) and density (rho1).
        mt (tuple): The properties of the transmitted medium, consisting of the vertical P-wave velocity (vp2), shear wave velocity (vs2), and density (rho2).

    Returns:
        Rpp (float or array-like): The reflection coefficient for P waves.
        Tpp (float or array-like): The transmission coefficient for P waves.
        Tps (float or array-like): The transmission coefficient for S waves.
    """
    p2 = np.array(p)**2
    # vp, density of incident medium, vp, vs, density of transmitted medium
    (vp1, rho1, vp2, vs2, rho2) = (mi[0], mi[-1], mt[0], mt[1], mt[2])
    q1 = csqrt(1/vp1**2 - p2)  # vertical slowness
    q2p = csqrt(1/vp2**2 - p2)
    q2s = csqrt(1/vs2**2 - p2)

    a = rho2*q1*((1-2*p2*vs2**2)**2 + 4*vs2**4*p2*q2p*q2s)
    b = rho1*q2p
    D = a + b
    Rpp = (a-b)/D
    Tpp = 2*rho1*q1*(1-2*(p*vs2)**2)*np.reciprocal(D)
    Tps = 4*rho1*q1*q2p*p*(vs2**2)*np.reciprocal(D)
    return Rpp, Tpp, Tps