




from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve




def    lowerCamelCase__			(		A__			:   np.ndarray    ,     A__			:   np.ndarray    ,     A__			:   np.ndarray      ):





       '''simple docstring'''



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

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

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

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

       # Set pixels at invalid locations to 0
       __lowerCamelCase  =		0

       return warped




def    lowerCamelCase__			(		A__			:   np.ndarray    ,     A__			:   np.ndarray    ,     A__			:   SupportsIndex    ,     A__			:   float | None = None    ,     ):





       '''simple docstring'''



       if alpha is None:
              __lowerCamelCase  =		0.1

       # Initialize flow
       __lowerCamelCase  =		np.zeros_like(A__      )
       __lowerCamelCase  =		np.zeros_like(A__      )

       # Prepare kernels for the calculation of the derivatives and the average velocity
       __lowerCamelCase  =		np.array([[-1, 1], [-1, 1]]      ) * 0.25
       __lowerCamelCase  =		np.array([[-1, -1], [1, 1]]      ) * 0.25
       __lowerCamelCase  =		np.array([[1, 1], [1, 1]]      ) * 0.25
       __lowerCamelCase  =		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(A__      ):
              __lowerCamelCase  =		warp(A__    ,     A__    ,     A__      )
              __lowerCamelCase  =		convolve(A__    ,     A__      ) + convolve(A__    ,     A__      )
              __lowerCamelCase  =		convolve(A__    ,     A__      ) + convolve(A__    ,     A__      )
              __lowerCamelCase  =		convolve(A__    ,     A__      ) + convolve(A__    ,     -kernel_t      )

              __lowerCamelCase  =		convolve(A__    ,     A__      )
              __lowerCamelCase  =		convolve(A__    ,     A__      )

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

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

       return horizontal_flow, vertical_flow


if __name__ == "__main__":
       import doctest

       doctest.testmod()







from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor    ,     A__			:   bool      ):





       '''simple docstring'''



       if not inplace:
              __lowerCamelCase  =		ma + ma
       else:
              ma += ma

       return ma




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   List[int]      ):





       '''simple docstring'''



       __lowerCamelCase  =		-1 * len(A__      )
       __lowerCamelCase  =		list(range(len(tensor.shape[:zero_index]      )      )      )
       return tensor.permute(first_inds + [zero_index + i for i in inds]      )




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   int      ):





       '''simple docstring'''



       return t.reshape(t.shape[:-no_dims] + (-1,)      )




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor    ,     A__			:   int    ,     A__			:   float = 1E-4      ):





       '''simple docstring'''



       __lowerCamelCase  =		mask.expand(*value.shape      )
       return torch.sum(mask * value    ,     dim=A__      ) / (eps + torch.sum(A__    ,     dim=A__      ))




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.types.Number = 2.3_125    ,     A__			:   torch.types.Number = 21.6_875    ,     A__			:   int = 64      ):





       '''simple docstring'''



       __lowerCamelCase  =		torch.linspace(A__    ,     A__    ,     no_bins - 1    ,     device=pts.device      )
       __lowerCamelCase  =		torch.sqrt(torch.sum((pts.unsqueeze(-2      ) - pts.unsqueeze(-3      )) ** 2    ,     dim=-1      )      )
       return torch.bucketize(A__    ,     A__      )




def    lowerCamelCase__			(		A__			:   Callable[[list], Any]    ,     A__			:   List[dict]      ):





       '''simple docstring'''



       __lowerCamelCase  =		dicts[0]
       __lowerCamelCase  =		{}
       for k, v in first.items():
              __lowerCamelCase  =		[d[k] for d in dicts]
              if isinstance(A__    ,     A__      ):
                     __lowerCamelCase  =		dict_multimap(A__    ,     A__      )
              else:
                     __lowerCamelCase  =		fn(A__      )

       return new_dict




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor      ):





       '''simple docstring'''



       __lowerCamelCase  =		v_bins.view(((1,) * len(x.shape      )) + (len(A__      ),)      )
       __lowerCamelCase  =		x[..., None] - reshaped_bins
       __lowerCamelCase  =		torch.argmin(torch.abs(A__      )    ,     dim=-1      )
       return nn.functional.one_hot(A__    ,     num_classes=len(A__      )      ).float()




def    lowerCamelCase__			(		A__			:   torch.Tensor    ,     A__			:   torch.Tensor    ,     A__			:   int = 0    ,     A__			:   int = 0      ):





       '''simple docstring'''



       __lowerCamelCase  =		[]
       for i, s in enumerate(data.shape[:no_batch_dims]      ):
              __lowerCamelCase  =		torch.arange(A__      )
              __lowerCamelCase  =		r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape      ) - i - 1)))      )
              ranges.append(A__      )

       __lowerCamelCase  =		[slice(A__      ) for _ in range(len(data.shape      ) - no_batch_dims      )]
       __lowerCamelCase  =		inds
       ranges.extend(A__      )
       # Matt note: Editing this to get around the behaviour of using a list as an array index changing
       # in recent Numpy versions
       return data[tuple(A__      )]


UpperCAmelCase_   							=      TypeVar('T')




def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   Dict[Any, Union[dict, list, tuple, T]]    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       __lowerCamelCase  =		{}
       for k, v in dic.items():
              if isinstance(A__    ,     A__      ):
                     __lowerCamelCase  =		dict_map(A__    ,     A__    ,     A__      )
              else:
                     __lowerCamelCase  =		tree_map(A__    ,     A__    ,     A__      )

       return new_dict




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   T    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   dict    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   list    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




@overload
def    lowerCamelCase__			(		A__			:   Callable[[T], Any]    ,     A__			:   tuple    ,     A__			:   Type[T]      ):





       '''simple docstring'''



       ...




def    lowerCamelCase__			(		A__			:   Dict    ,     A__			:   List[str]    ,     A__			:   List[str]      ):





       '''simple docstring'''



       if isinstance(A__    ,     A__      ):
              return dict_map(A__    ,     A__    ,     A__      )
       elif isinstance(A__    ,     A__      ):
              return [tree_map(A__    ,     A__    ,     A__      ) for x in tree]
       elif isinstance(A__    ,     A__      ):
              return tuple(tree_map(A__    ,     A__    ,     A__      ) for x in tree      )
       elif isinstance(A__    ,     A__      ):
              return fn(A__      )
       else:
              print(type(A__      )      )
              raise ValueError("""Not supported"""      )




UpperCAmelCase_   							=      partial(tree_map, leaf_type=torch.Tensor)

