code: # Copyright 2019 The OpenRadar Authors. All Rights Reserved.
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #     http://www.apache.org/licenses/LICENSE-2.0
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  
  import numpy as np
  from numba import njit, jit
  from . import compensation
  from . import utils
  
  
  def doppler_resolution(band_width, start_freq_const=77, ramp_end_time=62, idle_time_const=100, num_loops_per_frame=128,
                         num_tx_antennas=3):
      """Calculate the doppler resolution for the given radar configuration.
  
      Args:
          start_freq_const (float): Frequency chirp starting point.
          ramp_end_time (float): Frequency chirp end point.
          idle_time_const (int): Idle time between chirps.
          band_width (float): Radar config bandwidth.
          num_loops_per_frame (int): The number of loops in each frame.
          num_tx_antennas (int): The number of transmitting antennas (tx) on the radar.
  
      Returns:
          doppler_resolution (float): The doppler resolution for the given radar configuration.
  
      """
  
      light_speed_meter_per_sec = 299792458
  
      center_frequency = start_freq_const * 1e9 + band_width / 2
      chirp_interval = (ramp_end_time + idle_time_const) * 1e-6
      doppler_resolution = light_speed_meter_per_sec / (
                  2 * num_loops_per_frame * num_tx_antennas * center_frequency * chirp_interval)
  
      return doppler_resolution
  
  def separate_tx(signal, num_tx, vx_axis=1, axis=0):
      """Separate interleaved radar data from separate TX along a certain axis to account for TDM radars.
  
      Args:
          signal (ndarray): Received signal.
          num_tx (int): Number of transmit antennas.
          vx_axis (int): Axis in which to accumulate the separated data.
          axis (int): Axis in which the data is interleaved.
  
      Returns:
          ndarray: Separated received data in the
  
      """
      # Reorder the axes
      reordering = np.arange(len(signal.shape))
      reordering[0] = axis
      reordering[axis] = 0
      if not (reordering == np.arange(len(reordering))).all():  # check if has to reorder
          signal = signal.transpose(reordering)
  
      # if signal.shape[1] != num_tx * signal.shape[1]:
      #     pass
  
      out = np.concatenate([signal[i::num_tx, ...] for i in range(num_tx)], axis=vx_axis)
  
      return out.transpose(reordering)
  
  
  def doppler_processing(radar_cube,
                         num_tx_antennas=2,
                         clutter_removal_enabled=False,
                         interleaved=True,
                         window_type_2d=None,
                         accumulate=True):
      """Perform 2D FFT on the radar_cube.
  
      Interleave the radar_cube, perform optional windowing and 2D FFT on the radar_cube. Optional antenna couping
      signature removal can also be performed right before 2D FFT. In constrast to the original TI codes, CFAR and peak
      grouping are intentionally separated with 2D FFT for the easiness of debugging.
  
      Args:
          radar_cube (ndarray): Output of the 1D FFT. If not interleaved beforehand, it has the shape of
              (numChirpsPerFrame, numRxAntennas, numRangeBins). Otherwise, it has the shape of 
              (numRangeBins, numVirtualAntennas, num_doppler_bins). It is assumed that after interleaving the doppler
              dimension is located at the last axis.
          num_tx_antennas (int): Number of transmitter antennas. This affects how interleaving is performed.
          clutter_removal_enabled (boolean): Flag to enable naive clutter removal.
          interleaved (boolean): If the input radar_cube is interleaved before passing in. The default radar_cube is not
              interleaved, i.e. has the shape of (numChirpsPerFrame, numRxAntennas, numRangeBins). The interleaving
              process will transform it such that it becomes (numRangeBins, numVirtualAntennas, num_doppler_bins). Note
              that this interleaving is only applicable to TDM radar, i.e. each tx emits the chirp sequentially.
          window_type_2d (mmwave.dsp.utils.Window): Optional windowing type before doppler FFT.
          accumulate (boolean): Flag to reduce the numVirtualAntennas dimension.
      
      Returns:
          detMatrix (ndarray): (numRangeBins, num_doppler_bins) complete range-dopper information. Original datatype is
                               uint16_t. Note that azimuthStaticHeatMap can be extracted from zero-doppler index for
                               visualization.
          aoa_input (ndarray): (numRangeBins, numVirtualAntennas, num_doppler_bins) ADC data reorganized by vrx instead of
                               physical rx.
      """
  
      if interleaved:
          # radar_cube is interleaved in the first dimension (for 2 tx and 0-based indexing, odd are the chirps from tx1,
          # and even are from tx2) so it becomes ( , numVirtualAntennas, numADCSamples), where 
          # numChirpsPerFrame = num_doppler_bins * num_tx_antennas as designed.
          # Antennas associated to tx1 (Ping) are 0:4 and to tx2 (Pong) are 5:8.
          fft2d_in = separate_tx(radar_cube, num_tx_antennas, vx_axis=1, axis=0)
      else:
          fft2d_in = radar_cube
          
      # (Optional) Static Clutter Removal
      if clutter_removal_enabled:
          # fft2d_in = compensation.clutter_removal(fft2d_in, axis=0)
          fft2d_in[1:] = compensation.clutter_removal(fft2d_in[1:], axis=0)  # TODO this or above with static detection removal
  
      # transpose to (numRangeBins, numVirtualAntennas, num_doppler_bins)
      fft2d_in = np.transpose(fft2d_in, axes=(2, 1, 0))
  
      # Windowing 16x32
      if window_type_2d:
          fft2d_in = utils.windowing(fft2d_in, window_type_2d, axis=2)
  
      # It is assumed that doppler is at the last axis.
      # FFT 32x32
      fft2d_out = np.fft.fft(fft2d_in)
      aoa_input = fft2d_out
  
      # Save zero-Doppler as azimuthStaticHeatMap, watch out for the bit shift in
      # original code.
  
      # Log_2 Absolute Value
      fft2d_log_abs = np.log2(np.abs(fft2d_out))
  
      # Accumulate
      if accumulate:
          return np.sum(fft2d_log_abs, axis=1), aoa_input  # TODO divide by num_rx?
      else:
          return fft2d_log_abs, aoa_input
  
  
  def doppler_estimation(radar_cube,
                         beam_weights,
                         num_tx_antennas=2,
                         clutter_removal_enabled=False,
                         interleaved=False,
                         window_type_2d=None):
      """Perform doppler estimation on the weighted sum of range FFT output across all virtual antennas.
      
      In contrast to directly computing doppler FFT from the output of range FFT, this function combines it across all 
      the virtual receivers first using the weights generated from beamforming. Then FFT is performed and argmax is taken
      across each doppler axis to return the indices of max doppler values.
      
      Args:
          radar_cube (ndarray): Output of the 1D FFT with only ranges on detected objects. If not interleaved beforehand,
              it has the shape of (numChirpsPerFrame, numRxAntennas, numDetObjs). Otherwise, it has the shape of 
              (numDetObjs, numVirtualAntennas, num_doppler_bins). It is assumed that after interleaving the doppler
              dimension is located at the last axis.
          beam_weights (ndarray): Weights to sum up the radar_cube across the virtual receivers. It is from the
                                  beam-forming and has the shape of (numVirtualAntennas, numDetObjs)
          num_tx_antennas (int): Number of transmitter antennas. This affects how interleaving is performed.
          clutter_removal_enabled (boolean): Flag to enable naive clutter removal.
          interleaved (boolean): If the input radar_cube is interleaved before passing in. The default radar_cube is not
              interleaved, i.e. has the shape of (numChirpsPerFrame, numRxAntennas, numDetObjs). The interleaveing process
              will transform it such that it becomes (numDetObjs, numVirtualAntennas, num_doppler_bins). Note that this
              interleaving is only appliable to TDM radar, i.e. each tx emits the chirp sequentially.
          window_type_2d (string): Optional windowing type before doppler FFT.
      
      Returns:
          doppler_est (ndarray): (numDetObjs) Doppler index for each detected objects. Positive index means moving away
                                 from radar while negative index means moving towards the radar.
      """
      fft2d_in = None
      if not interleaved:
          num_doppler_bins = radar_cube.shape[0] / num_tx_antennas
          # radar_cube is interleaved in the first dimension (for 2 tx and 0-based indexing, odd are the chirps from tx1,
          # and even are from tx2) so it becomes (num_doppler_bins, numVirtualAntennas, numADCSamples), where
          # numChirpsPerFrame = num_doppler_bins * num_tx_antennas as designed.
          # Antennas associated to tx1 (Ping) are 0:4 and to tx2 (Pong) are 5:8.
          if num_tx_antennas == 2:
              fft2d_in = np.concatenate((radar_cube[0::2, ...], radar_cube[1::2, ...]), axis=1)
          elif num_tx_antennas == 3:
              fft2d_in = np.concatenate((radar_cube[0::3, ...], radar_cube[1::3, ...], radar_cube[2::3, ...]), axis=1)
  
          # transpose to (numRangeBins, numVirtualAntennas, num_doppler_bins)
          fft2d_in = np.transpose(fft2d_in, axes=(2, 1, 0))
      else:
          num_doppler_bins = radar_cube.shape[2]
  
      # (Optional) Static Clutter Removal
      if clutter_removal_enabled:
          fft2d_in = compensation.clutter_removal(fft2d_in)
  
      # Weighted sum across all virtual receivers.
      fft2d_in = np.einsum('ijk,jk->ik', fft2d_in, beam_weights)
  
      # Windowing 16x32
      if window_type_2d:
          fft2d_in = utils.windowing(fft2d_in, window_type_2d, axis=1)
  
      # It is assumed that doppler is at the last axis.
      # FFT 32x32
      doppler_est = np.fft.fft(fft2d_in)
      doppler_est = np.argmax(doppler_est, axis=1)
      doppler_est[doppler_est[:] >= num_doppler_bins] -= num_doppler_bins * 2
  
      return doppler_est
  
apis: numpy.sum numpy.fft.fft numpy.transpose numpy.abs numpy.argmax numpy.einsum numpy.concatenate