text: #!/usr/bin/env python
  # -*- coding: utf-8 -*-
  
  import pdb
  
  # 3rd party imports
  import numpy as np
  
  from scipy import constants
  
  # Local imports
  from ..pyrf import ts_tensor_xyz
  
  __author__ = "Louis Richard"
  __email__ = "louisr@irfu.se"
  __copyright__ = "Copyright 2020-2021"
  __license__ = "MIT"
  __version__ = "2.3.7"
  __status__ = "Prototype"
  
  
  def remove_idist_background(n_i, v_gse_i, p_gse_i, n_bg_i, p_bg_i):
      r"""Removes penetrating radiation background from ion moments.
  
      Parameters
      ----------
      n_i : xarray.DataArray
          Time series of the ion density.
      v_gse_i : xarray.DataArray
          Time series of the ion bulk velocity.
      p_gse_i : xarray.DataArray
          Time series of the ion pressure tensor.
      n_bg_i : xarray.DataArray
          Time series of the background ion number density.
      p_bg_i : xarray.DataArray
          Time series of the background ion pressure scalar.
  
      Returns
      -------
      n_i_new : xarray.DataArray
          Time series of the corrected ion number density.
      v_gse_i_new : xarray.DataArray
          Time series of the corrected ion bulk velocity.
      p_gse_i : xarray.DataArray
          Time series of the corrected ion pressure tensor.
  
      References
      ----------
      MMS DIS Penetrating radiation correction methods.
  
      """
  
      m_p = constants.proton_mass
  
      # Number density
      n_i_new = n_i - n_bg_i.data
  
      # Bulk velocity
      v_gse_i_new = v_gse_i.copy()
      v_gse_i_new.data *= n_i.data[:, None] / n_i_new.data[:, None]
  
      # Pressure tensor
      p_gse_i_new = np.zeros(p_gse_i.shape)
      n_old, v_old = [n_i.data, v_gse_i.data]
      n_new, v_new = [n_i_new.data, v_gse_i_new.data]
  
      for i, j in zip([0, 1, 2, 0, 0, 1], [0, 1, 2, 1, 2, 2]):
          p_gse_i_new[:, i, j] += p_gse_i.data[:, 0, 0]
          p_gse_i_new[:, i, j] -= p_bg_i.data
          p_gse_i_new[:, i, j] += m_p * n_old * np.multiply(v_old[:, i],
                                                            v_old[:, j])
          p_gse_i_new[:, i, j] -= m_p * n_new * np.multiply(v_new[:, i],
                                                            v_new[:, j])
  
      # Pressure tensor is symmetric
      p_gse_i_new[:, 1, 0] = p_gse_i_new[:, 0, 1]
      p_gse_i_new[:, 2, 0] = p_gse_i_new[:, 0, 2]
      p_gse_i_new[:, 2, 1] = p_gse_i_new[:, 1, 2]
  
      p_gse_i_new = ts_tensor_xyz(p_gse_i.time.data, p_gse_i_new)
  
      return n_i_new, v_gse_i_new, p_gse_i_new
  