code: def test_tensor_model():
      fdata, fbval, fbvec = get_fnames("small_25")
      data1 = load_nifti_data(fdata)
      gtab1 = grad.gradient_table(fbval, fbvec)
      data2, gtab2 = dsi_voxels()
      for data, gtab in zip([data1, data2], [gtab1, gtab2]):
          dm = dti.TensorModel(gtab, "LS")
          dtifit = dm.fit(data[0, 0, 0])
          npt.assert_equal(dtifit.fa < 0.9, True)
          dm = dti.TensorModel(gtab, "WLS")
          dtifit = dm.fit(data[0, 0, 0])
          npt.assert_equal(dtifit.fa < 0.9, True)
          npt.assert_equal(dtifit.fa > 0, True)
          sphere = create_unit_sphere(4)
          npt.assert_equal(
              len(dtifit.odf(sphere)), len(sphere.vertices)
          )
          # Check that the multivoxel case works:
          dtifit = dm.fit(data)
  
          # Check that it works on signal that has already been normalized to S0:
          dm_to_relative = dti.TensorModel(gtab)
          if np.any(gtab.b0s_mask):
              relative_data = data[0, 0, 0] / np.mean(
                  data[0, 0, 0, gtab.b0s_mask]
              )
  
              dtifit_to_relative = dm_to_relative.fit(relative_data)
              npt.assert_almost_equal(
                  dtifit.fa[0, 0, 0], dtifit_to_relative.fa, decimal=3
              )
  
      # And smoke-test that all these operations return sensibly-shaped arrays:
      npt.assert_equal(dtifit.fa.shape, data.shape[:3])
      npt.assert_equal(dtifit.ad.shape, data.shape[:3])
      npt.assert_equal(dtifit.md.shape, data.shape[:3])
      npt.assert_equal(dtifit.rd.shape, data.shape[:3])
      npt.assert_equal(dtifit.trace.shape, data.shape[:3])
      npt.assert_equal(dtifit.mode.shape, data.shape[:3])
      npt.assert_equal(dtifit.linearity.shape, data.shape[:3])
      npt.assert_equal(dtifit.planarity.shape, data.shape[:3])
      npt.assert_equal(dtifit.sphericity.shape, data.shape[:3])
  
      # Test for the shape of the mask
      npt.assert_raises(
          ValueError, dm.fit, np.ones((10, 10, 3)), np.ones((3, 3))
      )
  
      # Make some synthetic data
      b0 = 1000.0
      bvals, bvecs = read_bvals_bvecs(*get_fnames("55dir_grad"))
      gtab = grad.gradient_table_from_bvals_bvecs(bvals, bvecs)
      # The first b value is 0., so we take the second one:
      B = bvals[1]
      # Scale the eigenvalues and tensor by the B value so the units match
      D = np.array([1.0, 1.0, 1.0, 0.0, 0.0, 1.0, -np.log(b0) * B]) / B
      evals = np.array([2.0, 1.0, 0.0]) / B
      md = evals.mean()
      tensor = from_lower_triangular(D)
      A_squiggle = tensor - (1 / 3.0) * np.trace(tensor) * np.eye(3)
      mode = (
          3
          * np.sqrt(6)
          * np.linalg.det(A_squiggle / np.linalg.norm(A_squiggle))
      )
      evals_eigh, evecs_eigh = np.linalg.eigh(tensor)
      # Sort according to eigen-value from large to small:
      evecs = evecs_eigh[:, np.argsort(evals_eigh)[::-1]]
      # Check that eigenvalues and eigenvectors are properly sorted through
      # that previous operation:
      for i in range(3):
          npt.assert_array_almost_equal(
              np.dot(tensor, evecs[:, i]), evals[i] * evecs[:, i]
          )
      # Design Matrix
      X = dti.design_matrix(gtab)
      # Signals
      Y = np.exp(np.dot(X, D))
      npt.assert_almost_equal(Y[0], b0)
      Y.shape = (-1,) + Y.shape
  
      # Test fitting with different methods:
      for fit_method in ["OLS", "WLS", "NLLS"]:
          tensor_model = dti.TensorModel(
              gtab, fit_method=fit_method, return_S0_hat=True
          )
  
          tensor_fit = tensor_model.fit(Y)
          assert tensor_fit.model is tensor_model
          npt.assert_equal(tensor_fit.shape, Y.shape[:-1])
          npt.assert_array_almost_equal(tensor_fit.evals[0], evals)
          npt.assert_array_almost_equal(
              tensor_fit.S0_hat, b0, decimal=3
          )
          # Test that the eigenvectors are correct, one-by-one:
          for i in range(3):
              # Eigenvectors have intrinsic sign ambiguity
              # (see
              # http://prod.sandia.gov/techlib/access-control.cgi/2007/076422.pdf)
              # so we need to allow for sign flips. One of the following should
              # always be true:
              npt.assert_(
                  np.all(
                      np.abs(tensor_fit.evecs[0][:, i] - evecs[:, i])
                      < 10e-6
                  )
                  or np.all(
                      np.abs(-tensor_fit.evecs[0][:, i] - evecs[:, i])
                      < 10e-6
                  )
              )
              # We set a fixed tolerance of 10e-6, similar to array_almost_equal
  
          err_msg = "Calculation of tensor from Y does not compare to "
          err_msg += "analytical solution"
          npt.assert_array_almost_equal(
              tensor_fit.quadratic_form[0], tensor, err_msg=err_msg
          )
  
          npt.assert_almost_equal(tensor_fit.md[0], md)
          npt.assert_array_almost_equal(
              tensor_fit.mode, mode, decimal=5
          )
          npt.assert_equal(tensor_fit.directions.shape[-2], 1)
          npt.assert_equal(tensor_fit.directions.shape[-1], 3)
  
      # Test error-handling:
      npt.assert_raises(
          ValueError, dti.TensorModel, gtab, fit_method="crazy_method"
      )
  
      # Test custom fit tensor method
      try:
          model = dti.TensorModel(
              gtab, fit_method=lambda *args, **kwargs: 42
          )
          fit = model.fit_method()
      except Exception as exc:
          assert (
              False
          ), f"TensorModel should accept custom fit methods: {exc}"
      assert (
          fit == 42
      ), f"Custom fit method for TensorModel returned {fit}."
  
      # Test multi-voxel data
      data = np.zeros((3, Y.shape[1]))
      # Normal voxel
      data[0] = Y
      # High diffusion voxel, all diffusing weighted signal equal to zero
      data[1, gtab.b0s_mask] = b0
      data[1, ~gtab.b0s_mask] = 0
      # Masked voxel, all data set to zero
      data[2] = 0.0
  
      tensor_model = dti.TensorModel(gtab)
      fit = tensor_model.fit(data)
      npt.assert_array_almost_equal(fit[0].evals, evals)
  
      # Return S0_test
      tensor_model = dti.TensorModel(gtab, return_S0_hat=True)
      fit = tensor_model.fit(data)
      npt.assert_array_almost_equal(fit[0].evals, evals)
      npt.assert_array_almost_equal(fit[0].S0_hat, b0)
  
      # Evals should be high for high diffusion voxel
      assert all(fit[1].evals > evals[0] * 0.9)
  
      # Evals should be zero where data is masked
      npt.assert_array_almost_equal(fit[2].evals, 0.0)
  
imports: import warnings
  import numpy
  import numpy.testing
  import scipy.optimize
  import dipy.reconst.dti
  from dipy.io.gradients import read_bvals_bvecs
  from dipy.reconst.dti import ols_resort_msg, axial_diffusivity, color_fa, fractional_anisotropy, from_lower_triangular, geodesic_anisotropy, lower_triangular, mean_diffusivity, radial_diffusivity, TensorModel, trace, linearity, planarity, sphericity, decompose_tensor, _decompose_tensor_nan
  from dipy.io.image import load_nifti_data
  from dipy.data import get_fnames, dsi_voxels, get_sphere
  from dipy.core.subdivide_octahedron import create_unit_sphere
  import dipy.core.gradients
  import dipy.core.sphere
  from dipy.sims.voxel import single_tensor
  from dipy.testing.decorators import set_random_number_generator
label_desc: Use image node_id 4 for calling a global function with example usage: test_tensor_model() without return types