01_calc.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "82c9b707",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp calc\n",
    "from nbdev import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3e27a19",
   "metadata": {},
   "source": [
    "# TEF calculations\n",
    "\n",
    "> contains calculations of TEF framework"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "d9db8b3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from tqdm import tqdm\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "a6e7aa9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def convert_q_to_Q(var_q, q, var_q2 = None):\n",
    "    \"\"\"Converts transport per coordinate class `q` to the integrated transport `Q` with the respective coordinates. \n",
    "        Use if q is already computed separately.\"\"\"\n",
    "    if len(q.shape) == 1:\n",
    "        #no time axis\n",
    "        delta_var = var_q[1]-var_q[0]\n",
    "        out_Q = np.cumsum(q[::-1])[::-1]*delta_var\n",
    "        out_Q = np.append(out_Q, np.zeros(1,), axis=0)\n",
    "        var_Q = np.append(var_q-0.5*delta_var,\n",
    "                          var_q[-1] + 0.5*delta_var)\n",
    "\n",
    "        out = xr.Dataset({\n",
    "        \"Q\": ([\"var_Q\"], out_Q)},\n",
    "        coords={\n",
    "            \"var_Q\": ([\"var_Q\"], var_Q),\n",
    "        })\n",
    "\n",
    "        return out\n",
    "    \n",
    "    elif len(q.shape) == 2 and var_q2 is None:\n",
    "        #time axis existing, 1D TEF\n",
    "        delta_var = var_q[1]-var_q[0]\n",
    "        T = q.shape[0]\n",
    "        out_Q = np.cumsum(q[:,::-1], axis=1)[:,::-1]*delta_var\n",
    "        out_Q = np.append(out_Q,np.zeros((T,1)),axis=1)\n",
    "        var_Q = np.append(var_q-0.5*delta_var,var_q[-1]+0.5*delta_var)\n",
    "        \n",
    "        out = xr.Dataset({\n",
    "        \"Q\": ([\"time\", \"var_Q\"], out_Q)},\n",
    "        coords={\n",
    "            \"time\": ([\"time\"], _get_time_array(q)),\n",
    "            \"var_Q\": ([\"var_Q\"], var_Q),\n",
    "        })\n",
    "\n",
    "        return out\n",
    "\n",
    "    elif len(q.shape) == 2 and var_q2 is not None:\n",
    "        #no time axis existing, 2D TEF\n",
    "        N_1 = q.shape[1]\n",
    "        N_2 = q.shape[1]\n",
    "        delta_var = var_q[1]-var_q[0]\n",
    "        delta_var2 = var_q2[1]-var_q2[0]\n",
    "        out_Q = np.zeros((N_1+1, N_2+1))\n",
    "        out_Q_tmp = np.cumsum(np.cumsum(q[::-1,::-1],axis=0),axis=1)[::-1,::-1]*delta_var2*delta_var\n",
    "        out_Q[:-1,:-1] = out_Q_tmp\n",
    "        var_Q = np.append(var_q-0.5*delta_var,var_q[-1]+0.5*delta_var)\n",
    "        var_Q2 = np.append(var_q2-0.5*delta_var2,var_q2[-1]+0.5*delta_var2)\n",
    "        out = xr.Dataset({\n",
    "        \"Q2\": ([\"var_Q\", \"var_Q2\"], out_Q)},\n",
    "        coords={\n",
    "            \"var_Q\": ([\"var_Q\"], var_Q),\n",
    "            \"var_Q2\": ([\"var_Q2\"], var_Q),\n",
    "        })\n",
    "\n",
    "        return out \n",
    "\n",
    "    elif len(q.shape) == 3 and var_q2 is not None:\n",
    "        #time axis and 2D TEF\n",
    "        T = q.shape[0]\n",
    "        N_1 = q.shape[1]\n",
    "        N_2 = q.shape[1]\n",
    "        delta_var = var_q[1]-var_q[0]\n",
    "        delta_var2 = var_q2[1]-var_q2[0]\n",
    "        out_Q = np.zeros((T, N_1+1, N_2+1))\n",
    "        out_Q_tmp = np.cumsum(np.cumsum(q[:,::-1,::-1],\n",
    "                                        axis=1),\n",
    "                                        axis=2)[:,::-1,::-1]*delta_var2*delta_var\n",
    "        out_Q[:, :-1, :-1] = out_Q_tmp\n",
    "        var_Q = np.append(var_q-0.5*delta_var,\n",
    "                          var_q[-1]+0.5*delta_var)\n",
    "        var_Q2 = np.append(var_q2-0.5*delta_var2,\n",
    "                           var_q2[-1]+0.5*delta_var2)\n",
    "\n",
    "        out = xr.Dataset({\n",
    "        \"Q2\": ([\"time\", \"var_Q\", \"var_Q2\"], out_Q)},\n",
    "        coords={\n",
    "            \"time\": ([\"time\"], _get_time_array(q)),\n",
    "            \"var_Q\": ([\"var_Q\"], var_Q),\n",
    "            \"var_Q2\": ([\"var_Q2\"], var_Q),\n",
    "        })\n",
    "\n",
    "        return out     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "091fb881",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<h4 id=\"convert_q_to_Q\" class=\"doc_header\"><code>convert_q_to_Q</code><a href=\"__main__.py#L2\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
       "\n",
       "> <code>convert_q_to_Q</code>(**`var_q`**, **`q`**, **`var_q2`**=*`None`*)\n",
       "\n",
       "Converts transport per coordinate class `q` to the integrated transport `Q` with the respective coordinates. \n",
       "Use if q is already computed separately."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_doc(convert_q_to_Q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73b944e7",
   "metadata": {},
   "source": [
    "Converts the  transport per coordinate class `q` to the integrated transport `Q` with the respective coordinates.  \n",
    "`var_q` - coordinate bin centers of `q`(e.g. salinity)\n",
    "\n",
    "`q`     - transport per coordinates class (e.g. transport per salinity class)\n",
    "\n",
    "`var_q2`- for two dimensional TEF cases this corresponds to the second coordinate the transport is sorted by (e.g. temperature)\n",
    "\n",
    "**Returns**:\n",
    "\n",
    "New xarray dataset with `Q`, `var_Q`, and `var_Q2` (if two dimensional TEF)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "85354205",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def sort_1dim(constructorTEF,\n",
    "              N = 1024,\n",
    "              minmaxrange = None):\n",
    "    \"\"\"Performs coordinate transformation by given variable.\"\"\"\n",
    "    if constructorTEF.tracer is None:\n",
    "        raise ValueError(\"Please define a variable that you want to sort by.\")\n",
    "    if constructorTEF.transport is None:\n",
    "        raise ValueError(\"Please provide transport term.\")\n",
    "\n",
    "    if minmaxrange is None:\n",
    "        varmin = np.floor(constructorTEF.tracer.min().values)\n",
    "        varmax = np.ceil(constructorTEF.tracer.max().values)\n",
    "    else:\n",
    "        if minmaxrange[0] > constructorTEF.tracer.min().values:\n",
    "            print(\"Warning: Given minimum value is greater than the minimum value of the variable.\")\n",
    "            print(\"Warning: Given {}, minmum value of variable {}\".format(minmaxrange[0],\n",
    "                                                                 constructorTEF.tracer.min().values))\n",
    "        if minmaxrange[-1] < constructorTEF.tracer.max().values:\n",
    "            print(\"Warning: Given maximum value is smaller than the maximum value of the variable.\")\n",
    "            print(\"Warning: Given {}, maximum value of variable {}\".format(minmaxrange[-1],\n",
    "                                                                 constructorTEF.tracer.max().values))\n",
    "        if type(minmaxrange) != \"numpy.ndarray\" and type(minmaxrange) is not tuple:\n",
    "            raise ValueError(\"Please provide array range, e.g. np.arange(0,10), or a tuple, e.g. (0,10).\")\n",
    "\n",
    "        else:\n",
    "            varmin = minmaxrange[0]\n",
    "            varmax = minmaxrange[-1]\n",
    "\n",
    "    if type(minmaxrange) == \"numpy.ndarray\":\n",
    "        print('Using provided numpy array')\n",
    "\n",
    "        var_q = minmaxrange.copy()\n",
    "\n",
    "        #check if equidistant\n",
    "        if np.unique(np.diff(var_q)) != 1:\n",
    "            print('Warning: Provided array is not equidistant, but the function assumes equidistance!')\n",
    "        delta_var = var_q[1]-var_q[0]\n",
    "\n",
    "        var_Q = np.arange(var_q-0.5*delta_var,\n",
    "                          var_q[-1]+0.5*delta_var,\n",
    "                          delta_var)\n",
    "    else:\n",
    "        print('Constructing var_q and var_Q')\n",
    "        delta_var = ((varmax-varmin)/N)\n",
    "\n",
    "        var_q = np.linspace(varmin + 0.5*delta_var,\n",
    "                            varmax - 0.5*delta_var,\n",
    "                            N)\n",
    "        var_Q = np.linspace(varmin, varmax, N+1)\n",
    "\n",
    "    # Changelog: 27.05.2021: Change var_Q to var_q\n",
    "    # compute the index idx that will be used for sorting\n",
    "    idx = xr.apply_ufunc(np.digitize, constructorTEF.tracer, var_q) \n",
    "\n",
    "    out_q = np.zeros((len(constructorTEF.ds.time), N))\n",
    "\n",
    "    for i in tqdm(range(N)):\n",
    "        #Sorting into bins\n",
    "        out_q[:, i] = constructorTEF.transport.where(idx == i).sum([\"depth\",\n",
    "                                                                    \"lat\",\n",
    "                                                                    \"lon\"],\n",
    "                                                     dtype=np.float64) / delta_var\n",
    "    \n",
    "    out_Q = np.append(np.cumsum(out_q[:,::-1],\n",
    "                                axis=1)[:,::-1],\n",
    "                      np.zeros((len(constructorTEF.ds.time), 1)),axis=1)*delta_var\n",
    "    \n",
    "    out = xr.Dataset({\n",
    "    \"q\": ([\"time\", \"var_q\"], out_q),\n",
    "    \"Q\": ([\"time\", \"var_Q\"], out_Q)},\n",
    "    coords={\n",
    "        \"time\": ([\"time\"], constructorTEF.ds[\"time\"].data),\n",
    "        \"var_q\": ([\"var_q\"],var_q),\n",
    "        \"var_Q\": ([\"var_Q\"], var_Q),\n",
    "    })\n",
    "    \n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "f32b9b10",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<h4 id=\"sort_1dim\" class=\"doc_header\"><code>sort_1dim</code><a href=\"__main__.py#L2\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
       "\n",
       "> <code>sort_1dim</code>(**`constructorTEF`**, **`N`**=*`1024`*, **`minmaxrange`**=*`None`*)\n",
       "\n",
       "Performs coordinate transformation by given variable."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_doc(sort_1dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e81f22e4",
   "metadata": {},
   "source": [
    "`sort_by_variable` - Sorts by given variable array (e.g. salinity)\n",
    "\n",
    "`transport` - transport term; e.g. for volume the unit is $m^3 / s$\n",
    "\n",
    "`N` - Number of bins that is used by the sorting algorithm (default is 1024)\n",
    "\n",
    "`minmaxrange` - Defines the range of the bin axis that is considered by the sorting algorithm (type can be either a list, a tuple or a 1D-numpy array). If a 1D-numpy array is provided the coordinates will correspond to `var_q`. If `None` the minimum and maximum values are automatically defined by the minimum and maximum values of `sort_by_variable`.\n",
    "\n",
    "**Returns**:\n",
    "\n",
    "xarray dataset with computed `q` and `Q` terms and the their coordinates `var_q` and `var_Q`.\n",
    "\n",
    "### A short note on the binning for the terms `q` and `Q`\n",
    "\n",
    "`q` is defined on the axis that the transport is sorted by (center values of bins `var_q`). `Q` denotes the integral of `q` above a certain threshold and is therefore placed at the interface between to binds of `var_q`. These interface values are provided by `var_Q`.\n",
    "\n",
    "In the figure of [Lorenz et al (2019), their Figure 5](https://os.copernicus.org/articles/15/601/2019/) this is explained and illustrated as follows:\n",
    "\n",
    "Sketch of how `Q` (here denoted as $Q^c_n$) and `q` (here denoted as $q^c_n$) are located in a discrete salinity space. The salinity interval $[S_{1/2},S_{N+1/2}]$ is divided into N equidistant salinity classes of length $\\delta S$. The entries of $Q^c$ ($Q^c_n$) are located on the lines (bin interfaces), and the entries of $q^c$ ($q^c_n$) are located on the dots (bin centers).\n",
    "\n",
    "<img src=\"https://os.copernicus.org/articles/15/601/2019/os-15-601-2019-f05-web.png\" alt=\"drawing\" width=\"200\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "e595edc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def sort_2dim(constructorTEF,      \n",
    "              N = (1024, 1024),\n",
    "              minmaxrange = None,\n",
    "              minmaxrange2 = None):\n",
    "        \"\"\"Sort transport by two given variables\"\"\"\n",
    "        if constructorTEF.tracer[0] is None:\n",
    "            raise ValueError(\"Please define a variable that you want to sort by.\")\n",
    "        if constructorTEF.tracer[1] is None:\n",
    "            raise ValueError(\"Please define a second variable that you want to sort by.\")    \n",
    "\n",
    "        if constructorTEF.transport is None:\n",
    "            raise ValueError(\"Please provided transport term.\")\n",
    "\n",
    "        if minmaxrange is None:\n",
    "            varmin = np.floor(constructorTEF.tracer[0].min().values)\n",
    "            varmax = np.ceil(constructorTEF.tracer[0].max().values)      \n",
    "        else:\n",
    "            if minmaxrange[0] > constructorTEF.tracer[0].min().values:\n",
    "                print(\"Warning: Given minimum value is gretaer than the minimum value of the variable.\")\n",
    "                print(\"Warning: Given {}, minmum value of variable {}\".format(minmaxrange[0],\n",
    "                                                                     constructorTEF.tracer[0].min().values))\n",
    "            if minmaxrange[-1] < constructorTEF.tracer[0].max().values:\n",
    "                print(\"Warning: Given maximum value is smaller than the maximum value of the variable.\")\n",
    "                print(\"Warning: Given {}, maximum value of variable {}\".format(minmaxrange[-1],\n",
    "                                                                     constructorTEF.tracer[0].max().values))\n",
    "            if type(minmaxrange) != \"numpy.ndarray\" and type(minmaxrange) is not tuple:\n",
    "                raise ValueError(\"Please provide array range, e.g. np.arange(0,10), or a tuple, e.g. (0,10).\")\n",
    "            else: \n",
    "                print(\"minmaxrange is a tuple\")\n",
    "                varmin = minmaxrange[0]\n",
    "                varmax = minmaxrange[-1]\n",
    "\n",
    "        if minmaxrange2 is None:\n",
    "            varmin2 = np.floor(constructorTEF.tracer[1].min().values)\n",
    "            varmax2 = np.ceil(constructorTEF.tracer[1].max().values)\n",
    "        else:\n",
    "            if minmaxrange2[0] > constructorTEF.tracer[1].min().values:\n",
    "                print(\"Warning: Given minimum value is greater than the minimum value of the variable.\")\n",
    "                print(\"Warning: Given {}, minmum value of variable {}\".format(minmaxrange2[0],\n",
    "                                                                     constructorTEF.tracer[1].min().values))\n",
    "            if minmaxrange2[-1] < constructorTEF.tracer[1].max().values:\n",
    "                print(\"Warning: Given maximum value is smaller than the maximum value of the variable.\")\n",
    "                print(\"Warning: Given {}, maximum value of variable {}\".format(minmaxrange2[-1],\n",
    "                                                                     constructorTEF.tracer[1].max().values))\n",
    "            if type(minmaxrange2) != \"numpy.ndarray\" and type(minmaxrange2) is not tuple:\n",
    "                raise ValueError(\"Please provide array range, e.g. np.arange(0,10), or a tuple, e.g. (0,10).\")\n",
    "            else:    \n",
    "                print(\"minmaxrange2 is a tuple\")\n",
    "                varmin2 = minmaxrange2[0]\n",
    "                varmax2 = minmaxrange2[-1]\n",
    "     \n",
    "        if type(N) is tuple:\n",
    "            N1 = N[0]\n",
    "            N2 = N[1]\n",
    "\n",
    "        if type(minmaxrange) == \"numpy.ndarray\":\n",
    "            print('Using provided numpy array for variable 1')\n",
    "\n",
    "            var_q = minmaxrange\n",
    "\n",
    "            #check if equidistant\n",
    "            diff=np.diff(np.diff(var_q))\n",
    "            if len(diff[diff!=0]) != 0:\n",
    "                print('Warning: Provided array for variable1 is not equidistant, but the function assumes equidistance!')\n",
    "            delta_var=var_q[1]-var_q[0]\n",
    "\n",
    "            var_Q = np.arange(var_q-0.5*delta_var, var_q[-1]+0.5*delta_var, delta_var)\n",
    "        else:\n",
    "            #constructing\n",
    "            delta_var = ((varmax-varmin)/N1)           \n",
    "            var_q = np.linspace(varmin + 0.5*delta_var,\n",
    "                                varmax - 0.5*delta_var,\n",
    "                                N1)                           \n",
    "            var_Q = np.linspace(varmin, varmax, N1+1)\n",
    "\n",
    "        if type(minmaxrange2) == \"numpy.ndarray\":\n",
    "            print('Using provided numpy array for variable 2')\n",
    "\n",
    "            var_q2 = minmaxrange2\n",
    "\n",
    "            #check if equidistant\n",
    "            diff=np.diff(np.diff(var_q2))\n",
    "            if len(diff[diff!=0]) != 0:\n",
    "                print('Warning: Provided array for variable2 is not equidistant, but the function assumes equidistance!')\n",
    "            delta_var2=var_q2[1]-var_q2[0]\n",
    "\n",
    "            var_Q2 = np.arange(var_q2-0.5*delta_var2, var_q2[-1]+0.5*delta_var2, delta_var2)\n",
    "        else:\n",
    "            #contructing\n",
    "            delta_var2 = ((varmax2-varmin2)/N2)\n",
    "            var_q2 = np.linspace(varmin2 + 0.5*delta_var2,\n",
    "                                 varmax2 - 0.5*delta_var2,\n",
    "                                 N2)\n",
    "            var_Q2 = np.linspace(varmin2, varmax2, N2+1)\n",
    "\n",
    "        #sortingt\n",
    "        idx = xr.apply_ufunc(np.digitize, constructorTEF.tracer[0], var_Q)\n",
    "        idy = xr.apply_ufunc(np.digitize, constructorTEF.tracer[1], var_Q2)\n",
    "\n",
    "        out_q = np.zeros((len(constructorTEF.ds.time),N1, N2))\n",
    "        \n",
    "        for i in tqdm(range(N1)):\n",
    "            for j in range(N2):\n",
    "                indices = (idx == i) & (idy == j)\n",
    "                if indices.any():\n",
    "                    out_q[:, i, j] = constructorTEF.transport.where(indices).sum([\"depth\", \"lat\", \"lon\"],\n",
    "                                                                                 dtype=np.float64) / delta_var / delta_var2\n",
    "                else:\n",
    "                    out_q[:, i, j] = np.NaN\n",
    "        \n",
    "        out_Q = np.zeros((len(constructorTEF.ds.time), N1+1, N2+1))\n",
    "        out_Q_tmp = np.cumsum(np.cumsum(out_q[:,::-1,::-1],axis=1),axis=2)[:,::-1,::-1]*delta_var2*delta_var\n",
    "        out_Q[:,:-1,:-1] = out_Q_tmp\n",
    "        \n",
    "        out = xr.Dataset({\n",
    "        \"q2\": ([\"time\", \"var_q\", \"var_q2\"], out_q),\n",
    "        \"Q2\": ([\"time\", \"var_Q\", \"var_Q2\"], out_Q)},\n",
    "        coords={\n",
    "            \"time\": ([\"time\"], constructorTEF.ds[\"time\"].data),\n",
    "            \"var_q\": ([\"var_q\"],var_q),\n",
    "            \"var_q2\": ([\"var_q2\"], var_q2),\n",
    "            \"var_Q\": ([\"var_Q\"], var_Q),\n",
    "            \"var_Q2\": ([\"var_Q2\"], var_Q2),\n",
    "        })\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "e9f5b066",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<h4 id=\"sort_2dim\" class=\"doc_header\"><code>sort_2dim</code><a href=\"__main__.py#L2\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
       "\n",
       "> <code>sort_2dim</code>(**`constructorTEF`**, **`N`**=*`(1024, 1024)`*, **`minmaxrange`**=*`None`*, **`minmaxrange2`**=*`None`*)\n",
       "\n",
       "Sort transport by two given variables"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_doc(sort_2dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c064bf5b",
   "metadata": {},
   "source": [
    "`sort_by_variable` - Sorts by given variable array (e.g. salinity); Note: if plotted, this will be the y-axis\n",
    "\n",
    "`sort_by_variable2` - Sorts by given second variable array (e.g. temperature); Note: if plotted, this will be the x-axis\n",
    "\n",
    "`transport` - transport term; for volume the unit is $m^3 / s$\n",
    "\n",
    "`N` - Number of bins that is used by the sorting algorithm. N is a tuple (N1, N2; default) corresponding to the sorting dimensions.\n",
    "\n",
    "`minmaxrange` - Defines the range that is considered by the sorting algorithm (type can be either list, tuple or 1D-numpy array). If a 1D-numpy array is provided the coordinates will correspond to `var_q`. If `None` the min and max values are automatically defined by the min and max values of `sort_by_variable`.\n",
    "\n",
    "`minmaxrange2` - Defines the range that is considered by the sorting algorithm for the second variable `sort_by_variable2` (type can be either list, tuple or 1D-numpy array). If a 1D-numpy array is provided the coordinates will correspond to `var_q`. If `None` the min and max values are automatically defined by the min and max values of `sort_by_variable2`.\n",
    "\n",
    "**Returns**:\n",
    "\n",
    "xarray dataset with computed `q2` and `Q2` terms and the their coordinates `var_q`, `var_q2` and `var_Q`, `var_Q2`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "ce8da27e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def calc_bulk_values(coord,\n",
    "                     Q,\n",
    "                     Qc=None,\n",
    "                     Q_thresh=None,\n",
    "                     index=None,\n",
    "                     ):\n",
    "    \"\"\"Calculate the bulk values from a provided Q profile.\n",
    "\n",
    "    This methods uses the dividing salinity approach proposed by MacCready\n",
    "    et al. (2018) and described/tested in detail by Lorenz et al. (2019).\n",
    "\n",
    "    If a tracer transport Qc is provided, also its bulk values are computed.\n",
    "    \"\"\"\n",
    "    coord_min=coord[0]\n",
    "    delta_var=coord[1]-coord[0]\n",
    "\n",
    "    if len(Q.shape) > 1:\n",
    "        \n",
    "        #first dimension is time! -> keep this dimension!\n",
    "        #prepare storage arrays for Qin, Qout, consider multiple inflow/outflows! \n",
    "\n",
    "        Qin_ar = np.zeros((Q.shape[0],10)) #10 is the dummy length\n",
    "        Qout_ar = np.zeros((Q.shape[0],10))\n",
    "        if Qc is not None:\n",
    "            Qc_in_ar = np.zeros((Q.shape[0], 10))\n",
    "            Qc_out_ar = np.zeros((Q.shape[0], 10))\n",
    "        divval_ar = np.zeros((Q.shape[0],11)) #if there are 10 transports there would be 11 dividing salinities\n",
    "        indices = np.zeros((Q.shape[0],11))\n",
    "\n",
    "        for t in tqdm(np.arange(Q.shape[0])):\n",
    "            if Q_thresh is None:\n",
    "                #set a default thresh\n",
    "                Q_thresh=0.01*np.max(np.abs(Q[t]))\n",
    "            if index is None:\n",
    "                ind,minmax = _find_extrema(Q[t],Q_thresh)\n",
    "            else:\n",
    "                ind=np.copy(index[t])\n",
    "                ind=ind[ind!=0]\n",
    "\n",
    "            div_val=[]\n",
    "            i=0\n",
    "            for i in range(len(ind)):\n",
    "                div_val.append(coord_min+delta_var*ind[i])\n",
    "                i+=1\n",
    "                #calculate transports etc.\n",
    "            Q_in_m=[]\n",
    "            Q_out_m=[]\n",
    "            if Qc is not None:\n",
    "                Qc_in = []\n",
    "                Qc_out = []\n",
    "            index_del=[]\n",
    "            i=0\n",
    "            for i in range(len(ind)-1):\n",
    "                Q_i=-(Q[t,ind[i+1]]-Q[t,ind[i]])\n",
    "                if Qc is not None:\n",
    "                    Qc_i = -(Qc[t, ind[i+1]] - Qc[t, ind[i]])\n",
    "                if Q_i<0:\n",
    "                    Q_out_m.append(Q_i)\n",
    "                    if Qc is not None:\n",
    "                        Qc_out.append(Qc_i)\n",
    "                elif Q_i > 0:\n",
    "                    Q_in_m.append(Q_i)\n",
    "                    if Qc is not None:\n",
    "                        Qc_in.append(Qc_i)\n",
    "                else:\n",
    "                    index_del.append(i)\n",
    "                i+=1\n",
    "            div_val = np.delete(div_val, index_del)\n",
    "            ind = np.delete(ind, index_del)\n",
    "\n",
    "            #storing results\n",
    "            for i,qq in enumerate(Q_in_m):\n",
    "                Qin_ar[t,i] = qq\n",
    "            for i,qq in enumerate(Q_out_m):\n",
    "                Qout_ar[t,i] = qq\n",
    "            if Qc is not None:\n",
    "                for i,qq in enumerate(Qc_in):\n",
    "                    Qc_in_ar[t,i] = qq\n",
    "                for i,qq in enumerate(Qc_out):\n",
    "                    Qc_out_ar[t,i] = qq\n",
    "            for i,ss in enumerate(div_val):\n",
    "                divval_ar[t,i] = ss\n",
    "            for i,ss in enumerate(ind):\n",
    "                indices[t,i] = ss\n",
    "\n",
    "        #create a xarray Dataset for the results\n",
    "        out = xr.Dataset(\n",
    "        {\n",
    "            \"Qin\": ([\"time\", \"m\"], np.array(Qin_ar)),\n",
    "            \"Qout\": ([\"time\", \"n\"], np.array(Qout_ar)),\n",
    "            \"divval\": ([\"time\", \"o\"], np.array(divval_ar)),\n",
    "            \"index\": ([\"time\",\"o\"], np.array(indices).astype(int)),\n",
    "        },\n",
    "        coords={\n",
    "            \"time\": ([\"time\"], Q.time.data),\n",
    "            \"m\": ([\"m\"],np.arange(Qin_ar.shape[1])),\n",
    "            \"n\": ([\"n\"],np.arange(Qout_ar.shape[1])),\n",
    "            \"o\": ([\"o\"],np.arange(divval_ar.shape[1])),\n",
    "        },\n",
    "        )\n",
    "        if Qc is not None:\n",
    "            out[\"Qc_in\"] = ([\"time\", \"m\"], Qc_in_ar)\n",
    "            out[\"Qc_out\"] = ([\"time\", \"n\"], Qc_out_ar)\n",
    "\n",
    "    else:\n",
    "        #no time axis\n",
    "        if Q_thresh is None:\n",
    "        #set a default thresh\n",
    "            Q_thresh=0.01*np.max(np.abs(Q))\n",
    "        \n",
    "        if index is None:\n",
    "            ind,minmax = _find_extrema(Q,Q_thresh)\n",
    "        else:\n",
    "            ind=np.copy(index)\n",
    "        div_val=[]\n",
    "        i=0\n",
    "        while i < len(ind):\n",
    "                #print(Qvl[ind[i]])\n",
    "            div_val.append(coord_min+delta_var*ind[i])\n",
    "            i+=1\n",
    "                #print(smin+dss*ind[i])\n",
    "            #calculate transports etc.\n",
    "        Q_in_m=[]\n",
    "        Q_out_m=[]\n",
    "        if Qc is not None:\n",
    "            Qc_in = []\n",
    "            Qc_out = []\n",
    "        index_del=[]\n",
    "        i=0\n",
    "        for i in tqdm(range(len(ind)-1)):\n",
    "            Q_i=-(Q[ind[i+1]]-Q[ind[i]])\n",
    "            if Qc is not None:\n",
    "                Qc_i = -(Qc[ind[i+1]] - Qc[ind[i]])\n",
    "            if Q_i<0:\n",
    "                Q_out_m.append(Q_i)\n",
    "                if Qc is not None:\n",
    "                    Qc_out.append(Qc_i)\n",
    "            elif Q_i > 0:\n",
    "                Q_in_m.append(Q_i)\n",
    "                if Qc is not None:\n",
    "                    Qc_in.append(Qc_i)\n",
    "            else:\n",
    "                index_del.append(i)\n",
    "            i+=1\n",
    "        div_val = np.delete(div_val, index_del)\n",
    "        ind = np.delete(ind, index_del)\n",
    "\n",
    "        out = xr.Dataset(\n",
    "        {\n",
    "            \"Qin\": ([\"m\"], np.array(Q_in_m)),\n",
    "            \"Qout\": ([\"n\"], np.array(Q_out_m)),\n",
    "            \"divval\": ([\"o\"], np.array(div_val)),\n",
    "            \"index\": ([\"o\"], np.array(ind)),\n",
    "        },\n",
    "        coords={\n",
    "            \"m\": ([\"m\"],np.arange(len(Q_in_m))),\n",
    "            \"n\": ([\"n\"],np.arange(len(Q_out_m))),\n",
    "            \"o\": ([\"o\"],np.arange(len(div_val))),\n",
    "        }\n",
    "        )\n",
    "        if Qc is not None:\n",
    "            out[\"Qc_in\"] = ([\"m\"], Qc_in)\n",
    "            out[\"Qc_out\"] = ([\"n\"], Qc_out)\n",
    "    return(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "fec04b48",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<h4 id=\"calc_bulk_values\" class=\"doc_header\"><code>calc_bulk_values</code><a href=\"__main__.py#L2\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
       "\n",
       "> <code>calc_bulk_values</code>(**`coord`**, **`Q`**, **`Q_thresh`**=*`None`*, **`index`**=*`None`*, **\\*\\*`kwargs`**)\n",
       "\n",
       "Calculates the bulk values from a provided Q profile using the dividing salinity approach \n",
       "proposed by MacCready et al. (2018) and described/tested in detail by Lorenz et al. (2019)"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_doc(calc_bulk_values)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea120307",
   "metadata": {},
   "source": [
    "`coords` - coordinate corresponding to the interface values of `Q` (previously defined as `var_Q`\n",
    "\n",
    "`Q` - intergrated transport\n",
    "\n",
    "`Q_thresh` - threshold of the minimum integrated transport `Q` that will be considered. Default value if None: $0.01 * max(|Q|)$.\n",
    "\n",
    "`index` - array of indices of previously computed dividing values (e.g. salinities). If None the values will be computed. Hence, this index array can be only used if the dividing values are already known, but if they have been computed before, it can be used to save computation time.\n",
    "\n",
    "**Returns:**\n",
    "\n",
    "xarray dataset with `Qin, Qout` which are the bulk values of the inflow and outflow (multiple inflows and outflow are possible), respectively. Also returns `divval, index` which are the dividing values (e.g. salinities) that are separating inflows and outflows and their corresponding indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "3a9fbc90",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _get_time_array(x):\n",
    "    #check if x has a time property, i.e. is a dataset\n",
    "    if isinstance(x, np.ndarray):\n",
    "        print('numpy array -> creating artificial time axis')\n",
    "        time_array=np.arange(x.shape[0])\n",
    "    elif isinstance(x,xr.Dataset) or isinstance(x,xr.DataArray):\n",
    "        print('is xr.Dataset or xr.DataArray')\n",
    "        if not 'time' in x.dims:\n",
    "            print('has no time axis -> creating artificial one')\n",
    "            time_array=np.arange(x.shape[0])\n",
    "        else:\n",
    "            print('using existing time axis')\n",
    "            time_array = x['time']\n",
    "    else:\n",
    "        #create an artificial time array\n",
    "        time_array=np.arange(x.shape[0])\n",
    "    return time_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "699ff470",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _find_extrema(x, min_transport):\n",
    "    \"\"\"\n",
    "    internal function called by calc_bulk values to find the extrema in the transport function x\n",
    "    and label them correctly, see Appendix B in Lorenz et al. (2019).\n",
    "    x: Q(S)\n",
    "    min_transport: Q_thresh\n",
    "    \"\"\"\n",
    "    if np.count_nonzero(x)==0 or np.isnan(x).all():\n",
    "        indices=[0]\n",
    "        minmax=[0]\n",
    "        return(indices,minmax)\n",
    "    else:\n",
    "        ###\n",
    "        #set a minimum value to get rid of numerical noise\n",
    "        ###\n",
    "        if min_transport<=10**(-10):\n",
    "            min_transport=10**(-10)\n",
    "\n",
    "        ####\n",
    "        #finding all extrema by evaluating each data point\n",
    "        ####\n",
    "        comp=1\n",
    "        indices = []\n",
    "        minmax = []\n",
    "        i = 0\n",
    "        while i < np.shape(x)[0]:\n",
    "            if i-comp < 0:\n",
    "                a = 0\n",
    "            else:\n",
    "                a=i-comp\n",
    "            if i+comp+1>=len(x):\n",
    "                b=None\n",
    "                #c=i\n",
    "            else:\n",
    "                b=i+comp+1\n",
    "                #c=b\n",
    "            if x[i] == np.max(x[a:b]) and np.max(x[a:b]) != np.min(x[a:b]):# and x[i] != x[a]:\n",
    "                indices.append(i)\n",
    "                minmax.append('max')\n",
    "            elif x[i] == np.min(x[a:b]) and np.max(x[a:b]) != np.min(x[a:b]):# and (x[i] != x[c] or x[i] != x[a]):\n",
    "                indices.append(i)\n",
    "                minmax.append('min')\n",
    "            i+=1\n",
    "        #print(indices,minmax)\n",
    "        #print(x[indices])\n",
    "\n",
    "        ###\n",
    "        #correct consecutive extrema of the same kind, e.g., min min min or max max max (especially in the beginning and end of the salinity array)\n",
    "        ###\n",
    "\n",
    "        #index=[]\n",
    "        ii=1\n",
    "        while ii < len(indices):\n",
    "            index=[]\n",
    "            if minmax[ii] == minmax[ii-1]:\n",
    "                if minmax[ii] == 'max': #note the index of the smaller maximum\n",
    "                    if x[indices[ii]]>=x[indices[ii-1]]:\n",
    "                        index.append(ii-1)\n",
    "                    else:\n",
    "                        index.append(ii)\n",
    "                elif minmax[ii] == 'min': #note the index of the greater minimum\n",
    "                    if x[indices[ii]]<=x[indices[ii-1]]:\n",
    "                        index.append(ii-1)\n",
    "                    else:\n",
    "                        index.append(ii)\n",
    "                minmax = np.asarray(minmax)\n",
    "                indices = np.asarray(indices)\n",
    "                indices = np.delete(indices, index)\n",
    "                minmax = np.delete(minmax, index)\n",
    "            else:\n",
    "                ii+=1\n",
    "\n",
    "        ####\n",
    "        #delete too small transports\n",
    "        ####\n",
    "\n",
    "        ii=0\n",
    "        while ii < len(indices)-1: \n",
    "            index=[]\n",
    "            if np.abs(x[indices[ii+1]]-x[indices[ii]]) < min_transport:\n",
    "                if ii == 0: #if smin is involved and the transport is too small, smin has to change its min or max property\n",
    "                    index.append(ii+1)\n",
    "                    if minmax[ii] == 'min':\n",
    "                        minmax[ii] = 'max'\n",
    "                    else:\n",
    "                        minmax[ii] = 'min'\n",
    "                elif ii+1==len(indices)-1:#if smax is involved and the transport is too small, smin has to change its min or max property\n",
    "                    index.append(ii)\n",
    "                    if minmax[ii+1] == 'min':\n",
    "                        minmax[ii+1] = 'max'\n",
    "                    else:\n",
    "                        minmax[ii+1] = 'min'\n",
    "                else: #else both involved div sals are kicked out\n",
    "                    if ii+2 < len(indices)-1:\n",
    "                    #check and compare to i+2\n",
    "                        if minmax[ii]=='min':\n",
    "                            if x[indices[ii+2]]>x[indices[ii]]:\n",
    "                                index.append(ii+2)\n",
    "                                index.append(ii+1)\n",
    "                            else:\n",
    "                                index.append(ii)\n",
    "                                index.append(ii+1)\n",
    "                        elif minmax[ii]=='max':\n",
    "                            if x[indices[ii+2]]<x[indices[ii]]:\n",
    "                                index.append(ii+2)\n",
    "                                index.append(ii+1)\n",
    "                            else:\n",
    "                                index.append(ii)\n",
    "                                index.append(ii+1)\n",
    "                    else:\n",
    "                        index.append(ii)\n",
    "                        index.append(ii+1)\n",
    "                indices = np.delete(indices, index)\n",
    "                minmax = np.delete(minmax, index)\n",
    "            else:\n",
    "                ii+=1\n",
    "\n",
    "        ###\n",
    "        #so far the first and last minmax does not correspond to smin and smax of the data, expecially smin due to numerical errors (only makes sense)\n",
    "        #correct smin index\n",
    "        ###\n",
    "\n",
    "        if len(x)>4:\n",
    "            ii=1\n",
    "            while np.abs(np.abs(x[ii])-np.abs(x[0])) < 10**(-10) and ii < len(x)-1:\n",
    "                ii+=1\n",
    "            indices[0]=ii-1\n",
    "            #correct smax index\n",
    "            if x[-1]==0: #for low salinity classes Q[-1] might not be zero as supposed.\n",
    "                jj=-1\n",
    "                while x[jj] == 0 and np.abs(jj) < len(x)-1:\n",
    "                    jj -=1\n",
    "                indices[-1] = len(x)+jj+1\n",
    "        return indices,minmax"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "5ddaae1b1aacab30e6f9d299f885d3d1a71d6c802be31530f91bd978b8a5f689"
  },
  "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}