def create_motion_field(input_precip, motion_type):
    """
    Create idealized motion fields to be applied to the reference image.

    Parameters
    ----------

    input_precip: numpy array (lat, lon)

    motion_type: str
        The supported motion fields are:

            - linear_x: (u=2, v=0)
            - linear_y: (u=0, v=2)
            - rotor: rotor field

    Returns
    -------
    ideal_motion : numpy array (u, v)
    """

    # Create an imaginary grid on the image and create a motion field to be
    # applied to the image.
    ny, nx = input_precip.shape

    x_pos = np.arange(nx)
    y_pos = np.arange(ny)
    x, y = np.meshgrid(x_pos, y_pos, indexing="ij")

    ideal_motion = np.zeros((2, nx, ny))

    if motion_type == "linear_x":
        ideal_motion[0, :] = 2  # Motion along x
    elif motion_type == "linear_y":
        ideal_motion[1, :] = 2  # Motion along y
    elif motion_type == "rotor":
        x_mean = x.mean()
        y_mean = y.mean()
        norm = np.sqrt(x * x + y * y)
        mask = norm != 0
        ideal_motion[0, mask] = 2 * (y - y_mean)[mask] / norm[mask]
        ideal_motion[1, mask] = -2 * (x - x_mean)[mask] / norm[mask]
    else:
        raise ValueError("motion_type not supported.")

    # We need to swap the axes because the optical flow methods expect
    # (lat, lon) or (y,x) indexing convention.
    ideal_motion = ideal_motion.swapaxes(1, 2)
    return ideal_motion
Create idealized motion fields to be applied to the reference image.

Parameters
----------

input_precip: numpy array (lat, lon)

motion_type: str
The supported motion fields are:

- linear_x: (u=2, v=0)
- linear_y: (u=0, v=2)
- rotor: rotor field

Returns
-------
ideal_motion : numpy array (u, v)
examples/optical_flow_methods_convergence.py