def compute_angles()

in multiple_futures_prediction/my_utils.py [0:0]

• 8 lines of code
• 3 McCabe index (conditional complexity)

def compute_angles(x_mean: torch.Tensor, num_steps:int=3) -> torch.Tensor:
  """Compute the 2d angle of trajectories.
  Args:
    x_mean is [nSteps, nObjs, dim]
  """
  nSteps, nObjs, dim = x_mean.shape
  thetas = np.zeros( (nObjs, num_steps))
  for k in range(num_steps):
    for o in range(nObjs):
      diff = x_mean[k+1,o,:] - x_mean[k,o,:]
      thetas[o,k] = np.arctan2(diff[1], diff[0])
  return thetas.mean(axis=1)