code: def test_compute_em_demons_step_2d(rng):
      r"""
      Compares the output of the demons step in 2d against an analytical
      step. The fixed image is given by $F(x) = \frac{1}{2}||x - c_f||^2$, the
      moving image is given by $G(x) = \frac{1}{2}||x - c_g||^2$,
      $x, c_f, c_g \in R^{2}$
  
      References
      ----------
      [Vercauteren09] Vercauteren, T., Pennec, X., Perchant, A., & Ayache, N.
                      (2009). Diffeomorphic demons: efficient non-parametric
                      image registration. NeuroImage, 45(1 Suppl), S61-72.
                      doi:10.1016/j.neuroimage.2008.10.040
      """
      # Select arbitrary images' shape (same shape for both images)
      sh = (30, 20)
  
      # 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 = np.ndarray(sh + (2,), dtype=np.float64)
      O = np.ones(sh)
      X[..., 0] = x_0[:, None] * O
      X[..., 1] = x_1[None, :] * O
  
      # Compute the gradient fields of F and G
      grad_F = X - c_f
      grad_G = X - c_g
  
      # The squared norm of grad_G to be used later
      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
      delta_field = G - F
  
      # Now select an arbitrary parameter for
      # $\sigma_x$ (eq 4 in [Vercauteren09])
      sigma_x_sq = 1.5
  
      # Set arbitrary values for $\sigma_i$ (eq. 4 in [Vercauteren09])
      # The original Demons algorithm used simply |F(x) - G(x)| as an
      # estimator, so let's use it as well
      sigma_i_sq = (F - G) ** 2
  
      # Select some pixels to have special values
      random_labels = rng.integers(0, 5, sh[0] * sh[1])
      random_labels = random_labels.reshape(sh)
  
      # this label is used to set sigma_i_sq == 0 below
      random_labels[sigma_i_sq == 0] = 2
      # this label is used to set gradient == 0 below
      random_labels[sq_norm_grad_G == 0] = 2
  
      expected = np.zeros_like(grad_G)
      # Pixels with sigma_i_sq = inf
      sigma_i_sq[random_labels == 0] = np.inf
      expected[random_labels == 0, ...] = 0
  
      # Pixels with gradient!=0 and sigma_i_sq=0
      sqnrm = sq_norm_grad_G[random_labels == 1]
      sigma_i_sq[random_labels == 1] = 0
      expected[random_labels == 1, 0] = (
          delta_field[random_labels == 1]
          * grad_G[random_labels == 1, 0]
          / sqnrm
      )
      expected[random_labels == 1, 1] = (
          delta_field[random_labels == 1]
          * grad_G[random_labels == 1, 1]
          / sqnrm
      )
  
      # Pixels with gradient=0 and sigma_i_sq=0
      sigma_i_sq[random_labels == 2] = 0
      grad_G[random_labels == 2, ...] = 0
      expected[random_labels == 2, ...] = 0
  
      # Pixels with gradient=0 and sigma_i_sq!=0
      grad_G[random_labels == 3, ...] = 0
  
      # Directly compute the demons step according to eq. 4 in [Vercauteren09]
      num = (sigma_x_sq * (F - G))[random_labels >= 3]
      den = (sigma_x_sq * sq_norm_grad_G + sigma_i_sq)[
          random_labels >= 3
      ]
  
      # This is $J^{P}$ in eq. 4 [Vercauteren09]
      expected[random_labels >= 3] = -1 * np.array(
          grad_G[random_labels >= 3]
      )
      expected[random_labels >= 3, ...] *= (num / den)[..., None]
  
      # Now compute it using the implementation under test
  
      actual = np.empty_like(expected, dtype=floating)
      em.compute_em_demons_step_2d(
          np.array(delta_field, dtype=floating),
          np.array(sigma_i_sq, dtype=floating),
          np.array(grad_G, dtype=floating),
          sigma_x_sq,
          actual,
      )
  
      # Test sigma_i_sq == inf
      try:
          assert_array_almost_equal(
              actual[random_labels == 0], expected[random_labels == 0]
          )
      except AssertionError:
          raise AssertionError("Failed for sigma_i_sq == inf")
  
      # Test sigma_i_sq == 0 and gradient != 0
      try:
          assert_array_almost_equal(
              actual[random_labels == 1], expected[random_labels == 1]
          )
      except AssertionError:
          raise AssertionError(
              "Failed for sigma_i_sq == 0 and gradient != 0"
          )
  
      # Test sigma_i_sq == 0 and gradient == 0
      try:
          assert_array_almost_equal(
              actual[random_labels == 2], expected[random_labels == 2]
          )
      except AssertionError:
          raise AssertionError(
              "Failed for sigma_i_sq == 0 and gradient == 0"
          )
  
      # Test sigma_i_sq != 0 and gradient == 0
      try:
          assert_array_almost_equal(
              actual[random_labels == 3], expected[random_labels == 3]
          )
      except AssertionError:
          raise AssertionError(
              "Failed for sigma_i_sq != 0 and gradient == 0 "
          )
  
      # Test sigma_i_sq != 0 and gradient != 0
      try:
          assert_array_almost_equal(
              actual[random_labels == 4], expected[random_labels == 4]
          )
      except AssertionError:
          raise AssertionError(
              "Failed for sigma_i_sq != 0 and gradient != 0"
          )
  
file_path: dipy/dipy/align/tests/test_expectmax.py
imports: import numpy
  from numpy.testing import assert_equal, assert_array_equal, assert_array_almost_equal, assert_raises
  from dipy.align import floating
  from dipy.align import expectmax
  from dipy.testing.decorators import set_random_number_generator
label_desc: Use image node_id 1 for calling a global function with example usage: test_compute_em_demons_step_2d(rng) without return types