solldavid / twistpy Goto Github PK

  # number of stations
  N = coordinates.shape[0]

  # compute steering_vectors
  frequency = np.mean(freq_band)
  n_vel, n_inc, n_azi, steering_vectors = compute_steering_vectors(
      coordinates,
      reference_receiver,
      dest_epsg,
      frequency,
      velocity,
      inclination,
      azimuth,
  )

  # compute covariance matrix
  C = np.ones((N, N), dtype=np.complex128)

  # compute beamforming power
  P: np.ndarray = np.einsum(
      "sn, nk, sk->s",
      steering_vectors.conj(),
      C,
      steering_vectors,
      optimize=True,
  )

  # reshape beamforming power
  P_reshape = np.reshape(P, (n_inc, n_azi, n_vel))
  P = np.real(P_reshape)

def compute_steering_vectors(
    coordinates,
    reference_receiver,
    dest_epsg,
    frequency: float,
    velocity: Union[float, tuple] = 6000,
    inclination: Union[float, tuple] = (-90, 90, 1),
    azimuth: tuple = (0, 360, 1),
) -> None:
    r"""Precompute the steering vectors

    Compute the steering vectors for the specified parameter range. For parameters that are specified as a tuple,
    the grid search is performed over the range: (min_value, max_value, increment)

    Parameters
    ----------
        frequency : :obj:`float`
            Discrete frequency at which beamforming is performed
        velocity : :obj:`float` or :obj:`tuple`
            Specifies the velocity as a float (if known) or grid over which search is performed
        inclination : :obj:`tuple`
            Specifies inclination grid over which search is performed
        azimuth : :obj:`tuple`
            Specifies azimuth grid over which search is performed

    """
    # number of stations
    N = coordinates.shape[0]

    # compute inclination, azimuth and velocity vectors
    if isinstance(velocity, tuple):
        velocity_gridded = np.arange(
            velocity[0],
            velocity[1] + velocity[2],
            velocity[2],
        )
    else:
        velocity_gridded = np.array([velocity])

    if isinstance(inclination, tuple):
        inclination_gridded = np.radians(
            np.arange(
                inclination[0],
                inclination[1] + inclination[2],
                inclination[2],
            )
        )
    else:
        inclination_gridded = np.radians(np.array([inclination]))

    azimuth_gridded = np.radians(
        np.arange(azimuth[0], azimuth[1] + azimuth[2], azimuth[2])
    )
    n_vel = len(velocity_gridded)
    n_inc = len(inclination_gridded)
    n_azi = len(azimuth_gridded)

    # create grid
    inclination_gridded, azimuth_gridded, velocity_gridded = np.meshgrid(
        inclination_gridded, azimuth_gridded, velocity_gridded, indexing="ij"
    )

    # compute relative coordinates, and convert lat/lon/elev to y/x/z
    utm_x, utm_y = latlon_2_utm(
        coordinates[:, 0], coordinates[:, 1], dest_epsg=dest_epsg
    )
    coords_utm = np.column_stack((utm_x, utm_y, coordinates[:, 2]))
    coords = coords_utm - np.tile(coords_utm[reference_receiver, :], (N, 1))
    coords = np.asmatrix(coords)

    # compute wave vector and wave number
    wave_vector_x = (np.sin(inclination_gridded) * np.cos(azimuth_gridded)).ravel()
    wave_vector_y = (np.sin(inclination_gridded) * np.sin(azimuth_gridded)).ravel()
    wave_vector_z = (np.cos(inclination_gridded)).ravel()
    wave_vector_x, wave_vector_y, wave_vector_z = (
        np.asmatrix(wave_vector_x).T,
        np.asmatrix(wave_vector_y).T,
        np.asmatrix(wave_vector_z).T,
    )
    wave_number = (-2 * np.pi * frequency / velocity_gridded).ravel()
    wave_number = np.asmatrix(wave_number).T

    # compute steering vectors, steering_vectors = exp(i * k * np.dot(wave_vector * coords))
    steering_vectors: np.ndarray = np.exp(
        1j
        * np.multiply(
            np.tile(wave_number, (1, N)),
            (
                wave_vector_x * coords[:, 0].T
                + wave_vector_y * coords[:, 1].T
                + wave_vector_z * coords[:, 2].T
            ),
        )
    )

    # ensure that steering vectors are unit vectors
    steering_vectors = steering_vectors / np.sqrt(N)

    return n_vel, n_inc, n_azi, steering_vectors