content: import pytest
  import numpy as np
  from scipy.integrate import trapezoid
  import itertools
  from scipy.special import laguerre
  from numpy.random import rand
  from numpy.testing import assert_equal, assert_almost_equal, assert_allclose
  
  import qutip
  from qutip.core.states import coherent, fock, ket, bell_state
  from qutip.wigner import wigner, wigner_transform, _parity
  from qutip.random_objects import rand_dm, rand_ket
  
  
  class TestHusimiQ:
      @pytest.mark.parametrize('xs', ["", 1, None], ids=['str', 'int', 'none'])
      def test_failure_if_non_arraylike_coordinates(self, xs):
          state = qutip.rand_ket(4)
          valid = np.linspace(-1, 1, 5)
          with pytest.raises(TypeError) as e:
              qutip.qfunc(state, xs, valid)
          assert "must be array-like" in e.value.args[0]
          with pytest.raises(TypeError) as e:
              qutip.qfunc(state, valid, xs)
          assert "must be array-like" in e.value.args[0]
          with pytest.raises(TypeError) as e:
              qutip.QFunc(xs, valid)
          assert "must be array-like" in e.value.args[0]
          with pytest.raises(TypeError) as e:
              qutip.QFunc(valid, xs)
          assert "must be array-like" in e.value.args[0]
  
      @pytest.mark.parametrize('ndim', [2, 3])
      def test_failure_if_coordinates_not_1d(self, ndim):
          state = qutip.rand_ket(4)
          valid = np.linspace(-1, 1, 5)
          bad = valid.reshape((-1,) + (1,)*(ndim - 1))
          with pytest.raises(ValueError) as e:
              qutip.qfunc(state, bad, valid)
          assert "must be 1D" in e.value.args[0]
          with pytest.raises(ValueError) as e:
              qutip.qfunc(state, valid, bad)
          assert "must be 1D" in e.value.args[0]
          with pytest.raises(ValueError) as e:
              qutip.QFunc(bad, valid)
          assert "must be 1D" in e.value.args[0]
          with pytest.raises(ValueError) as e:
              qutip.QFunc(valid, bad)
          assert "must be 1D" in e.value.args[0]
  
      @pytest.mark.parametrize('dm', [True, False], ids=['dm', 'ket'])
      def test_failure_if_tensor_hilbert_space(self, dm):
          if dm:
              state = qutip.rand_dm([2, 2])
          else:
              state = qutip.rand_ket([2, 2])
          xs = np.linspace(-1, 1, 5)
          with pytest.raises(ValueError) as e:
              qutip.qfunc(state, xs, xs)
          assert "must not have tensor structure" in e.value.args[0]
          with pytest.raises(ValueError) as e:
              qutip.QFunc(xs, xs)(state)
          assert "must not have tensor structure" in e.value.args[0]
  
      def test_QFunc_raises_if_insufficient_memory(self):
          xs = np.linspace(-1, 1, 11)
          state = qutip.rand_ket(4)
          qfunc = qutip.QFunc(xs, xs, memory=0)
          with pytest.raises(MemoryError) as e:
              qfunc(state)
          assert e.value.args[0].startswith("Refusing to precompute")
  
      def test_qfunc_warns_if_insufficient_memory(self):
          xs = np.linspace(-1, 1, 11)
          state = qutip.rand_dm(4)
          with pytest.warns(UserWarning) as e:
              qutip.qfunc(state, xs, xs, precompute_memory=0)
          assert (
              e[0].message.args[0]
              .startswith("Falling back to iterative algorithm")
          )
  
      @pytest.mark.parametrize('obj', [
          pytest.param(np.eye(2, dtype=np.complex128), id='ndarray'),
          pytest.param([[1, 0], [0, 1]], id='list'),
          pytest.param(1, id='int'),
      ])
      def test_failure_if_not_a_Qobj(self, obj):
          xs = np.linspace(-1, 1, 11)
          with pytest.raises(TypeError) as e:
              qutip.qfunc(obj, xs, xs)
          assert e.value.args[0].startswith("state must be Qobj")
          qfunc = qutip.QFunc(xs, xs)
          with pytest.raises(TypeError) as e:
              qfunc(obj)
          assert e.value.args[0].startswith("state must be Qobj")
  
      # Use indirection so that the tests can still be collected if there's a bug
      # in the generating QuTiP functions.
      @pytest.mark.parametrize('state', [
          pytest.param(lambda: qutip.rand_super(2), id='super'),
          pytest.param(lambda: qutip.rand_ket(2).dag(), id='bra'),
          pytest.param(lambda: 1j*qutip.rand_dm(2), id='non-dm operator'),
          pytest.param(lambda: qutip.Qobj([[1, 0], [0, 0]], dims=[[2], [2, 1]]),
                       id='nonsquare dm'),
          pytest.param(lambda: qutip.operator_to_vector(qutip.qeye(2)),
                       id='operator-ket'),
          pytest.param(lambda: qutip.operator_to_vector(qutip.qeye(2)).dag(),
                       id='operator-bra'),
      ])
      def test_failure_if_not_a_state(self, state):
          xs = np.linspace(-1, 1, 11)
          state = state()
          with pytest.raises(ValueError) as e:
              qutip.qfunc(state, xs, xs)
          assert (
              e.value.args[0].startswith("state must be a ket or density matrix")
          )
          qfunc = qutip.QFunc(xs, xs)
          with pytest.raises(ValueError) as e:
              qfunc(state)
          assert (
              e.value.args[0].startswith("state must be a ket or density matrix")
          )
  
      @pytest.mark.parametrize('g', [
          pytest.param(np.sqrt(2), id='natural units'),
          pytest.param(1, id='arb units'),
      ])
      @pytest.mark.parametrize('n_ys', [5, 101])
      @pytest.mark.parametrize('n_xs', [5, 101])
      @pytest.mark.parametrize('dm', [True, False], ids=['dm', 'ket'])
      @pytest.mark.parametrize('size', [5, 32])
      def test_function_and_class_are_equivalent(self, size, dm, n_xs, n_ys, g):
          xs = np.linspace(-1, 1, n_xs)
          ys = np.linspace(0, 2, n_ys)
          state = qutip.rand_dm(size) if dm else qutip.rand_ket(size)
          function = qutip.qfunc(state, xs, ys, g)
          class_ = qutip.QFunc(xs, ys, g)(state)
          np.testing.assert_allclose(function, class_)
  
      @pytest.mark.parametrize('g', [
          pytest.param(np.sqrt(2), id='natural units'),
          pytest.param(1, id='arb units'),
      ])
      @pytest.mark.parametrize('n_ys', [5, 101])
      @pytest.mark.parametrize('n_xs', [5, 101])
      @pytest.mark.parametrize('size', [5, 32])
      def test_iterate_and_precompute_are_equivalent(self, size, n_xs, n_ys, g):
          xs = np.linspace(-1, 1, n_xs)
          ys = np.linspace(0, 2, n_ys)
          state = qutip.rand_dm(size)
          iterate = qutip.qfunc(state, xs, ys, g, precompute_memory=None)
          precompute = qutip.qfunc(state, xs, ys, g, precompute_memory=np.inf)
          np.testing.assert_allclose(iterate, precompute)
  
      @pytest.mark.parametrize('initial_size', [5, 8])
      @pytest.mark.parametrize('dm', [True, False], ids=['dm', 'ket'])
      def test_same_class_can_take_many_sizes(self, dm, initial_size):
          xs = np.linspace(-1, 1, 11)
          ys = np.linspace(0, 2, 11)
          shape = np.meshgrid(xs, ys)[0].shape
          sizes = initial_size + np.array([0, 1, -1, 4])
          qfunc = qutip.QFunc(xs, ys)
          for size in sizes:
              state = qutip.rand_dm(size) if dm else qutip.rand_ket(size)
              out = qfunc(state)
              assert isinstance(out, np.ndarray)
              assert out.shape == shape
  
      @pytest.mark.parametrize('dm_first', [True, False])
      def test_same_class_can_mix_ket_and_dm(self, dm_first):
          dms = [True, False, True, False]
          if not dm_first:
              dms = dms[::-1]
          xs = np.linspace(-1, 1, 11)
          ys = np.linspace(0, 2, 11)
          shape = np.meshgrid(xs, ys)[0].shape
          qfunc = qutip.QFunc(xs, ys)
          for dm in dms:
              state = qutip.rand_dm(4) if dm else qutip.rand_ket(4)
              out = qfunc(state)
              assert isinstance(out, np.ndarray)
              assert out.shape == shape
  
      @pytest.mark.parametrize('n_ys', [5, 101])
      @pytest.mark.parametrize('n_xs', [5, 101])
      @pytest.mark.parametrize('mix', [0.1, 0.5])
      def test_qfunc_is_linear(self, n_xs, n_ys, mix):
          xs = np.linspace(-1, 1, n_xs)
          ys = np.linspace(-1, 1, n_ys)
          qfunc = qutip.QFunc(xs, ys)
          left, right = qutip.rand_dm(5), qutip.rand_dm(5)
          qleft, qright = qfunc(left), qfunc(right)
          qboth = qfunc(mix*left + (1-mix)*right)
          np.testing.assert_allclose(mix*qleft + (1-mix)*qright, qboth)
  
      @pytest.mark.parametrize('n_ys', [5, 101])
      @pytest.mark.parametrize('n_xs', [5, 101])
      @pytest.mark.parametrize('size', [5, 32])
      def test_ket_and_dm_give_same_result(self, n_xs, n_ys, size):
          xs = np.linspace(-1, 1, n_xs)
          ys = np.linspace(-1, 1, n_ys)
          state = qutip.rand_ket(size)
          qfunc = qutip.QFunc(xs, ys)
          np.testing.assert_allclose(qfunc(state), qfunc(state.proj()))
  
      @pytest.mark.parametrize('g', [
          pytest.param(np.sqrt(2), id='natural units'),
          pytest.param(1, id='arb units'),
      ])
      @pytest.mark.parametrize('ys', [
          pytest.param(np.linspace(-1, 1, 5), id='(-1,1,5)'),
          pytest.param(np.linspace(0, 2, 3), id='(0,2,3)'),
      ])
      @pytest.mark.parametrize('xs', [
          pytest.param(np.linspace(-1, 1, 5), id='(-1,1,5)'),
          pytest.param(np.linspace(0, 2, 3), id='(0,2,3)'),
      ])
      @pytest.mark.parametrize('size', [3, 5])
      def test_against_naive_implementation(self, xs, ys, g, size):
          state = qutip.rand_dm(size)
          state_np = state.full()
          x, y = np.meshgrid(xs, ys)
          alphas = 0.5*g * (x + 1j*y)
          naive = np.empty(alphas.shape, dtype=np.float64)
          for i, alpha in enumerate(alphas.flat):
              coh = qutip.coherent(size, alpha, method='analytic').full()
              naive.flat[i] = (coh.conj().T @ state_np @ coh).real[0, 0]
          naive *= (0.5*g)**2 / np.pi
          np.testing.assert_allclose(naive, qutip.qfunc(state, xs, ys, g))
          np.testing.assert_allclose(naive, qutip.QFunc(xs, ys, g)(state))
  
  
  def test_wigner_bell1_su2parity():
      """wigner: testing the SU2 parity of the first Bell state.
      """
      psi = bell_state('00')
  
      steps = 25
      theta = np.tile(np.linspace(0, np.pi, steps), 2).reshape(2, steps)
      phi = np.tile(np.linspace(0, 2 * np.pi, steps), 2).reshape(2, steps)
      slicearray = ['l', 'l']
  
      wigner_analyt = np.zeros((steps, steps))
      for t in range(steps):
          for p in range(steps):
              wigner_analyt[t, p] = np.real(((1 + np.sqrt(3)
                                              * np.cos(theta[0, t]))
                                             * (1 + np.sqrt(3)
                                             * np.cos(theta[1, t]))
                                             + 3 * (np.sin(theta[0, t])
                                             * np.exp(-1j * phi[0, p])
                                             * np.sin(theta[1, t])
                                             * np.exp(-1j * phi[1, p])
                                             + np.sin(theta[0, t])
                                             * np.exp(1j * phi[0, p])
                                             * np.sin(theta[1, t])
                                             * np.exp(1j * phi[1, p]))
                                             + (1 - np.sqrt(3)
                                             * np.cos(theta[0, t]))
                                             * (1 - np.sqrt(3)
                                             * np.cos(theta[1, t]))) / 8.)
  
      wigner_theo = wigner_transform(psi, 0.5, False, steps, slicearray)
  
      assert (np.sum(np.abs(wigner_analyt - wigner_theo)) < 1e-11)
  
  
  @pytest.mark.slow
  def test_wigner_bell4_su2parity():
      """wigner: testing the SU2 parity of the fourth Bell state.
      """
      psi = bell_state('11')
  
      steps = 25
      slicearray = ['l', 'l']
  
      wigner_analyt = np.zeros((steps, steps))
      for t in range(steps):
          for p in range(steps):
              wigner_analyt[t, p] = -0.5
  
      wigner_theo = wigner_transform(psi, 0.5, False, steps, slicearray)
  
      assert (np.sum(np.abs(wigner_analyt - wigner_theo)) < 1e-11)
  
  
  @pytest.mark.slow
  def test_wigner_bell4_fullparity():
      """wigner: testing the parity of the fourth Bell state using the parity of
      the full space.
      """
      psi = bell_state('11')
  
      steps = 25
      slicearray = ['l', 'l']
  
      wigner_analyt = np.zeros((steps, steps))
      for t in range(steps):
          for p in range(steps):
              wigner_analyt[t, p] = -0.30901699
  
      wigner_theo = wigner_transform(psi, 0.5, True, steps, slicearray)
      assert (np.sum(np.abs(wigner_analyt - wigner_theo)) < 1e-4)
  
  
  def test_parity():
      """wigner: testing the parity function.
      """
      j = 0.5
      assert (_parity(2, j)[0, 0] - (1 - np.sqrt(3)) / 2. < 1e-11)
      assert (_parity(2, j)[0, 1] < 1e-11)
      assert (_parity(2, j)[1, 1] - (1 + np.sqrt(3)) / 2. < 1e-11)
      assert (_parity(2, j)[1, 0] < 1e-11)
  
  
  @pytest.mark.slow
  def test_wigner_pure_su2():
      """wigner: testing the SU2 wigner transformation of a pure state.
      """
      psi = (ket([1]))
      steps = 25
      theta = np.linspace(0, np.pi, steps)
      phi = np.linspace(0, 2 * np.pi, steps)
      theta = theta[None, :]
      phi = phi[None, :]
      slicearray = ['l']
  
      wigner_analyt = np.zeros((steps, steps))
      for t in range(steps):
          for p in range(steps):
              wigner_analyt[t, p] = (1 + np.sqrt(3) * np.cos(theta[0, t])) / 2.
  
      wigner_theo = wigner_transform(psi, 0.5, False, steps, slicearray)
  
      assert (np.sum(np.abs(wigner_analyt - wigner_theo)) < 1e-11)
  
  
  @pytest.mark.slow
  def test_wigner_ghz_su2parity():
      """wigner: testing the SU2 wigner transformation of the GHZ state.
      """
      psi = (ket([0, 0, 0]) + ket([1, 1, 1])) / np.sqrt(2)
  
      steps = 25
      N = 3
      theta = np.tile(np.linspace(0, np.pi, steps), N).reshape(N, steps)
      phi = np.tile(np.linspace(0, 2 * np.pi, steps), N).reshape(N, steps)
      slicearray = ['l', 'l', 'l']
  
      wigner_analyt = np.zeros((steps, steps))
      for t in range(steps):
          for p in range(steps):
              wigner_analyt[t, p] = np.real(((1 + np.sqrt(3)*np.cos(theta[0, t]))
                                             * (1 + np.sqrt(3)
                                             * np.cos(theta[1, t]))
                                             * (1 + np.sqrt(3)
                                             * np.cos(theta[2, t]))
                                             + 3**(3 / 2) * (np.sin(theta[0, t])
                                             * np.exp(-1j * phi[0, p])
                                             * np.sin(theta[1, t])
                                             * np.exp(-1j * phi[1, p])
                                             * np.sin(theta[2, t])
                                             * np.exp(-1j * phi[2, p])
                                             + np.sin(theta[0, t])
                                             * np.exp(1j * phi[0, p])
                                             * np.sin(theta[1, t])
                                             * np.exp(1j * phi[1, p])
                                             * np.sin(theta[2, t])
                                             * np.exp(1j * phi[2, p]))
                                             + (1 - np.sqrt(3)
                                             * np.cos(theta[0, t]))
                                             * (1 - np.sqrt(3)
                                             * np.cos(theta[1, t]))
                                             * (1 - np.sqrt(3)
                                             * np.cos(theta[2, t]))) / 16.)
  
      wigner_theo = wigner_transform(psi, 0.5, False, steps, slicearray)
  
      assert (np.sum(np.abs(wigner_analyt - wigner_theo)) < 1e-11)
  
  
  @pytest.mark.slow
  def test_angle_slicing():
      """wigner: tests angle slicing.
      """
      psi1 = bell_state('00')
      psi2 = bell_state('01')
      psi3 = bell_state('10')
      psi4 = bell_state('11')
  
      steps = 25
      j = 0.5
  
      wigner1 = wigner_transform(psi1, j, False, steps, ['l', 'l'])
      wigner2 = wigner_transform(psi2, j, False, steps, ['l', 'z'])
      wigner3 = wigner_transform(psi3, j, False, steps, ['l', 'x'])
      wigner4 = wigner_transform(psi4, j, False, steps, ['l', 'y'])
  
      assert (np.sum(np.abs(wigner2 - wigner1)) < 1e-11)
      assert (np.sum(np.abs(wigner3 - wigner2)) < 1e-11)
      assert (np.sum(np.abs(wigner4 - wigner3)) < 1e-11)
      assert (np.sum(np.abs(wigner4 - wigner1)) < 1e-11)
  
  
  def test_wigner_coherent():
      "wigner: test wigner function calculation for coherent states"
      xvec = np.linspace(-5.0, 5.0, 100)
      yvec = xvec
  
      X, Y = np.meshgrid(xvec, yvec)
  
      a = X + 1j * Y  # consistent with g=2 option to wigner function
  
      dx = xvec[1] - xvec[0]
      dy = yvec[1] - yvec[0]
  
      N = 20
      beta = rand() + rand() * 1.0j
      psi = coherent(N, beta)
  
      # calculate the wigner function using qutip and analytic formula
      W_qutip = wigner(psi, xvec, yvec, g=2)
      W_qutip_cl = wigner(psi, xvec, yvec, g=2, method='clenshaw')
      W_analytic = 2 / np.pi * np.exp(-2 * abs(a - beta) ** 2)
  
      # check difference
      assert (np.sum(abs(W_qutip - W_analytic) ** 2) < 1e-4)
      assert (np.sum(abs(W_qutip_cl - W_analytic) ** 2) < 1e-4)
  
      # check normalization
      assert (np.sum(W_qutip) * dx * dy - 1.0 < 1e-8)
      assert (np.sum(W_qutip_cl) * dx * dy - 1.0 < 1e-8)
      assert (np.sum(W_analytic) * dx * dy - 1.0 < 1e-8)
  
  
  def test_wigner_fock():
      "wigner: test wigner function calculation for Fock states"
  
      xvec = np.linspace(-5.0, 5.0, 100)
      yvec = xvec
  
      X, Y = np.meshgrid(xvec, yvec)
  
      a = X + 1j * Y  # consistent with g=2 option to wigner function
  
      dx = xvec[1] - xvec[0]
      dy = yvec[1] - yvec[0]
  
      N = 15
  
      for n in [2, 3, 4, 5, 6]:
  
          psi = fock(N, n)
  
          # calculate the wigner function using qutip and analytic formula
          W_qutip = wigner(psi, xvec, yvec, g=2)
          W_qutip_cl = wigner(psi, xvec, yvec, g=2, method='clenshaw')
          W_qutip_sparse = wigner(psi, xvec, yvec, g=2, sparse=True, method='clenshaw')
          W_analytic = 2 / np.pi * (-1) ** n * \
              np.exp(-2 * abs(a) ** 2) * np.polyval(laguerre(n), 4 * abs(a) ** 2)
  
          # check difference
          assert (np.sum(abs(W_qutip - W_analytic)) < 1e-4)
          assert (np.sum(abs(W_qutip_cl - W_analytic)) < 1e-4)
          assert (np.sum(abs(W_qutip_sparse - W_analytic)) < 1e-4)
  
          # check normalization
          assert (np.sum(W_qutip) * dx * dy - 1.0 < 1e-8)
          assert (np.sum(W_qutip_cl) * dx * dy - 1.0 < 1e-8)
          assert (np.sum(W_qutip_sparse) * dx * dy - 1.0 < 1e-8)
          assert (np.sum(W_analytic) * dx * dy - 1.0 < 1e-8)
  
  
  def test_wigner_compare_methods_dm():
      "wigner: compare wigner methods for random density matrices"
  
      xvec = np.linspace(-5.0, 5.0, 100)
      yvec = xvec
  
      X, Y = np.meshgrid(xvec, yvec)
  
      # a = X + 1j * Y  # consistent with g=2 option to wigner function
  
      dx = xvec[1] - xvec[0]
      dy = yvec[1] - yvec[0]
  
      N = 15
  
      for n in range(10):
          # try ten different random density matrices
  
          rho = rand_dm(N, density=0.5 + rand() / 2)
  
          # calculate the wigner function using qutip and analytic formula
          W_qutip1 = wigner(rho, xvec, yvec, g=2)
          W_qutip2 = wigner(rho, xvec, yvec, g=2, method='laguerre')
  
          # check difference
          assert (np.sum(abs(W_qutip1 - W_qutip1)) < 1e-4)
  
          # check normalization
          assert (np.sum(W_qutip1) * dx * dy - 1.0 < 1e-8)
          assert (np.sum(W_qutip2) * dx * dy - 1.0 < 1e-8)
  
  
  def test_wigner_compare_methods_ket():
      "wigner: compare wigner methods for random state vectors"
  
      xvec = np.linspace(-5.0, 5.0, 100)
      yvec = xvec
  
      X, Y = np.meshgrid(xvec, yvec)
  
      # a = X + 1j * Y  # consistent with g=2 option to wigner function
  
      dx = xvec[1] - xvec[0]
      dy = yvec[1] - yvec[0]
  
      N = 15
  
      for n in range(10):
          # try ten different random density matrices
  
          psi = rand_ket(N, density=0.5 + rand() / 2)
  
          # calculate the wigner function using qutip and analytic formula
          W_qutip1 = wigner(psi, xvec, yvec, g=2)
          W_qutip2 = wigner(psi, xvec, yvec, g=2, sparse=True)
  
          # check difference
          assert (np.sum(abs(W_qutip1 - W_qutip2)) < 1e-4)
  
          # check normalization
          assert (np.sum(W_qutip1) * dx * dy - 1.0 < 1e-8)
          assert (np.sum(W_qutip2) * dx * dy - 1.0 < 1e-8)
  
  
  def test_wigner_fft_comparse_ket():
      "Wigner: Compare Wigner fft and iterative for rand. ket"
      N = 20
      xvec = np.linspace(-10, 10, 128)
      for i in range(3):
          rho = rand_ket(N)
  
          Wfft, yvec = wigner(rho, xvec, xvec, method='fft')
          W = wigner(rho, xvec, yvec, method='iterative')
  
          Wdiff = abs(W - Wfft)
          assert_equal(np.sum(abs(Wdiff)) < 1e-7, True)
  
  
  def test_wigner_fft_comparse_dm():
      "Wigner: Compare Wigner fft and iterative for rand. dm"
      N = 20
      xvec = np.linspace(-10, 10, 128)
      for i in range(3):
          rho = rand_dm(N)
  
          Wfft, yvec = wigner(rho, xvec, xvec, method='fft')
          W = wigner(rho, xvec, yvec, method='iterative')
  
          Wdiff = abs(W - Wfft)
          assert_equal(np.sum(abs(Wdiff)) < 1e-7, True)
  
  
  def test_wigner_clenshaw_iter_dm():
      "Wigner: Compare Wigner clenshaw and iterative for rand. dm"
      N = 20
      xvec = np.linspace(-10, 10, 128)
      for i in range(3):
          rho = rand_dm(N)
  
          Wclen = wigner(rho, xvec, xvec, method='clenshaw')
          W = wigner(rho, xvec, xvec, method='iterative')
  
          Wdiff = abs(W - Wclen)
          assert_equal(np.sum(abs(Wdiff)) < 1e-7, True)
  
  
  def test_wigner_clenshaw_sp_iter_dm():
      "Wigner: Compare Wigner sparse clenshaw and iterative for rand. dm"
      N = 20
      xvec = np.linspace(-10, 10, 128)
      for i in range(3):
          rho = rand_dm(N)
  
          Wclen = wigner(rho, xvec, xvec, method='clenshaw', sparse=True)
          W = wigner(rho, xvec, xvec, method='iterative')
  
          Wdiff = abs(W - Wclen)
          assert_equal(np.sum(abs(Wdiff)) < 1e-7, True)
  
  @pytest.mark.parametrize(['spin'], [
      pytest.param(1/2, id="spin-one-half"),
      pytest.param(3, id="spin-three"),
      pytest.param(13/2, id="spin-thirteen-half"),
      pytest.param(7, id="spin-seven")
  ])
  @pytest.mark.parametrize("pure", [
      pytest.param("pure", id="pure"),
      pytest.param("herm", id="mixed")
  ])
  def test_spin_q_function(spin, pure):
      d = int(2*spin + 1)
      rho = rand_dm(d, distribution=pure)
  
      # Points at which to evaluate the spin Q function
      theta = np.linspace(0, np.pi, 16, endpoint=True)
      phi = np.linspace(-np.pi, np.pi, 32, endpoint=True)
      Q, _, _ = qutip.spin_q_function(rho, theta, phi)
  
      for k, (phi_prime, theta_prime) in enumerate(itertools.product(phi, theta)):
          state = qutip.spin_coherent(spin, theta_prime, phi_prime)
          direct_Q = abs(state.dag() * rho * state)
          assert_almost_equal(Q.flat[k], direct_Q, decimal=9)
  
  @pytest.mark.parametrize(['spin'], [
      pytest.param(1/2, id="spin-one-half"),
      pytest.param(3, id="spin-three"),
      pytest.param(13/2, id="spin-thirteen-half"),
      pytest.param(7, id="spin-seven")
  ])
  @pytest.mark.parametrize("pure", [
      pytest.param("pure", id="pure"),
      pytest.param("herm", id="mixed")
  ])
  def test_spin_q_function_normalized(spin, pure):
      d = int(2 * spin + 1)
      rho = rand_dm(d, distribution=pure)
  
      # Points at which to evaluate the spin Q function
      theta = np.linspace(0, np.pi, 128, endpoint=True)
      phi = np.linspace(-np.pi, np.pi, 256, endpoint=True)
      Q, THETA, _ = qutip.spin_q_function(rho, theta, phi)
  
      norm = d / (4 * np.pi) * trapezoid(
          trapezoid(Q * np.sin(THETA), theta), phi
      )
      assert_allclose(norm, 1, atol=2e-4)
  
  
  @pytest.mark.parametrize(["spin"], [
      pytest.param(1/2, id="spin-one-half"),
      pytest.param(1, id="spin-one"),
      pytest.param(3/2, id="spin-three-half"),
      pytest.param(2, id="spin-two")
  ])
  @pytest.mark.parametrize("pure", [
      pytest.param("pure", id="pure"),
      pytest.param("herm", id="mixed")
  ])
  def test_spin_wigner_normalized(spin, pure):
      d = int(2*spin + 1)
      rho = rand_dm(d, distribution=pure)
  
      # Points at which to evaluate the spin Wigner function
      theta = np.linspace(0, np.pi, 256, endpoint=True)
      phi = np.linspace(-np.pi, np.pi, 512, endpoint=True)
      W, THETA, PHI = qutip.spin_wigner(rho, theta, phi)
  
      norm = trapezoid(
          trapezoid(W * np.sin(THETA) * np.sqrt(d / (4*np.pi)), theta), phi
      )
      assert_almost_equal(norm, 1, decimal=4)
  
  @pytest.mark.parametrize(['spin'], [
      pytest.param(1 / 2, id="spin-one-half"),
      pytest.param(1, id="spin-one"),
      pytest.param(3 / 2, id="spin-three-half"),
      pytest.param(2, id="spin-two")
  ])
  @pytest.mark.parametrize("pure", [
      pytest.param("pure", id="pure"),
      pytest.param("herm", id="mixed")
  ])
  def test_spin_wigner_overlap(spin, pure, n=5):
      d = int(2*spin + 1)
      rho = rand_dm(d, distribution=pure)
  
      # Points at which to evaluate the spin Wigner function
      theta = np.linspace(0, np.pi, 256, endpoint=True)
      phi = np.linspace(-np.pi, np.pi, 512, endpoint=True)
      W, THETA, _ = qutip.spin_wigner(rho, theta, phi)
  
      for k in range(n):
          test_state = rand_dm(d)
          state_overlap = (test_state*rho).tr().real
  
          W_state, _, _ = qutip.spin_wigner(test_state, theta, phi)
          W_overlap = trapezoid(
              trapezoid(W_state * W * np.sin(THETA), theta), phi).real
          assert_almost_equal(W_overlap, state_overlap, decimal=4)
  