{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2D Parcels Test Model (depth and lon) - Matt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load packages and functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from parcels import Field, FieldSet, ParticleSet, Variable, JITParticle, ErrorCode\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import rc, animation\n",
    "import xarray as xr\n",
    "from datetime import datetime\n",
    "import os\n",
    "\n",
    "from IPython.display import Image\n",
    "rc('animation', html='html5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build parcels particle simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Define parcels FieldSet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define SalishSeaCast results filepath\n",
    "def make_prefix(date, res='h'):\n",
    "    \"\"\"Construct path prefix for local SalishSeaCast results given date object and paths dict\n",
    "    e.g., /results2/SalishSea/nowcast-green.201905/daymonthyear/SalishSea_1h_yyyymmdd_yyyymmdd\n",
    "    \"\"\"\n",
    "    path = '/results2/SalishSea/nowcast-green.202111/'\n",
    "    datestr = '_'.join(np.repeat(date.strftime('%Y%m%d'), 2))\n",
    "    folder = date.strftime(\"%d%b%y\").lower()\n",
    "    prefix = os.path.join(path, f'{folder}/SalishSea_1{res}_{datestr}')\n",
    "    \n",
    "    return prefix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pick hour, day and location for simulation\n",
    "hour = 0 # Choose start hour\n",
    "path_NEMO = make_prefix(datetime(2018, 8, 1)) # Choose start year, month, day\n",
    "lati = 49.224563 # Choose coords\n",
    "loni = -123.605357\n",
    "jjii = xr.open_dataset('/home/sallen/MEOPAR/grid/grid_from_lat_lon_mask999.nc') # This is a file that gives the closest SalishSeaCast grid point (i and j) for a given lat and lon \n",
    "j = [jjii.jj.sel(lats=lati, lons=loni, method='nearest').item()][0] # j is the index along the SalishSeaCast grid (think y axis of model domain, or similar to latitude), while i is the index across the grid\n",
    "i = [jjii.ii.sel(lats=lati, lons=loni, method='nearest').item()][0]\n",
    "\n",
    "# Specify some simulation parameters\n",
    "T = 6 * 86400 #s (run time) (86400 s = 24 hours/1 day) - initially 5e3 * 100 (this is about 5.7 days)\n",
    "dt = 5 #s (timestep)\n",
    "N = 10e3 #number of particles\n",
    "outputdt = 3600 #s (3600 = 1 hr) (how often do you want output?) initially 500\n",
    "outputpath = '/ocean/mattmiller/MOAD/results/parcels/test/Outputmix_sink_and_swim_6_days.zarr' # this makes sure the results are saved in my /results folder located beside /analysis-matt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'depth' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Define domain, velocity fields and Kz \u001b[39;00m\n\u001b[1;32m      2\u001b[0m dim \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m100\u001b[39m\n\u001b[0;32m----> 3\u001b[0m dep \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(\u001b[43mdepth\u001b[49m)\n\u001b[1;32m      4\u001b[0m lon \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mlinspace(\u001b[38;5;241m0.\u001b[39m, \u001b[38;5;241m2e3\u001b[39m, dim, dtype\u001b[38;5;241m=\u001b[39mnp\u001b[38;5;241m.\u001b[39mfloat32)\n\u001b[1;32m      6\u001b[0m \u001b[38;5;66;03m# Define variables for parcels FieldSet\u001b[39;00m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'depth' is not defined"
     ]
    }
   ],
   "source": [
    "# Define domain, velocity fields and Kz \n",
    "dim = 100\n",
    "dep = len(depth)\n",
    "lon = np.linspace(0., 2e3, dim, dtype=np.float32)\n",
    "\n",
    "# Define variables for parcels FieldSet\n",
    "U = Field('U', np.zeros((dep, dim), dtype=np.float32), lon=lon, depth=depth) # in parcels, 'U' represents the zonal flow velocity (zonal = east-to-west/west-to-east)\n",
    "V = Field('V', np.zeros((dep, dim), dtype=np.float32), lon=lon, depth=depth) # in parcels, 'V' represents the meridional flow velocity (meridional = north-to-south/south-to-north)\n",
    "Kz_data = np.zeros((dep, dim), dtype=np.float32)\n",
    "for i in range(dim):\n",
    "    Kz_data[:,i]=Kz_col   \n",
    "Kz = Field('Kz', Kz_data, grid=U.grid)\n",
    "\n",
    "# Build parcels FieldSet\n",
    "fieldset = FieldSet(U,V)\n",
    "fieldset.add_field(Kz)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Define parcels ParticleSet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "additional Variables can be added to the particles (e.g. temperature, to keep track of the temperature that particles experience)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a new particleclass with Variable 'age' with initial value 0.\n",
    "AgeParticle = parcels.JITParticle.add_variable(parcels.Variable(\"age\", initial=0))\n",
    "\n",
    "pset = parcels.ParticleSet(\n",
    "    fieldset=fieldset,  # the fields that the particleset uses\n",
    "    pclass=AgeParticle,  # define the type of particle\n",
    "    lon=29,  # release longitude\n",
    "    lat=-33,  # release latitude\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Define parcels particle kernels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Create a custom kernel which displaces each particle southward\n",
    "def NorthVel(particle, fieldset, time):\n",
    "    if time > 10 * 86400 and time < 10.2 * 86400:\n",
    "        vvel = -1e-4\n",
    "        particle_dlat += vvel * particle.dt\n",
    "\n",
    "\n",
    "# Create a custom kernel which keeps track of the particle age (minutes)\n",
    "def Age(particle, fieldset, time):\n",
    "    particle.age += particle.dt / 3600\n",
    "\n",
    "\n",
    "# define all kernels to be executed on particles using an (ordered) list\n",
    "kernels = [Age, NorthVel, parcels.AdvectionRK4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Execute parcels simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_file = pset.ParticleFile(\n",
    "    name = \"Outputmix_sink_and_swim.zarr\",  # the name of the output file\n",
    "    outputdt = 3600,  # the time period between consecutive output steps\n",
    "    chunks = (1, 10),  # the chunking of the output file (number of particles, timesteps)\n",
    ")\n",
    "\n",
    "pset.execute(\n",
    "    kernels,  # the list of kernels (which defines how particles move)\n",
    "    runtime = 86400 * 24,  # the total length of the run in seconds\n",
    "    dt = 300,  # the timestep of the kernel in seconds\n",
    "    output_file = output_file,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load results from previous parcels simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dat = xr.load_dataset('/ocean/mattmiller/MOAD/results/parcels/test/Outputmix_sink_and_swim.zarr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View data\n",
    "print(dat)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "analysis-matt",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}