NAAI_Nemo_Adjacency_AI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d50f6b3e-56b2-4e01-be4b-0f2cae9fcb19",
   "metadata": {},
   "source": [
    "# _**NAAI!**_ — Nemo Adjacency AI !"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ede1351e-268b-48fd-b5a1-964206404985",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import tqdm\n",
    "\n",
    "# 🚀 AI!\n",
    "from sklearn.linear_model import LinearRegression  \n",
    "# 🚀 AI!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "84da5a0d-3c80-4607-8c66-4d9c6d8e8f50",
   "metadata": {},
   "outputs": [],
   "source": [
    "mm_path_v = \"/home/jovyan/shared_materials/model_data/mask/VIKING20X.L46-KKG36107B/mesh_mask.nc\"\n",
    "mm_path_i = \"/home/jovyan/shared_materials/model_data/mask/INALT20.L46-KFS104/mesh_mask.nc\"\n",
    "lon_lat_prec_degrees = 1e-5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5a51b7b0-bfad-4cf4-83e0-3332f07dae94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_adjacent_points_north(mesh_mask_path=None, lon_lat_precision_degrees=None):\n",
    "\n",
    "    # load mesh mask\n",
    "    ds_mm = xr.open_dataset(mesh_mask_path)\n",
    "    ds_mm = ds_mm.squeeze(drop=True)\n",
    "    ds_mm = ds_mm.assign_coords(\n",
    "        x=np.arange(ds_mm.sizes[\"x\"]),\n",
    "        y=np.arange(ds_mm.sizes[\"y\"]),\n",
    "    )\n",
    "    \n",
    "    # extract T-point lon and lat\n",
    "    lon, lat = ds_mm.glamt, ds_mm.gphit\n",
    "\n",
    "\n",
    "    # Cast lon lat to int\n",
    "    ilon = (lon / lon_lat_prec_degrees).astype(int)\n",
    "    ilat = (lat / lon_lat_prec_degrees).astype(int)\n",
    "\n",
    "    # Find out which row corresponds to the last one y\n",
    "    if (\n",
    "        abs(np.sort(ilon.isel(y=-1)) - np.sort(ilon.isel(y=-2))).mean()\n",
    "        < abs(np.sort(ilon.isel(y=-1)) - np.sort(ilon.isel(y=-3))).mean()\n",
    "    ):\n",
    "        corresponds_to_redundant = -2\n",
    "    else:\n",
    "        corresponds_to_redundant = -3\n",
    "    print(\"corresponding row is at y = \", corresponds_to_redundant)\n",
    "\n",
    "    # extract redundant (last row) and correspoding row coords\n",
    "    ilon_redundant = ilon.isel(x=slice(1, -1)).isel(y=-1, drop=True)\n",
    "    ilon_corresponds = ilon.isel(x=slice(1, -1)).isel(y=corresponds_to_redundant, drop=True)\n",
    "    ilat_redundant = ilon.isel(x=slice(1, -1)).isel(y=-1, drop=True)\n",
    "    ilat_corresponds = ilon.isel(x=slice(1, -1)).isel(y=corresponds_to_redundant, drop=True)\n",
    "\n",
    "    # find corresponding x\n",
    "    x_corr = []\n",
    "    for x_r, lon_r, lat_r in tqdm.tqdm(list(zip(ilon_redundant.x.data, ilon_redundant.data, ilat_redundant.data))):\n",
    "        for x_c, lon_c, lat_c in zip(ilon_corresponds.x.data, ilon_corresponds.data, ilat_corresponds.data):\n",
    "            if lon_r.data[()] == lon_c.data[()]:\n",
    "                if lat_r.data[()] == lat_c.data[()]:\n",
    "                    if x_c not in x_corr:\n",
    "                        x_corr.append(x_c)\n",
    "                        break\n",
    "\n",
    "    # filter adjacent x for outliers\n",
    "    adjacent_x = pd.Series(x_corr, name=\"adjacent_x\", index=pd.Series(ilon_redundant.x.data, name=\"reference_x\"))\n",
    "    adjacent_x_sanitized = adjacent_x.where(abs(adjacent_x.diff()) < 1.5 * abs(adjacent_x.diff()).mean()).dropna()\n",
    "\n",
    "    # fit clean adjacent indices\n",
    "    lr = LinearRegression()\n",
    "    lr.fit(np.array(adjacent_x_sanitized.index).reshape(-1, 1), np.array(adjacent_x_sanitized).reshape(-1, 1))\n",
    "    adjacent_x_fit = pd.Series(lr.predict(np.array(adjacent_x.index).reshape(-1, 1)).reshape(-1), index=adjacent_x.index, name=\"adjacent_x\").round().astype(int)\n",
    "\n",
    "    # test if we were successful\n",
    "    np.testing.assert_array_almost_equal(\n",
    "        ds_mm.glamt.isel(y=corresponds_to_redundant).sel(x=adjacent_x_fit.to_xarray()),\n",
    "        ds_mm.glamt.isel(y=-1).sel(x=adjacent_x_fit.index.to_series().to_xarray()),\n",
    "    )\n",
    "\n",
    "    # return adjacent indices\n",
    "    return adjacent_x_fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3fa8ee70-d43d-4fb3-97dc-0b9b7c0504a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "corresponding row is at y =  -3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1440/1440 [00:00<00:00, 3348.92it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "reference_x\n",
       "1       1441\n",
       "2       1440\n",
       "3       1439\n",
       "4       1438\n",
       "5       1437\n",
       "        ... \n",
       "1436       6\n",
       "1437       5\n",
       "1438       4\n",
       "1439       3\n",
       "1440       2\n",
       "Name: adjacent_x, Length: 1440, dtype: int64"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "find_adjacent_points_north(mesh_mask_path=mm_path_i, lon_lat_precision_degrees=lon_lat_prec_degrees)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "726629c7-7810-4361-bc0a-4dd34eae51f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "corresponding row is at y =  -3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1440/1440 [00:00<00:00, 3212.13it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "reference_x\n",
       "1       1441\n",
       "2       1440\n",
       "3       1439\n",
       "4       1438\n",
       "5       1437\n",
       "        ... \n",
       "1436       6\n",
       "1437       5\n",
       "1438       4\n",
       "1439       3\n",
       "1440       2\n",
       "Name: adjacent_x, Length: 1440, dtype: int64"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "find_adjacent_points_north(mesh_mask_path=mm_path_v, lon_lat_precision_degrees=lon_lat_prec_degrees)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "parcels",
   "language": "python",
   "name": "parcels"
  },
  "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}