{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import time as timing # renamed to not conflict with other variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hello Class\n",
    "#---------------------------------------------------------\n",
    "# Calculate and plot w vs t, z vs t, and w vs z given B(z)\n",
    "\n",
    "# Parameters to change: dt, alt_buoy\n",
    "\n",
    "dt = 0.1 # time step (s)  converged (true solution)\n",
    "#dt = 10 # time step (s)  Use dt from 10 to 0.001 s \n",
    "alt_buoy = False\n",
    "do_heun = False\n",
    "\n",
    "# alt_buoy = True changes the buoyancy profile for the parcel below its EL\n",
    "# after it first reaches its maximum height. The buoyancy profile is \n",
    "# similar to that for a parcel undergoing dry adiabatic descent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----Do not change any parameters in this cell----\n",
    "\n",
    "# Use Euler Forward scheme and a quadratic B(z) function,\n",
    "# B(z) = a * z^2 + b * z + c  (J/kg, z in meters)\n",
    "a = -2.3438e-08\n",
    "b = 2.3437e-04\n",
    "c = -0.2109\n",
    "\n",
    "q = 9000 # EL (m) (needed for alt_buoy = T)\n",
    "\n",
    "# Stopping conditions\n",
    "zmax = 20000 # maximum height allowed (m)\n",
    "zmin = -1    # minimum height allowed (m)\n",
    "tmax = 2000  # maximum time allowed (s)\n",
    "\n",
    "# Initial values\n",
    "t = 0           # time (s)\n",
    "w0 = 14.2548    # m/s    14.2548 just barely reaches LFC\n",
    "# w0 = 14\n",
    "w0 = 14.3\n",
    "z0 = 0          # m\n",
    "\n",
    "# Initialize Arrays\n",
    "w = []\n",
    "z = []\n",
    "time = []\n",
    "\n",
    "i = 0\n",
    "irrev = True # Irreversible condition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "While loop took: ~0.027 Seconds\n"
     ]
    }
   ],
   "source": [
    "# Here is the main computation\n",
    "\n",
    "start_time = timing.time()\n",
    "while (z0 < zmax) and (z0 > zmin) and (t < tmax):\n",
    "    \n",
    "    i += 1 #increment time step counter\n",
    "\n",
    "    # Save into arrays for plotting\n",
    "    w.append(w0)\n",
    "    z.append(z0)\n",
    "    time.append(t)\n",
    "\n",
    "    # Calculate buoyancy at current height, z0\n",
    "    B = (a * z0**2) + (b * z0) + c\n",
    "    if alt_buoy:\n",
    "        \n",
    "        # AKA if reversible\n",
    "        if not irrev:\n",
    "            if (B > 0) and (w0 < 0):\n",
    "                irrev = True\n",
    "\n",
    "        # Irreversible buoyancy profile differs below EL = q:\n",
    "        # B = dB/dz(EL) * (z - q) = (2 * a * q + b) * (z - q)\n",
    "        else:\n",
    "            if z0 < q:\n",
    "                B = ((2 * a * q) + b) * (z0 - q)  # Larger only below EL\n",
    "\n",
    "    # update w and z\n",
    "    w1 = w0 + (B * dt)  # dw/dt = B\n",
    "    z1 = z0 + (w0 * dt) # dz/dt = w\n",
    "\n",
    "    if do_heun:\n",
    "        zstar = z1\n",
    "        wstar = w1\n",
    "        # Calculate buoyancy at current height\n",
    "        Bstar = (a * zstar**2) + (b * zstar) + c\n",
    "        if alt_buoy:\n",
    "            if zstar < q:\n",
    "                Bstar = (2 * a * q + b) * (zstar - q)\n",
    "        \n",
    "        # Heun Scheme\n",
    "        w1 = w0 + 0.5 * (B + Bstar) * dt\n",
    "        z1 = z0 + 0.5 * (w0 + wstar) * dt\n",
    "\n",
    "    # Update time\n",
    "    t = t + dt\n",
    "\n",
    "    # copy w1 and w0 over for use in next iteration\n",
    "    w0 = w1\n",
    "    z0 = z1\n",
    "\n",
    "end_time = timing.time()\n",
    "\n",
    "# f-strings are amazing, look them up\n",
    "print(f'While loop took: ~{end_time-start_time:.3f} Seconds')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting - Feel free to change the plotting cells below\n",
    "\n",
    "# For convenience, let's convert python lists to numpy ~~arrays~~\n",
    "w = np.array(w)\n",
    "z = np.array(z)\n",
    "time = np.array(time)\n",
    "\n",
    "# Now we can operate on our arrays with math equations\n",
    "# We couldn't have done that with lists\n",
    "B_plot = (a * z**2) + (b * z) + c # Buoyancy acceleration (m/s^2)\n",
    "Bfac = 28 # Converts buoyancy acceleration to temperature difference (K)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now plot graphs of the variables using the plt.plot() function\n",
    "# See https://matplotlib.org/stable/tutorials/pyplot.html for a good tutorial"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}