<commit_before># Incorporates Zoeppritz code from:
#       Written by: <NAME>
#       Modified by: <NAME>

import numpy as np
from numpy import log, tan, sin, cos, arcsin, arccosh, radians, \
                  degrees

def zoeppritz(vp1, vs1, rho1, vp0, vs0, rho0, theta1):
    '''
    Full Zoeppritz solution, considered the definitive solution.
    Calculates the angle dependent p-wave reflectivity of an interface
    between two mediums.

    :param vp1: The p-wave velocity of the upper medium.
    :param vs1: The s-wave velocity of the upper medium.
    :param rho1: The density of the upper medium.

    :param vp0: The p-wave velocity of the lower medium.
    :param vs0: The s-wave velocity of the lower medium.
    :param rho0: The density of the lower medium.
    
    :param theta1: An array of incident angles to use for reflectivity
                   calculation [degrees].

    :returns: a vector of len(theta1) containing the reflectivity
             value corresponding to each angle.
    '''
    
    # Set the ray paramter, p
    p = sin(radians(theta1)) / vp1 # ray parameter
    
    # Calculate reflection & transmission angles for Zoeppritz
    theta1 = radians(theta1)   # Convert theta1 to radians
    theta2 = arcsin(p * vp0);      # Trans. angle of P-wave
    phi1 = arcsin(p * vs1);      # Refl. angle of converted S-wave
    phi2 = arcsin(p * vs0);      # Trans. angle of converted S-wave
                
    # Matrix form of Zoeppritz Equations... M & N are matricies
    M = np.array([ \
        [-sin(theta1), -cos(phi1), sin(theta2), cos(phi2)], \
        [cos(theta1), -sin(phi1), cos(theta2), -sin(phi2)], \
        [2 * rho1 * vs1 * sin(phi1) * cos(theta1), rho1 * vs1 *\
            (1 - 2 * sin(phi1) ** 2), \
            2 * rho0 * vs0 * sin(phi2) * cos(theta2), \
            rho0 * vs0 * (1 - 2 * sin(phi2) ** 2)], \
        [-rho1 * vp1 * (1 - 2 * sin(phi1) ** 2), rho1 * vs1 * \
             sin(2 * phi1), \
            rho0 * vp0 * (1 - 2 * sin(phi2) ** 2), -rho0 * \
              vs0 * sin(2 * phi2)]
        ], dtype='float')
    
    N = np.array([ \
        [sin(theta1), cos(phi1), -sin(theta2), -cos(phi2)], \
        [cos(theta1), -sin(phi1), cos(theta2), -sin(phi2)], \
        [2 * rho1 * vs1 * sin(phi1) * cos(theta1), rho1 * vs1 * (1 - 2 * sin(phi1) ** 2), \
            2 * rho0 * vs0 * sin(phi2) * cos(theta2), rho0 * vs0 * (1 - 2 * sin(phi2) ** 2)], \
        [rho1 * vp1 * (1 - 2 * sin(phi1) ** 2), -rho1 * vs1 * sin(2 * phi1), \
            - rho0 * vp0 * (1 - 2 * sin(phi2) ** 2), rho0 * vs0 * sin(2 * phi2)]\
        ], dtype='float')
    
    # This is the important step, calculating
    # coefficients for all modes
    # and rays result is a 4x4 matrix, we want the R[0][0]
    # element for Rpp reflectivity only
    
    if M.ndim == 3:
        zoep = np.zeros([M.shape[-1]])
        for i in range(M.shape[-1]): 
            Mi = M[..., i]
            Ni = N[..., i]
            dt = np.dot(np.linalg.inv(Mi), Ni)
            zoep[i] = dt[0][0]
    else:
        dt = np.dot(np.linalg.inv(M), N)
        zoep = dt[0][0]
    
    return zoep

def akirichards(vp1, vs1, rho1, vp2, vs2, rho2, theta1):
    """
    This is the formulation from Avseth et al.,
    Quantitative seismic interpretation,
    Cambridge University Press, 2006. Adapted for a 4-term formula. 
    See http://subsurfwiki.org/wiki/Aki-Richards_equation

    :param vp1: The p-wave velocity of the upper medium.
    :param vs1: The s-wave velocity of the upper medium.
    :param rho1: The density of the upper medium.

    :param vp2: The p-wave velocity of the lower medium.<commit_msg>scattering matrix and scattering dictionary lookup<commit_after># Incorporates Zoeppritz code from:
#       Written by: <NAME>
#       Modified by: <NAME>

import numpy as np
from numpy import log, tan, sin, cos, arcsin, arccosh, radians, \
                  degrees

def zoeppritz(vp1, vs1, rho1, vp0, vs0, rho0, theta1):
    '''
    Returns the PP reflection coefficients from the Zoeppritz 
    scattering matrix.
    Calculates the angle dependent p-wave reflectivity of an interface
    between two mediums.

    :param vp1: The p-wave velocity of the upper medium.
    :param vs1: The s-wave velocity of the upper medium.
    :param rho1: The density of the upper medium.

    :param vp0: The p-wave velocity of the lower medium.
    :param vs0: The s-wave velocity of the lower medium.
    :param rho0: The density of the lower medium.
    
    :param theta1: An array of incident angles to use for reflectivity
                   calculation [degrees].

    :returns: a vector of len(theta1) containing the reflectivity
             value corresponding to each angle.
    '''
    
    # Set the ray paramter, p
    p = sin(radians(theta1)) / vp1 # ray parameter
    
    # Calculate reflection & transmission angles for Zoeppritz
    theta1 = radians(theta1)   # Convert theta1 to radians
    theta2 = arcsin(p * vp0);      # Trans. angle of P-wave
    phi1 = arcsin(p * vs1);      # Refl. angle of converted S-wave
    phi2 = arcsin(p * vs0);      # Trans. angle of converted S-wave
                
    # Matrix form of Zoeppritz Equations... M & N are matricies
    M = np.array([ \
        [-sin(theta1), -cos(phi1), sin(theta2), cos(phi2)], \
        [cos(theta1), -sin(phi1), cos(theta2), -sin(phi2)], \
        [2 * rho1 * vs1 * sin(phi1) * cos(theta1), rho1 * vs1 *\
            (1 - 2 * sin(phi1) ** 2), \
            2 * rho0 * vs0 * sin(phi2) * cos(theta2), \
            rho0 * vs0 * (1 - 2 * sin(phi2) ** 2)], \
        [-rho1 * vp1 * (1 - 2 * sin(phi1) ** 2), rho1 * vs1 * \
             sin(2 * phi1), \
            rho0 * vp0 * (1 - 2 * sin(phi2) ** 2), -rho0 * \
              vs0 * sin(2 * phi2)]
        ], dtype='float')
    
    N = np.array([ \
        [sin(theta1), cos(phi1), -sin(theta2), -cos(phi2)], \
        [cos(theta1), -sin(phi1), cos(theta2), -sin(phi2)], \
        [2 * rho1 * vs1 * sin(phi1) * cos(theta1), rho1 * vs1 * (1 - 2 * sin(phi1) ** 2), \
            2 * rho0 * vs0 * sin(phi2) * cos(theta2), rho0 * vs0 * (1 - 2 * sin(phi2) ** 2)], \
        [rho1 * vp1 * (1 - 2 * sin(phi1) ** 2), -rho1 * vs1 * sin(2 * phi1), \
            - rho0 * vp0 * (1 - 2 * sin(phi2) ** 2), rho0 * vs0 * sin(2 * phi2)]\
        ], dtype='float')
    
    # This is the important step, calculating
    # coefficients for all modes
    # and rays result is a 4x4 matrix, we want the R[0][0]
    # element for Rpp reflectivity only
    
    if M.ndim == 3:
        zoep = np.zeros([M.shape[-1]])
        for i in range(M.shape[-1]): 
            Mi = M[..., i]
            Ni = N[..., i]
            dt = np.dot(np.linalg.inv(Mi), Ni)
            zoep[i] = dt[0][0]
    else:
        dt = np.dot(np.linalg.inv(M), N)
        zoep = dt[0][0]
    
    return zoep

def scattering_mtx(vp1, vs1, rho1, vp0, vs0, rho0, theta1):
    '''
    Full Zoeppritz solution, considered the definitive solution.
    Calculates the angle dependent p-wave reflectivity of an interface
    between two mediums.

    :param vp1: The p-wave velocity of the upper medium.
    :param vs1: The s-wave velocity of the upper medium.
    :param rho1: The density of the upper medium.

    :param vp0: The p-wave velocity of the lower medium.
    :param vs0: The s-wave velocity of the lower medium.
    :param rho0: The density of the lower medium.
    
    :param theta1: A scalar  [degrees].
    
    :returns: a 4x4 array representing the scattering matrix
                at the incident angle theta1.
    '''
    if theta1.size == 1:
        # add a dimension
        theta1 = np.expand_dims(theta1, axis=1)
    
    # Set the ray paramter, p
    p = sin(radians(theta1)) / vp1 # ray parameter
    
    # Calculate reflection & transmission angles for Zoeppritz
    theta1 = radians(theta1)   # Convert theta1 to radians
    theta2 = arcsin(p * vp0);      # Trans. angle of P-wave
    phi1 = arcsin(p * vs1);      # Refl. angle of converted S-wave
    phi2 = arcsin(p * vs0);      # Trans. angle of converted S-wave
                
    # Matrix form of Zoeppritz Equations... M & N are matricies
    M = np.array([ \
        [-sin(theta1), -cos(phi1), sin(theta2), cos(phi2)], \
        [cos(theta1), -sin(phi1), cos(theta2), -sin(phi2)], \
        [2 * rho1 * vs1 * sin(phi1) * cos(theta1), rho1 * vs1 *\
            (1 - 2 * sin(phi1) ** 2), \
            2 * rho0 * vs0 * sin(phi2) * cos(theta2), \
            rho0 * vs0 * (1 - 2 * sin(phi2) ** 2)], \
        [-rho1 * vp1 * (1 - 2 * sin(phi1) ** 2), rho1 * vs1 * \
             sin(2 * phi1), \
            rho0 * vp0 * (1 - 2 * sin(phi2) ** 2), -rho0 * \
              vs0 * sin(2 * phi2)]
        ], dtype='float')
    
    N = np.array([ \
        [sin(theta1), cos(phi1), -sin(theta2), -cos(phi2)], \
        [cos(theta1), -sin(phi1), cos(theta2), -sin(phi2)], \
        [2 * rho1 * vs1 * sin(phi1) * cos(theta1), rho1 * vs1 * (1 - 2 * sin(phi1) ** 2), \
            2 * rho0 * vs0 * sin(phi2) * cos(theta2), rho0 * vs0 * (1 - 2 * sin(phi2) ** 2)], \
        [rho1 * vp1 * (1 - 2 * sin(phi1) ** 2), -rho1 * vs1 * sin(2 * phi1), \
            - rho0 * vp0 * (1 - 2 * sin(phi2) ** 2), rho0 * vs0 * sin(2 * phi2)]\
        ], dtype='float')
    
    #print M.shape
    #print N.shape
    # This is the important step, calculating
    # coefficients for all modes
    # and rays result is a 4x4 matrix, we want to return
    # all 16 elements of the Zoeppritz scattering matrix 
    # for each angle
    
    zoep = np.zeros((4,4,M.shape[-1]))
    for i in range(M.shape[-1]): 
        Mi = M[..., i]
        Ni = N[..., i]
        dt = np.dot(np.linalg.inv(Mi), Ni)
        zoep[...,i] = dt
    
    return zoep

def scattering_dict(Q):
    """Takes the Zoepprittz scattering matrix as input and
    maps each element to a keyword in a dictionary
    :param: Q - the 4x4 scattering matrix
    """
    
    scattering_keys = {
    # downgoing incident P-waves
    'dPuP' : Q[0,0,:],  # Rpp
    'dPuS' : Q[1,0,:],  # Rps
    'dPdP' : Q[2,0,:],
    'dPdS' : Q[3,0,:],
    # downgoing incident S-waves
    'dSuP' : Q[0,1,:],
    'dSuS' : Q[1,1,:],
    'dSdP' : Q[2,1,:],
    'dSdS' : Q[3,1,:],
    # upgoing incident P-waves
    'uPuP' : Q[0,2,:],
    'uPuS' : Q[1,2,:],
    'uPdP' : Q[2,2,:],
    'uPdS' : Q[3,2,:],
    # upgoing incident S-waves
    'uSuP' : Q[0,3,:],
    'uSuS' : Q[1,3,:],
    'uSdP' : Q[2,3,:],
    'uSdS' : Q[3,3,:]
    }
    return scattering_keys

def akirichards(vp1, vs1, rho1, vp2, vs2, rho2, theta1):
    """
    This is the formulation from Avseth et al.,
    Quantitative seismic interpretation,
    Cambridge University Press, 2006. Adapted for a 4-term formula. 
    See http://subsurfwiki.org/wiki/Aki-Richards_equation

    :param vp1: The p-wave velocity of the upper medium.
    :param vs1: The s-wave velocity of the upper medium.
    :param rho1: The density of the upper medium.

    :param vp2: The p-wave velocity of the lower medium.
{'fraction_non_alphanumeric': 0.14354765161878705, 'fraction_numerical': 0.04386684906520748, 'mean_word_length': 2.0435747987787956, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, '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}}
{'dir': 'git-commits-cleaned', 'id': '5687514', 'n_tokens_mistral': 4177, 'n_tokens_neox': 3981, 'n_words': 1222}