





from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve





def 						__SCREAMING_SNAKE_CASE  (    UpperCamelCase__      ,      UpperCamelCase__      ,      UpperCamelCase__			)  ->    np.ndarray:







			'''simple docstring'''

			UpperCAmelCase       		=    np.stack((horizontal_flow, vertical_flow)      ,      2			)

			# Create a grid of all pixel coordinates and subtract the flow to get the
			# target pixels coordinates
			UpperCAmelCase       		=    np.stack(
			    np.meshgrid(np.arange(0      ,      image.shape[1]			)      ,      np.arange(0      ,      image.shape[0]			)			)      ,      2			)
			UpperCAmelCase       		=    np.round(grid - flow			).astype(np.intaa			)

			# Find the locations outside of the original image
			UpperCAmelCase       		=    (grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]			))
			UpperCAmelCase       		=    0

			UpperCAmelCase       		=    image[grid[:, :, 1], grid[:, :, 0]]

			# Set pixels at invalid locations to 0
			UpperCAmelCase       		=    0

			return warped





def 						__SCREAMING_SNAKE_CASE  (    UpperCamelCase__      ,      UpperCamelCase__      ,      UpperCamelCase__      ,      UpperCamelCase__ = None      ,      )  ->    tuple[np.ndarray, np.ndarray]:







			'''simple docstring'''

			if alpha is None:
						UpperCAmelCase       		=    0.1

			# Initialize flow
			UpperCAmelCase       		=    np.zeros_like(UpperCamelCase__			)
			UpperCAmelCase       		=    np.zeros_like(UpperCamelCase__			)

			# Prepare kernels for the calculation of the derivatives and the average velocity
			UpperCAmelCase       		=    np.array([[-1, 1], [-1, 1]]			) * 0.25
			UpperCAmelCase       		=    np.array([[-1, -1], [1, 1]]			) * 0.25
			UpperCAmelCase       		=    np.array([[1, 1], [1, 1]]			) * 0.25
			UpperCAmelCase       		=    np.array(
			    [[1 / 12, 1 / 6, 1 / 12], [1 / 6, 0, 1 / 6], [1 / 12, 1 / 6, 1 / 12]]			)

			# Iteratively refine the flow
			for _ in range(UpperCamelCase__			):
						UpperCAmelCase       		=    warp(UpperCamelCase__      ,      UpperCamelCase__      ,      UpperCamelCase__			)
						UpperCAmelCase       		=    convolve(UpperCamelCase__      ,      UpperCamelCase__			) + convolve(UpperCamelCase__      ,      UpperCamelCase__			)
						UpperCAmelCase       		=    convolve(UpperCamelCase__      ,      UpperCamelCase__			) + convolve(UpperCamelCase__      ,      UpperCamelCase__			)
						UpperCAmelCase       		=    convolve(UpperCamelCase__      ,      UpperCamelCase__			) + convolve(UpperCamelCase__      ,      -kernel_t			)

						UpperCAmelCase       		=    convolve(UpperCamelCase__      ,      UpperCamelCase__			)
						UpperCAmelCase       		=    convolve(UpperCamelCase__      ,      UpperCamelCase__			)

						# This updates the flow as proposed in the paper (Step 12)
						UpperCAmelCase       		=    (
						    derivative_x * avg_horizontal_velocity
						    + derivative_y * avg_vertical_velocity
						    + derivative_t
						)
						UpperCAmelCase       		=    update / (alpha**2 + derivative_x**2 + derivative_y**2)

						UpperCAmelCase       		=    avg_horizontal_velocity - derivative_x * update
						UpperCAmelCase       		=    avg_vertical_velocity - derivative_y * update

			return horizontal_flow, vertical_flow


if __name__ == "__main__":
				import doctest

				doctest.testmod()










import textwrap

import pyarrow as pa
import pytest

from datasets import Features, Value
from datasets.packaged_modules.json.json import Json





@pytest.fixture
def 						__SCREAMING_SNAKE_CASE  (    UpperCamelCase__			)  ->    str:







			'''simple docstring'''

			UpperCAmelCase       		=    tmp_path / '''file.jsonl'''
			UpperCAmelCase       		=    textwrap.dedent(
			    '''\
        {"col_1": -1}
        {"col_1": 1, "col_2": 2}
        {"col_1": 10, "col_2": 20}
        '''			)
			with open(UpperCamelCase__      ,      '''w'''			) as f:
						f.write(UpperCamelCase__			)
			return str(UpperCamelCase__			)





@pytest.fixture
def 						__SCREAMING_SNAKE_CASE  (    UpperCamelCase__			)  ->    Union[str, Any]:







			'''simple docstring'''

			UpperCAmelCase       		=    tmp_path / '''file_utf16_encoded.jsonl'''
			UpperCAmelCase       		=    textwrap.dedent(
			    '''\
        {"col_1": -1}
        {"col_1": 1, "col_2": 2}
        {"col_1": 10, "col_2": 20}
        '''			)
			with open(UpperCamelCase__      ,      '''w'''      ,      encoding='''utf-16'''			) as f:
						f.write(UpperCamelCase__			)
			return str(UpperCamelCase__			)





@pytest.fixture
def 						__SCREAMING_SNAKE_CASE  (    UpperCamelCase__			)  ->    Optional[Any]:







			'''simple docstring'''

			UpperCAmelCase       		=    tmp_path / '''file_with_list_of_dicts.json'''
			UpperCAmelCase       		=    textwrap.dedent(
			    '''\
        [
            {"col_1": -1},
            {"col_1": 1, "col_2": 2},
            {"col_1": 10, "col_2": 20}
        ]
        '''			)
			with open(UpperCamelCase__      ,      '''w'''			) as f:
						f.write(UpperCamelCase__			)
			return str(UpperCamelCase__			)





@pytest.fixture
def 						__SCREAMING_SNAKE_CASE  (    UpperCamelCase__			)  ->    int:







			'''simple docstring'''

			UpperCAmelCase       		=    tmp_path / '''file_with_list_of_dicts_field.json'''
			UpperCAmelCase       		=    textwrap.dedent(
			    '''\
        {
            "field1": 1,
            "field2": "aabb",
            "field3": [
                {"col_1": -1},
                {"col_1": 1, "col_2": 2},
                {"col_1": 10, "col_2": 20}
            ]
        }
        '''			)
			with open(UpperCamelCase__      ,      '''w'''			) as f:
						f.write(UpperCamelCase__			)
			return str(UpperCamelCase__			)





@pytest.mark.parametrize(
    '''file_fixture, config_kwargs'''      ,      [
        ('''jsonl_file''', {}),
        ('''jsonl_file_utf16_encoded''', {'''encoding''': '''utf-16'''}),
        ('''json_file_with_list_of_dicts''', {}),
        ('''json_file_with_list_of_dicts_field''', {'''field''': '''field3'''}),
    ]      ,      )
def 						__SCREAMING_SNAKE_CASE  (    UpperCamelCase__      ,      UpperCamelCase__      ,      UpperCamelCase__			)  ->    Dict:







			'''simple docstring'''

			UpperCAmelCase       		=    Json(**UpperCamelCase__			)
			UpperCAmelCase       		=    json._generate_tables([[request.getfixturevalue(UpperCamelCase__			)]]			)
			UpperCAmelCase       		=    pa.concat_tables([table for _, table in generator]			)
			assert pa_table.to_pydict() == {"col_1": [-1, 1, 10], "col_2": [None, 2, 20]}






@pytest.mark.parametrize(
    '''file_fixture, config_kwargs'''      ,      [
        (
            '''jsonl_file''',
            {'''features''': Features({'''col_1''': Value('''int64'''			), '''col_2''': Value('''int64'''			), '''missing_col''': Value('''string'''			)}			)},
        ),
        (
            '''json_file_with_list_of_dicts''',
            {'''features''': Features({'''col_1''': Value('''int64'''			), '''col_2''': Value('''int64'''			), '''missing_col''': Value('''string'''			)}			)},
        ),
        (
            '''json_file_with_list_of_dicts_field''',
            {
                '''field''': '''field3''',
                '''features''': Features(
                    {'''col_1''': Value('''int64'''			), '''col_2''': Value('''int64'''			), '''missing_col''': Value('''string'''			)}			),
            },
        ),
    ]      ,      )
def 						__SCREAMING_SNAKE_CASE  (    UpperCamelCase__      ,      UpperCamelCase__      ,      UpperCamelCase__			)  ->    List[Any]:







			'''simple docstring'''

			UpperCAmelCase       		=    Json(**UpperCamelCase__			)
			UpperCAmelCase       		=    json._generate_tables([[request.getfixturevalue(UpperCamelCase__			)]]			)
			UpperCAmelCase       		=    pa.concat_tables([table for _, table in generator]			)
			assert pa_table.to_pydict() == {"col_1": [-1, 1, 10], "col_2": [None, 2, 20], "missing_col": [None, None, None]}



