def calculate_total_infections(new_infections):
  """Calculate the total infections at each location.

  Args:
    new_infections: a xr.DataArray containing new_infections and dimension
      (locations, time)

  Returns:
    total_infections: a xr.DataArray containing the summed infections of
    dimension (locations)
  """

  return new_infections.sum('time')