code: def test_qti_fit(rng):
      """Test the QTI fit class."""
  
      # Generate a diffusion tensor distribution
      DTD = np.concatenate(
          (
              _isotropic_DTD(),
              _anisotropic_DTD(),
              np.array([[[3, 0, 0], [0, 0, 0], [0, 0, 0]]]),
          )
      )
  
      # Calculate the ground-truth parameter values
      S0 = 1000
      D = np.mean(DTD, axis=0)
      C = qti.dtd_covariance(DTD)
      params = np.concatenate(
          (
              np.log(S0)[np.newaxis, np.newaxis],
              qti.from_3x3_to_6x1(D),
              qti.from_6x6_to_21x1(C),
          )
      ).T
      evals, evecs = np.linalg.eig(DTD)
      avg_eval_var = np.mean(np.var(evals, axis=1))
      md = np.mean(evals)
      fa = fractional_anisotropy(np.linalg.eig(D)[0])
      v_md = np.var(np.mean(evals, axis=1))
      v_shear = avg_eval_var - np.var(np.linalg.eig(D)[0])
      v_iso = v_md + v_shear
      d_sq = qti.from_3x3_to_6x1(D) @ qti.from_3x3_to_6x1(D).T
      mean_d_sq = np.mean(
          np.matmul(
              qti.from_3x3_to_6x1(DTD),
              np.swapaxes(qti.from_3x3_to_6x1(DTD), -2, -1),
          ),
          axis=0,
      )
      c_md = v_md / np.mean(np.mean(evals, axis=1) ** 2)
      c_m = fa**2
      c_mu = 1.5 * avg_eval_var / np.mean(evals**2)
      ufa = np.sqrt(c_mu)
      c_c = c_m / c_mu
      k_bulk = (
          3
          * np.matmul(
              np.swapaxes(qti.from_6x6_to_21x1(C), -1, -2),
              qti.from_6x6_to_21x1(qti.E_bulk),
          )
          / np.matmul(
              np.swapaxes(qti.from_6x6_to_21x1(d_sq), -1, -2),
              qti.from_6x6_to_21x1(qti.E_bulk),
          )
      )[0, 0]
      k_shear = (
          6
          / 5
          * np.matmul(
              np.swapaxes(qti.from_6x6_to_21x1(C), -1, -2),
              qti.from_6x6_to_21x1(qti.E_shear),
          )
          / np.matmul(
              np.swapaxes(qti.from_6x6_to_21x1(d_sq), -1, -2),
              qti.from_6x6_to_21x1(qti.E_bulk),
          )
      )[0, 0]
      mk = k_bulk + k_shear
      k_mu = (
          6
          / 5
          * np.matmul(
              np.swapaxes(qti.from_6x6_to_21x1(mean_d_sq), -1, -2),
              qti.from_6x6_to_21x1(qti.E_shear),
          )
          / np.matmul(
              np.swapaxes(qti.from_6x6_to_21x1(d_sq), -1, -2),
              qti.from_6x6_to_21x1(qti.E_bulk),
          )
      )[0, 0]
  
      # Fit QTI
      gtab = _qti_gtab(rng)
  
      if have_cvxpy:
          for fit_method in ["OLS", "WLS", "SDPdc"]:
              qtimodel = qti.QtiModel(gtab, fit_method)
              data = qtimodel.predict(params)
              npt.assert_raises(
                  ValueError, qtimodel.fit, data, np.ones(2)
              )
              npt.assert_raises(
                  ValueError, qtimodel.fit, data, np.ones(data.shape)
              )
              for mask in [None, np.ones(data.shape[0:-1])]:
                  qtifit = qtimodel.fit(data, mask)
                  npt.assert_raises(
                      ValueError,
                      qtifit.predict,
                      gradient_table(np.zeros(3), np.zeros((3, 3))),
                  )
                  npt.assert_almost_equal(
                      qtifit.predict(gtab), data, decimal=2
                  )
                  npt.assert_almost_equal(qtifit.S0_hat, S0, decimal=2)
                  npt.assert_almost_equal(qtifit.md, md, decimal=2)
                  npt.assert_almost_equal(qtifit.v_md, v_md, decimal=2)
                  npt.assert_almost_equal(
                      qtifit.v_shear, v_shear, decimal=2
                  )
                  npt.assert_almost_equal(
                      qtifit.v_iso, v_iso, decimal=2
                  )
                  npt.assert_almost_equal(qtifit.c_md, c_md, decimal=2)
                  npt.assert_almost_equal(qtifit.c_mu, c_mu, decimal=2)
                  npt.assert_almost_equal(qtifit.ufa, ufa, decimal=2)
                  npt.assert_almost_equal(qtifit.c_m, c_m, decimal=2)
                  npt.assert_almost_equal(qtifit.fa, fa, decimal=2)
                  npt.assert_almost_equal(qtifit.c_c, c_c, decimal=2)
                  npt.assert_almost_equal(qtifit.mk, mk, decimal=2)
                  npt.assert_almost_equal(
                      qtifit.k_bulk, k_bulk, decimal=2
                  )
                  npt.assert_almost_equal(
                      qtifit.k_shear, k_shear, decimal=2
                  )
                  npt.assert_almost_equal(qtifit.k_mu, k_mu, decimal=2)
      else:
          for fit_method in ["OLS", "WLS"]:
              qtimodel = qti.QtiModel(gtab, fit_method)
              data = qtimodel.predict(params)
              npt.assert_raises(
                  ValueError, qtimodel.fit, data, np.ones(2)
              )
              npt.assert_raises(
                  ValueError, qtimodel.fit, data, np.ones(data.shape)
              )
              for mask in [None, np.ones(data.shape[0:-1])]:
                  qtifit = qtimodel.fit(data, mask)
                  npt.assert_raises(
                      ValueError,
                      qtifit.predict,
                      gradient_table(np.zeros(3), np.zeros((3, 3))),
                  )
                  npt.assert_almost_equal(qtifit.predict(gtab), data)
                  npt.assert_almost_equal(qtifit.S0_hat, S0)
                  npt.assert_almost_equal(qtifit.md, md)
                  npt.assert_almost_equal(qtifit.v_md, v_md)
                  npt.assert_almost_equal(qtifit.v_shear, v_shear)
                  npt.assert_almost_equal(qtifit.v_iso, v_iso)
                  npt.assert_almost_equal(qtifit.c_md, c_md)
                  npt.assert_almost_equal(qtifit.c_mu, c_mu)
                  npt.assert_almost_equal(qtifit.ufa, ufa)
                  npt.assert_almost_equal(qtifit.c_m, c_m)
                  npt.assert_almost_equal(qtifit.fa, fa)
                  npt.assert_almost_equal(qtifit.c_c, c_c)
                  npt.assert_almost_equal(qtifit.mk, mk)
                  npt.assert_almost_equal(qtifit.k_bulk, k_bulk)
                  npt.assert_almost_equal(qtifit.k_shear, k_shear)
                  npt.assert_almost_equal(qtifit.k_mu, k_mu)
  
imports: import numpy
  import numpy.testing
  import dipy.reconst.qti
  from dipy.core.gradients import gradient_table
  from dipy.sims.voxel import vec2vec_rotmat
  from dipy.core.sphere import disperse_charges, HemiSphere
  from dipy.reconst.dti import fractional_anisotropy
  from dipy.testing.decorators import set_random_number_generator
  from dipy.utils.optpkg import optional_package
label_desc: Use image node_id 16 for calling a global function with example usage: test_qti_fit(rng) without return types