{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "83ddc01a-271a-4cc0-af8a-980496592fcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate test input for perturbing solar flux applied in POP\n",
    "\n",
    "# import libraries\n",
    "import numpy as np \n",
    "from netCDF4 import Dataset\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Polygon\n",
    "from mpl_toolkits.basemap import Basemap, shiftgrid\n",
    "\n",
    "# define anomaly and domain where flux is pertubed \n",
    "flux_anomaly = 30 \n",
    "lonRange=[-40,0]\n",
    "latRange=[0,30] \n",
    "\n",
    "# read coordinates and land mask \n",
    "nc = Dataset('grid_gx1.nc','r')\n",
    "lon = nc['plon'][:]\n",
    "lat = nc['plat'][:]\n",
    "msk = nc['pmask'][:]\n",
    "nc.close()\n",
    "lon1 = lon.flatten()\n",
    "lat1 = lat.flatten()\n",
    "msk1 = msk.flatten()\n",
    "for i in range(len(msk1)): \n",
    "    if not (msk1[i] == 1 and lon1[i] >= lonRange[0] and lon1[i] <= lonRange[1] and lat1[i] >= latRange[0] and lat1[i] <= latRange[1]): msk1[i] = 0 \n",
    "\n",
    "# generate test input for Jan 2023 \n",
    "for day in range(1,32): \n",
    "    file_name = f'popforcing_2023-01-{day:0>2d}-00000.nc' \n",
    "    nc = Dataset(file_name, 'w', format='NETCDF4_CLASSIC')\n",
    "    nc.createDimension('time',0)\n",
    "    nc.createDimension('x',msk.shape[-1])\n",
    "    nc.createDimension('y',msk.shape[-2])\n",
    "    nc.createVariable('SHF_QSW','f4',['time','y','x'])[0,:,:] = msk1.reshape(msk.shape)*flux_anomaly \n",
    "    nc.close()\n",
    "\n",
    "# plot grid coordinates for domain\n",
    "fig, ax = plt.subplots(figsize=(8, 8))    \n",
    "m = Basemap(projection='ortho',lon_0=-10,lat_0=20,resolution='c')\n",
    "m.fillcontinents(color='gray',lake_color='gray')\n",
    "x, y = m(lon1,lat1)\n",
    "xvec = []\n",
    "yvec = []\n",
    "for i in range(len(msk1)):\n",
    "    if msk1[i] == 1: \n",
    "        xvec.append(x[i])\n",
    "        yvec.append(y[i])\n",
    "        \n",
    "m.scatter(xvec,yvec,0.1,marker='o',color='r')\n",
    "plt.savefig(f'grid_gx1.png',format='png',dpi=600,bbox_inches='tight')\n",
    "plt.close()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5b688b9b-3574-470a-a7f2-a9ed80a5a56c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-39.43749108 -38.31249105 -37.18749102 ... -42.81249117 -41.68749114\n",
      "  -40.56249111]\n",
      " [-39.43749108 -38.31249105 -37.18749102 ... -42.81249117 -41.68749114\n",
      "  -40.56249111]\n",
      " [-39.43749108 -38.31249105 -37.18749102 ... -42.81249117 -41.68749114\n",
      "  -40.56249111]\n",
      " ...\n",
      " [-39.74866914 -39.24619887 -38.74422675 ... -41.25575544 -40.75378332\n",
      "  -40.25130857]\n",
      " [-39.76540523 -39.29641051 -38.82792558 ... -41.17205661 -40.70357168\n",
      "  -40.23457279]\n",
      " [-39.78349101 -39.35066969 -38.91836527 ... -41.08161692 -40.6493125\n",
      "  -40.21648733]]\n"
     ]
    }
   ],
   "source": [
    "print(lon)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab2475ad-49ea-4b70-8f12-2168bbca99b3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}