code: def test_nlls_fit_tensor():
      """
      Test the implementation of NLLS and RESTORE
      """
  
      b0 = 1000.0
      bvals, bvecs = read_bvals_bvecs(*get_fnames("55dir_grad"))
      gtab = grad.gradient_table(bvals, bvecs)
      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)
  
      # Design Matrix
      X = dti.design_matrix(gtab)
  
      # Signals
      Y = np.exp(np.dot(X, D))
      Y.shape = (-1,) + Y.shape
  
      # Estimate tensor from test signals and compare against expected result
      # using non-linear least squares:
      tensor_model = dti.TensorModel(gtab, fit_method="NLLS")
      tensor_est = tensor_model.fit(Y)
      npt.assert_equal(tensor_est.shape, Y.shape[:-1])
      npt.assert_array_almost_equal(tensor_est.evals[0], evals)
      npt.assert_array_almost_equal(
          tensor_est.quadratic_form[0], tensor
      )
      npt.assert_almost_equal(tensor_est.md[0], md)
  
      # You can also do this without the Jacobian (though it's slower):
      tensor_model = dti.TensorModel(gtab, fit_method="NLLS", jac=False)
      tensor_est = tensor_model.fit(Y)
      npt.assert_equal(tensor_est.shape, Y.shape[:-1])
      npt.assert_array_almost_equal(tensor_est.evals[0], evals)
      npt.assert_array_almost_equal(
          tensor_est.quadratic_form[0], tensor
      )
      npt.assert_almost_equal(tensor_est.md[0], md)
  
      # Using the gmm weighting scheme:
      tensor_model = dti.TensorModel(
          gtab, fit_method="NLLS", weighting="gmm", sigma=1
      )
      tensor_est = tensor_model.fit(Y)
      npt.assert_equal(tensor_est.shape, Y.shape[:-1])
      npt.assert_array_almost_equal(tensor_est.evals[0], evals)
      npt.assert_array_almost_equal(
          tensor_est.quadratic_form[0], tensor
      )
      npt.assert_almost_equal(tensor_est.md[0], md)
  
      # If you use sigma weighting, you'd better provide a sigma:
      tensor_model = dti.TensorModel(
          gtab, fit_method="NLLS", weighting="sigma"
      )
      npt.assert_raises(ValueError, tensor_model.fit, Y)
  
      # Use NLLS with some actual 4D data:
      data, bvals, bvecs = get_fnames("small_25")
      gtab = grad.gradient_table(bvals, bvecs)
      tm1 = dti.TensorModel(gtab, fit_method="NLLS")
      dd = load_nifti_data(data)
      tf1 = tm1.fit(dd)
      tm2 = dti.TensorModel(gtab)
      tf2 = tm2.fit(dd)
  
      npt.assert_array_almost_equal(tf1.fa, tf2.fa, decimal=1)
  
      # Reduce amount of data, to cause NLLS to fail
      gtab_less = grad.gradient_table(
          gtab.bvals[0:3], gtab.bvecs[0:3, :]
      )
      Y_less = Y[..., 0:3].copy()
  
      # Test warning for failure of NLLS method, resort to OLS result
      # (reason for failure: too few data points for NLLS)
      tensor_model = dti.TensorModel(
          gtab_less, fit_method="NLLS", sigma=1.0, return_S0_hat=True
      )
      tmf = npt.assert_warns(UserWarning, tensor_model.fit, Y_less)
  
      # Test fail_is_nan=True, failed NLLS method gives NaN
      tensor_model = dti.TensorModel(
          gtab_less,
          fit_method="NLLS",
          sigma=1.0,
          return_S0_hat=True,
          fail_is_nan=True,
      )
      tmf = npt.assert_warns(UserWarning, tensor_model.fit, Y_less)
      npt.assert_equal(tmf[0].S0_hat, np.nan)
  
      # Test sigma with an array
      sigma = np.ones(Y_less.shape[-1])
      tensor_model = dti.TensorModel(
          gtab_less, fit_method="NLLS", weighting="sigma", sigma=sigma
      )
      with warnings.catch_warnings():
          warnings.filterwarnings(
              "ignore", message=ols_resort_msg, category=UserWarning
          )
          tmf = tensor_model.fit(Y_less)
  
      # Test sigma with an array of wrong size
      sigma = np.ones(Y_less.shape[-1] + 10)
      tensor_model = dti.TensorModel(
          gtab_less, fit_method="NLLS", weighting="sigma", sigma=sigma
      )
      tmf = npt.assert_raises(ValueError, tensor_model.fit, Y_less)
  
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 19 for calling a global function with example usage: test_nlls_fit_tensor() without return types