code: def test_compute_residual_displacement_field_ssd_3d(rng):
      # Select arbitrary images' shape (same shape for both images)
      sh = (20, 15, 10)
  
      # Select arbitrary centers
      c_f = np.asarray(sh) / 2
      c_g = c_f + 0.5
  
      # Compute the identity vector field I(x) = x in R^2
      x_0 = np.asarray(range(sh[0]))
      x_1 = np.asarray(range(sh[1]))
      x_2 = np.asarray(range(sh[2]))
      X = np.ndarray(sh + (3,), dtype=np.float64)
      O = np.ones(sh)
      X[..., 0] = x_0[:, None, None] * O
      X[..., 1] = x_1[None, :, None] * O
      X[..., 2] = x_2[None, None, :] * O
  
      # Compute the gradient fields of F and G
      grad_F = X - c_f
      grad_G = X - c_g
  
      Fnoise = (
          rng.random(np.size(grad_F)).reshape(grad_F.shape)
          * grad_F.max()
          * 0.1
      )
      Fnoise = Fnoise.astype(floating)
      grad_F += Fnoise
  
      Gnoise = (
          rng.random(np.size(grad_G)).reshape(grad_G.shape)
          * grad_G.max()
          * 0.1
      )
      Gnoise = Gnoise.astype(floating)
      grad_G += Gnoise
  
      # The squared norm of grad_G
      sq_norm_grad_G = np.sum(grad_G**2, -1)
  
      # Compute F and G
      F = 0.5 * np.sum(grad_F**2, -1)
      G = 0.5 * sq_norm_grad_G
  
      Fnoise = rng.random(np.size(F)).reshape(F.shape) * F.max() * 0.1
      Fnoise = Fnoise.astype(floating)
      F += Fnoise
  
      Gnoise = rng.random(np.size(G)).reshape(G.shape) * G.max() * 0.1
      Gnoise = Gnoise.astype(floating)
      G += Gnoise
  
      delta_field = np.array(F - G, dtype=floating)
  
      sigma_field = rng.random(delta_field.size).reshape(
          delta_field.shape
      )
      sigma_field = sigma_field.astype(floating)
  
      # Select some pixels to force sigma_field = infinite
      inf_sigma = rng.integers(0, 2, sh[0] * sh[1] * sh[2])
      inf_sigma = inf_sigma.reshape(sh)
      sigma_field[inf_sigma == 1] = np.inf
  
      # Select an initial displacement field
      d = rng.random(grad_G.size).reshape(grad_G.shape).astype(floating)
      lambda_param = 1.5
  
      # Implementation under test
      iut = ssd.compute_residual_displacement_field_ssd_3d
  
      # In the first iteration we test the case target=None
      # In the second iteration, target is not None
      target = None
      rtol = 1e-9
      atol = 1e-4
      for it in range(2):
          # Sum of differences with the neighbors
          s = np.zeros_like(d, dtype=np.float64)
          s[:, :, :-1] += d[:, :, :-1] - d[:, :, 1:]  # right
          s[:, :, 1:] += d[:, :, 1:] - d[:, :, :-1]  # left
          s[:, :-1, :] += d[:, :-1, :] - d[:, 1:, :]  # down
          s[:, 1:, :] += d[:, 1:, :] - d[:, :-1, :]  # up
          s[:-1, :, :] += d[:-1, :, :] - d[1:, :, :]  # below
          s[1:, :, :] += d[1:, :, :] - d[:-1, :, :]  # above
          s *= lambda_param
  
          # Dot product of displacement and gradient
          dp = (
              d[..., 0] * grad_G[..., 0]
              + d[..., 1] * grad_G[..., 1]
              + d[..., 2] * grad_G[..., 2]
          )
  
          # Compute expected residual
          if target is None:
              expected = np.zeros_like(grad_G)
              for i in range(3):
                  expected[..., i] = delta_field * grad_G[..., i]
          else:
              expected = target.copy().astype(np.float64)
  
          # Expected residuals when sigma != infinity
          for i in range(3):
              expected[inf_sigma == 0, i] -= (
                  grad_G[inf_sigma == 0, i] * dp[inf_sigma == 0]
                  + sigma_field[inf_sigma == 0] * s[inf_sigma == 0, i]
              )
          # Expected residuals when sigma == infinity
          expected[inf_sigma == 1] = -1.0 * s[inf_sigma == 1]
  
          # Test residual field computation starting with residual = None
          actual = iut(
              delta_field,
              sigma_field,
              grad_G.astype(floating),
              target,
              lambda_param,
              d,
              None,
          )
          assert_allclose(actual, expected, rtol=rtol, atol=atol)
          # destroy previous result
          actual = np.ndarray(actual.shape, dtype=floating)
  
          # Test residual field computation starting with residual is not None
          iut(
              delta_field,
              sigma_field,
              grad_G.astype(floating),
              target,
              lambda_param,
              d,
              actual,
          )
          assert_allclose(actual, expected, rtol=rtol, atol=atol)
  
          # Set target for next iteration
          target = actual
  
          # Test Gauss-Seidel step with residual=None and residual=target
          for residual in [None, target]:
              expected = d.copy()
              iterate_residual_field_ssd_3d(
                  delta_field,
                  sigma_field,
                  grad_G.astype(floating),
                  residual,
                  lambda_param,
                  expected,
              )
  
              actual = d.copy()
              ssd.iterate_residual_displacement_field_ssd_3d(
                  delta_field,
                  sigma_field,
                  grad_G.astype(floating),
                  residual,
                  lambda_param,
                  actual,
              )
  
              # the numpy linear solver may differ from our custom implementation
              # we need to increase the tolerance a bit
              assert_allclose(
                  actual, expected, rtol=rtol, atol=atol * 10
              )
  
file_path: dipy/dipy/align/tests/test_sumsqdiff.py
imports: import numpy
  from dipy.align import floating
  from dipy.align import sumsqdiff
  from numpy.testing import assert_equal, assert_almost_equal, assert_array_almost_equal, assert_allclose
  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_compute_residual_displacement_field_ssd_3d(rng) without return types