{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test filters on masked arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from openpiv import filters\n",
    "from openpiv.lib import replace_nans\n",
    "from typing import Optional, Tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def replace_outliers(\n",
    "    u: np.ndarray,\n",
    "    v: np.ndarray,\n",
    "    invalid_mask: np.ndarray,\n",
    "    grid_mask: np.ndarray,\n",
    "    w: Optional[np.ndarray]=None,\n",
    "    method: str=\"localmean\",\n",
    "    max_iter: int=5,\n",
    "    tol: float=1e-3,\n",
    "    kernel_size: int=1,\n",
    "    )-> Tuple[np.ndarray, np.ndarray,np.ndarray]:\n",
    "    \"\"\"Replace invalid vectors in an velocity field using an iterative image\n",
    "        inpainting algorithm.\n",
    "\n",
    "    The algorithm is the following:\n",
    "\n",
    "    1) For each element in the arrays of the ``u`` and ``v`` components,\n",
    "       replace it by a weighted average\n",
    "       of the neighbouring elements which are not invalid themselves. The\n",
    "       weights depends of the method type. If ``method=localmean`` weight\n",
    "       are equal to 1/( (2*kernel_size+1)**2 -1 )\n",
    "\n",
    "    2) Several iterations are needed if there are adjacent invalid elements.\n",
    "       If this is the case, inforation is \"spread\" from the edges of the\n",
    "       missing regions iteratively, until the variation is below a certain\n",
    "       threshold.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "\n",
    "    u : 2d or 3d np.ndarray\n",
    "        the u velocity component field\n",
    "\n",
    "    v : 2d or 3d  np.ndarray\n",
    "        the v velocity component field\n",
    "\n",
    "    w : 2d or 3d  np.ndarray\n",
    "        the w velocity component field\n",
    "\n",
    "    max_iter : int\n",
    "        the number of iterations\n",
    "\n",
    "    kernel_size : int\n",
    "        the size of the kernel, default is 1\n",
    "\n",
    "    method : str\n",
    "        the type of kernel used for repairing missing vectors\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    uf : 2d or 3d np.ndarray\n",
    "        the smoothed u velocity component field, where invalid vectors have\n",
    "        been replaced\n",
    "\n",
    "    vf : 2d or 3d np.ndarray\n",
    "        the smoothed v velocity component field, where invalid vectors have\n",
    "        been replaced\n",
    "\n",
    "    wf : 2d or 3d np.ndarray\n",
    "        the smoothed w velocity component field, where invalid vectors have\n",
    "        been replaced\n",
    "\n",
    "    \"\"\"\n",
    "    # we shall now replace NaNs only at invalid_mask positions,\n",
    "    # regardless the grid_mask (which is a user-provided masked region)\n",
    "\n",
    "    u[invalid_mask] = np.nan\n",
    "    v[invalid_mask] = np.nan\n",
    "    wf = np.empty_like(u)\n",
    "    \n",
    "    uf = replace_nans(\n",
    "        u, method=method, max_iter=max_iter, tol=tol,\n",
    "        kernel_size=kernel_size\n",
    "    )\n",
    "    vf = replace_nans(\n",
    "        v, method=method, max_iter=max_iter, tol=tol,\n",
    "        kernel_size=kernel_size\n",
    "    )\n",
    "\n",
    "    if isinstance(w, np.ndarray):\n",
    "        w[invalid_mask] = np.nan\n",
    "        wf = replace_nans(\n",
    "            w, method=method, max_iter=max_iter, tol=tol,\n",
    "            kernel_size=kernel_size\n",
    "        )\n",
    "\n",
    "    \n",
    "    # reinforce grid_mask\n",
    "    uf = np.ma.masked_array(uf, mask=grid_mask)\n",
    "    vf = np.ma.masked_array(vf, mask=grid_mask)\n",
    "    wf = np.ma.masked_array(wf, mask=grid_mask)\n",
    "\n",
    "    return uf, vf, wf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(masked_array(\n",
       "   data=[[1.0, 1.0, 1.0, 1.0, 1.0],\n",
       "         [1.0, nan, 1.0, 1.0, 1.0],\n",
       "         [1.0, 1.0, 1.0, 1.0, 1.0],\n",
       "         [1.0, 1.0, 1.0, --, --],\n",
       "         [1.0, 1.0, 1.0, --, --]],\n",
       "   mask=[[False, False, False, False, False],\n",
       "         [False, False, False, False, False],\n",
       "         [False, False, False, False, False],\n",
       "         [False, False, False,  True,  True],\n",
       "         [False, False, False,  True,  True]],\n",
       "   fill_value=1e+20),\n",
       " array([[False, False, False, False, False],\n",
       "        [False, False, False, False, False],\n",
       "        [False, False, False, False, False],\n",
       "        [False, False, False,  True,  True],\n",
       "        [False, False, False,  True,  True]]))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# bottom right corner is masked\n",
    "u = np.ma.copy(np.ones((5,5)))\n",
    "u[3:,3:] = np.ma.masked\n",
    "u[1,1] = np.nan\n",
    "grid_mask = u.mask.copy()\n",
    "u, grid_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "masked_array(\n",
       "  data=[[1.0, 1.0, 1.0, 1.0, 1.0],\n",
       "        [1.0, 1.0, 1.0, 1.0, 1.0],\n",
       "        [1.0, 1.0, 1.0, 1.0, 1.0],\n",
       "        [1.0, 1.0, 1.0, --, --],\n",
       "        [1.0, 1.0, 1.0, --, --]],\n",
       "  mask=[[False, False, False, False, False],\n",
       "        [False, False, False, False, False],\n",
       "        [False, False, False, False, False],\n",
       "        [False, False, False,  True,  True],\n",
       "        [False, False, False,  True,  True]],\n",
       "  fill_value=1e+20)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# verify that replace_nans does not know about the mask\n",
    "replace_nans(u, 1, 1e-3) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(masked_array(\n",
       "   data=[[1.0, 1.0, 1.0, 1.0, 1.0],\n",
       "         [1.0, nan, 1.0, 1.0, 1.0],\n",
       "         [1.0, 1.0, 1.0, 1.0, 1.0],\n",
       "         [1.0, 1.0, 1.0, --, --],\n",
       "         [1.0, 1.0, 1.0, --, --]],\n",
       "   mask=[[False, False, False, False, False],\n",
       "         [False, False, False, False, False],\n",
       "         [False, False, False, False, False],\n",
       "         [False, False, False,  True,  True],\n",
       "         [False, False, False,  True,  True]],\n",
       "   fill_value=1e+20),\n",
       " array([[False, False, False, False, False],\n",
       "        [False, False, False, False, False],\n",
       "        [False, False, False, False, False],\n",
       "        [False, False, False,  True,  True],\n",
       "        [False, False, False,  True,  True]]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u, grid_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(masked_array(\n",
       "   data=[[1.0, 1.0, 1.0, 1.0, 1.0],\n",
       "         [1.0, nan, 1.0, 1.0, 1.0],\n",
       "         [1.0, 1.0, 1.0, 1.0, 1.0],\n",
       "         [1.0, 1.0, 1.0, nan, --],\n",
       "         [1.0, 1.0, 1.0, --, --]],\n",
       "   mask=[[False, False, False, False, False],\n",
       "         [False, False, False, False, False],\n",
       "         [False, False, False, False, False],\n",
       "         [False, False, False, False,  True],\n",
       "         [False, False, False,  True,  True]],\n",
       "   fill_value=1e+20),\n",
       " array([[False, False, False, False, False],\n",
       "        [False, False, False, False, False],\n",
       "        [False, False, False, False, False],\n",
       "        [False, False, False,  True,  True],\n",
       "        [False, False, False,  True,  True]]))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# but if we put nan on the masked region, it \n",
    "# destroys the mask and the replace_nans also works there\n",
    "# we shall later re-enforce the mask\n",
    "u[3,3] = np.nan\n",
    "u, grid_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[False, False, False, False, False],\n",
       "       [False,  True, False, False, False],\n",
       "       [False, False, False, False, False],\n",
       "       [False, False, False,  True, False],\n",
       "       [False, False, False, False, False]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "invalid_mask  = np.isnan(u.data)\n",
    "invalid_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "u \n",
      " [[1.0 1.0 1.0 1.0 1.0]\n",
      " [1.0 nan 1.0 1.0 1.0]\n",
      " [1.0 1.0 1.0 1.0 1.0]\n",
      " [1.0 1.0 1.0 nan --]\n",
      " [1.0 1.0 1.0 -- --]]\n",
      "uf \n",
      " [[1.0 1.0 1.0 1.0 1.0]\n",
      " [1.0 1.0 1.0 1.0 1.0]\n",
      " [1.0 1.0 1.0 1.0 1.0]\n",
      " [1.0 1.0 1.0 -- --]\n",
      " [1.0 1.0 1.0 -- --]]\n"
     ]
    }
   ],
   "source": [
    "uf,_,_ = replace_outliers(u,u,invalid_mask,grid_mask)\n",
    "print(f'u \\n {u}')\n",
    "print(f'uf \\n',uf)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.12 ('echopiv')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "f83b0c3a4910470a1212112b1707d582432916ed4ba8aec962241a050aa18fae"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}