to_wavespec_estimator.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "403d193b-0105-4df8-900d-9d162aeb94c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This works only if you restore it to the 1d spec Jake originally sent.\n",
    "# Putting this on hold because APL people would rather compare stats computed from wave components a1 etc.\n",
    "\n",
    "import numpy as np\n",
    "from roguewave.wavespectra.estimators.mem2 import mem2\n",
    "from roguewave.wavespectra.estimators.mem import mem\n",
    "from numba_progress import ProgressBar\n",
    "from typing import Literal\n",
    "\n",
    "import os\n",
    "import pickle\n",
    "\n",
    "import pandas as pd\n",
    "import wavespectra \n",
    "import xarray as xr\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "Estimators = Literal[\"mem\", \"mem2\"]\n",
    "\n",
    "# These were copied from roguewave\n",
    "# -----------------------------------------------------------------------------\n",
    "#                       Boilerplate Interfaces\n",
    "# -----------------------------------------------------------------------------\n",
    "def estimate_directional_spectrum_from_moments(\n",
    "    e: np.ndarray,\n",
    "    a1: np.ndarray,\n",
    "    b1: np.ndarray,\n",
    "    a2: np.ndarray,\n",
    "    b2: np.ndarray,\n",
    "    direction: np.ndarray,\n",
    "    method: Estimators = \"mem2\",\n",
    "    **kwargs,\n",
    ") -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Construct a 2D directional distribution based on the directional moments and a spectral\n",
    "    reconstruction method.\n",
    "\n",
    "    :param number_of_directions: length of the directional vector for the\n",
    "    2D spectrum. Directions returned are in degrees\n",
    "\n",
    "    :param method: Choose a method in ['mem','mem2']\n",
    "        mem: maximum entrophy (in the Boltzmann sense) method\n",
    "        Lygre, A., & Krogstad, H. E. (1986). Explicit expression and\n",
    "        fast but tends to create narrow spectra anderroneous secondary peaks.\n",
    "\n",
    "        mem2: use entrophy (in the Shannon sense) to maximize. Likely\n",
    "        best method see- Benoit, M. (1993).\n",
    "\n",
    "    REFERENCES:\n",
    "    Benoit, M. (1993). Practical comparative performance survey of methods\n",
    "        used for estimating directional wave spectra from heave-pitch-roll data.\n",
    "        In Coastal Engineering 1992 (pp. 62-75).\n",
    "\n",
    "    Lygre, A., & Krogstad, H. E. (1986). Maximum entropy estimation of the\n",
    "        directional distribution in ocean wave spectra.\n",
    "        Journal of Physical Oceanography, 16(12), 2052-2060.\n",
    "\n",
    "    \"\"\"\n",
    "    return (\n",
    "        estimate_directional_distribution(a1, b1, a2, b2, direction, method, **kwargs)\n",
    "        * e[..., None]\n",
    "    )\n",
    "\n",
    "\n",
    "def estimate_directional_distribution(\n",
    "    a1: np.ndarray,\n",
    "    b1: np.ndarray,\n",
    "    a2: np.ndarray,\n",
    "    b2: np.ndarray,\n",
    "    direction: np.ndarray,\n",
    "    method: Estimators = \"mem2\",\n",
    "    **kwargs,\n",
    ") -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Construct a 2D directional distribution based on the directional moments and a spectral\n",
    "    reconstruction method.\n",
    "\n",
    "    :param number_of_directions: length of the directional vector for the\n",
    "    2D spectrum. Directions returned are in degrees\n",
    "\n",
    "    :param method: Choose a method in ['mem','mem2']\n",
    "        mem: maximum entrophy (in the Boltzmann sense) method\n",
    "        Lygre, A., & Krogstad, H. E. (1986). Explicit expression and\n",
    "        fast but tends to create narrow spectra anderroneous secondary peaks.\n",
    "\n",
    "        mem2: use entrophy (in the Shannon sense) to maximize. Likely\n",
    "        best method see- Benoit, M. (1993).\n",
    "\n",
    "    REFERENCES:\n",
    "    Benoit, M. (1993). Practical comparative performance survey of methods\n",
    "        used for estimating directional wave spectra from heave-pitch-roll data.\n",
    "        In Coastal Engineering 1992 (pp. 62-75).\n",
    "\n",
    "    Lygre, A., & Krogstad, H. E. (1986). Maximum entropy estimation of the\n",
    "        directional distribution in ocean wave spectra.\n",
    "        Journal of Physical Oceanography, 16(12), 2052-2060.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    # Jacobian to transform distribution as function of radian angles into\n",
    "    # degrees.\n",
    "    Jacobian = np.pi / 180\n",
    "    direction_radians = direction * Jacobian\n",
    "\n",
    "    if method.lower() in [\"maximum_entropy_method\", \"mem\"]:\n",
    "        # reconstruct the directional distribution using the maximum entropy\n",
    "        # method.\n",
    "        function = mem\n",
    "    elif method.lower() in [\"maximum_entrophy_method2\", \"mem2\"]:\n",
    "        function = mem2\n",
    "    else:\n",
    "        raise Exception(f\"unsupported spectral estimator method: {method}\")\n",
    "\n",
    "    output_shape = list(a1.shape) + [len(direction)]\n",
    "    if a1.ndim == 1:\n",
    "        input_shape = [1, a1.shape[-1]]\n",
    "    else:\n",
    "        input_shape = [np.prod(a1.shape[0:-1]), a1.shape[-1]]\n",
    "\n",
    "    a1 = a1.reshape(input_shape)\n",
    "    b1 = b1.reshape(input_shape)\n",
    "    a2 = a2.reshape(input_shape)\n",
    "    b2 = b2.reshape(input_shape)\n",
    "\n",
    "    number_of_iterations = a1.shape[0]\n",
    "    if number_of_iterations < 10:\n",
    "        disable = True\n",
    "    else:\n",
    "        disable = False\n",
    "\n",
    "    if method != \"mem2\":\n",
    "        msg = f\"Reconstructing 2d spectrum with {method} using implementation: \"\n",
    "    else:\n",
    "        solution_method = kwargs.get(\"solution_method\", \"scipy\")\n",
    "        msg = f\"Reconstructing 2d spectrum with {method} using solution_method {solution_method}\"\n",
    "\n",
    "    with ProgressBar(total=number_of_iterations, disable=disable, desc=msg) as progress:\n",
    "        res = function(direction_radians, a1, b1, a2, b2, progress, **kwargs)\n",
    "\n",
    "    return res.reshape(output_shape) * Jacobian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8c52f700-2ee2-4920-a2a5-e3057b8f23fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIRECTORY = '/vortexfs1/home/csherwood/proj/NOPP/buoy_data/'\n",
    "DATA_FILENAME = 'hurricane_ian_spotter_data_v1.pickle'\n",
    "\n",
    "with open(os.path.join(DATA_DIRECTORY, DATA_FILENAME), 'rb') as handle:\n",
    "    spotter = pickle.load(handle)\n",
    "\n",
    "# `spotter` is a python dictionary of Pandas DataFrames, keyed by\n",
    "# each drifter ID. The drifter ids can then be accessed as follows:\n",
    "spotter_ids = list(spotter.keys())\n",
    "spotter_id = spotter_ids[0]\n",
    "\n",
    "# Extract the observation times that contain spectral data.\n",
    "only_waves = spotter[spotter_id]['energy_density'].notnull()\n",
    "drifter = spotter[spotter_id][only_waves]\n",
    "\n",
    "# Exract the coordinate arrays; note that the frequency array is\n",
    "# uniform across the Spotter observations, so we can just \n",
    "# use the array in the first index of the DataFrame.\n",
    "time = drifter.index.to_numpy()\n",
    "freq = drifter['frequency'][0] \n",
    "\n",
    "# Extract the variable arrays.\n",
    "efth = np.stack(drifter['energy_density'])\n",
    "lat = drifter['latitude']\n",
    "lon = drifter['longitude']\n",
    "\n",
    "# Construct the dataset. This must match the conventions used by the\n",
    "# wavespectra package:\n",
    "# (https://wavespectra.readthedocs.io/en/latest/conventions.html#)\n",
    "# The directional spectrum is computed using the\n",
    "# directional moments and an estimator (e.g. MEM).\n",
    "directions = np.arange(0., 360., 10)\n",
    "efth_array = np.zeros( (len(time), 1, 1, len(freq), len(directions)))\n",
    "for i in range(len(time)):\n",
    "    #print(np.sum(drifter.a1[i]))\n",
    "    ea = estimate_directional_distribution( drifter.a1[i], drifter.b1[i], \\\n",
    "                                                     drifter.a2[i], drifter.b2[i], directions, method='mem2' )\n",
    "    #print(np.shape(ea))\n",
    "    efth_array[i,0,0,:,:] = ea\n",
    "    #print( np.sum(ea) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "61529edc-a2aa-410e-ab1f-abc7a86ac363",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(72, 1, 1, 39, 36)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(efth_array)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2b0babe7-5109-4bfd-80ac-969edfdd2456",
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "dimension 'lat' already exists as a scalar variable",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 43\u001b[0m\n\u001b[1;32m     39\u001b[0m     efth_array[i,\u001b[38;5;241m0\u001b[39m,\u001b[38;5;241m0\u001b[39m,:,:] \u001b[38;5;241m=\u001b[39m ea\n\u001b[1;32m     40\u001b[0m     \u001b[38;5;66;03m#print( np.sum(ea) )\u001b[39;00m\n\u001b[1;32m     41\u001b[0m \n\u001b[1;32m     42\u001b[0m \u001b[38;5;66;03m# efth(time, freq) and not efth(time, freq, dir).\u001b[39;00m\n\u001b[0;32m---> 43\u001b[0m ds \u001b[38;5;241m=\u001b[39m \u001b[43mxr\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mDataset\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     44\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdata_vars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m     45\u001b[0m \u001b[43m        \u001b[49m\u001b[43mefth\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtime\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mlat\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mlon\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mfreq\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mdir\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mefth_array\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     46\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;66;43;03m# efth=([\"time\", \"freq\"], efth), # Non-directinal...delete this line if using above (Note: efth may need to be transposed)\u001b[39;49;00m\n\u001b[1;32m     47\u001b[0m \u001b[43m        \u001b[49m\u001b[43msite\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     48\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     49\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcoords\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m     50\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtime\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtime\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     51\u001b[0m \u001b[43m        \u001b[49m\u001b[43mlat\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     52\u001b[0m \u001b[43m        \u001b[49m\u001b[43mlon\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     53\u001b[0m \u001b[43m        \u001b[49m\u001b[43mfreq\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfreq\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     54\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mdir\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdirections\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;66;43;03m#TODO: need to compute this from directional moments\u001b[39;49;00m\n\u001b[1;32m     55\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     56\u001b[0m \u001b[43m    \u001b[49m\u001b[43mattrs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m     57\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;66;43;03m# Attributes here; wavespectra would put significant wave\u001b[39;49;00m\n\u001b[1;32m     58\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;66;43;03m# height, etc., here but these are a function of time so it\u001b[39;49;00m\n\u001b[1;32m     59\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;66;43;03m# might be reasonable to set them as data_vars instead (as \u001b[39;49;00m\n\u001b[1;32m     60\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;66;43;03m# a function of the time coordinate)\u001b[39;49;00m\n\u001b[1;32m     61\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     62\u001b[0m \u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/mambaforge/envs/NOPP/lib/python3.9/site-packages/xarray/core/dataset.py:605\u001b[0m, in \u001b[0;36mDataset.__init__\u001b[0;34m(self, data_vars, coords, attrs)\u001b[0m\n\u001b[1;32m    602\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(coords, Dataset):\n\u001b[1;32m    603\u001b[0m     coords \u001b[38;5;241m=\u001b[39m coords\u001b[38;5;241m.\u001b[39mvariables\n\u001b[0;32m--> 605\u001b[0m variables, coord_names, dims, indexes, _ \u001b[38;5;241m=\u001b[39m \u001b[43mmerge_data_and_coords\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    606\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdata_vars\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcoords\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcompat\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mbroadcast_equals\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\n\u001b[1;32m    607\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    609\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_attrs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mdict\u001b[39m(attrs) \u001b[38;5;28;01mif\u001b[39;00m attrs \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    610\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_close \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/mambaforge/envs/NOPP/lib/python3.9/site-packages/xarray/core/merge.py:575\u001b[0m, in \u001b[0;36mmerge_data_and_coords\u001b[0;34m(data_vars, coords, compat, join)\u001b[0m\n\u001b[1;32m    573\u001b[0m objects \u001b[38;5;241m=\u001b[39m [data_vars, coords]\n\u001b[1;32m    574\u001b[0m explicit_coords \u001b[38;5;241m=\u001b[39m coords\u001b[38;5;241m.\u001b[39mkeys()\n\u001b[0;32m--> 575\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmerge_core\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    576\u001b[0m \u001b[43m    \u001b[49m\u001b[43mobjects\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    577\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcompat\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    578\u001b[0m \u001b[43m    \u001b[49m\u001b[43mjoin\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    579\u001b[0m \u001b[43m    \u001b[49m\u001b[43mexplicit_coords\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mexplicit_coords\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    580\u001b[0m \u001b[43m    \u001b[49m\u001b[43mindexes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mIndexes\u001b[49m\u001b[43m(\u001b[49m\u001b[43mindexes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcoords\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    581\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/mambaforge/envs/NOPP/lib/python3.9/site-packages/xarray/core/merge.py:761\u001b[0m, in \u001b[0;36mmerge_core\u001b[0;34m(objects, compat, join, combine_attrs, priority_arg, explicit_coords, indexes, fill_value)\u001b[0m\n\u001b[1;32m    756\u001b[0m prioritized \u001b[38;5;241m=\u001b[39m _get_priority_vars_and_indexes(aligned, priority_arg, compat\u001b[38;5;241m=\u001b[39mcompat)\n\u001b[1;32m    757\u001b[0m variables, out_indexes \u001b[38;5;241m=\u001b[39m merge_collected(\n\u001b[1;32m    758\u001b[0m     collected, prioritized, compat\u001b[38;5;241m=\u001b[39mcompat, combine_attrs\u001b[38;5;241m=\u001b[39mcombine_attrs\n\u001b[1;32m    759\u001b[0m )\n\u001b[0;32m--> 761\u001b[0m dims \u001b[38;5;241m=\u001b[39m \u001b[43mcalculate_dimensions\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvariables\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    763\u001b[0m coord_names, noncoord_names \u001b[38;5;241m=\u001b[39m determine_coords(coerced)\n\u001b[1;32m    764\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m explicit_coords \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "File \u001b[0;32m~/mambaforge/envs/NOPP/lib/python3.9/site-packages/xarray/core/variable.py:3201\u001b[0m, in \u001b[0;36mcalculate_dimensions\u001b[0;34m(variables)\u001b[0m\n\u001b[1;32m   3199\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m dim, size \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(var\u001b[38;5;241m.\u001b[39mdims, var\u001b[38;5;241m.\u001b[39mshape):\n\u001b[1;32m   3200\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m dim \u001b[38;5;129;01min\u001b[39;00m scalar_vars:\n\u001b[0;32m-> 3201\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m   3202\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdimension \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mdim\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m already exists as a scalar variable\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   3203\u001b[0m         )\n\u001b[1;32m   3204\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m dim \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m dims:\n\u001b[1;32m   3205\u001b[0m         dims[dim] \u001b[38;5;241m=\u001b[39m size\n",
      "\u001b[0;31mValueError\u001b[0m: dimension 'lat' already exists as a scalar variable"
     ]
    }
   ],
   "source": [
    "# efth(time, freq) and not efth(time, freq, dir).\n",
    "ds = xr.Dataset(\n",
    "    data_vars=dict(\n",
    "        efth=([\"time\", \"lat\", \"lon\", \"freq\", \"dir\"], efth_array),\n",
    "        # efth=([\"time\", \"freq\"], efth), # Non-directinal...delete this line if using above (Note: efth may need to be transposed)\n",
    "        site=\"\",\n",
    "    ),\n",
    "    coords=dict(\n",
    "        time=time,\n",
    "        lat=0,\n",
    "        lon=0,\n",
    "        freq=freq,\n",
    "        dir=directions, #TODO: need to compute this from directional moments\n",
    "    ),\n",
    "    attrs=dict(\n",
    "        # Attributes here; wavespectra would put significant wave\n",
    "        # height, etc., here but these are a function of time so it\n",
    "        # might be reasonable to set them as data_vars instead (as \n",
    "        # a function of the time coordinate)\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51de91c8-7815-4d48-b5a4-908b3e84700c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d83b343-0d43-4c20-b3f5-1651153335fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.efth.spec.hs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3737737c-17d9-44ed-bdd1-10b2570432b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "efth_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "191a5f7b-7e4b-4018-ae79-b47d4fb7d51c",
   "metadata": {},
   "outputs": [],
   "source": [
    "directions = np.arange(0., 360., 10)\n",
    "efth_array = np.zeros( (len(time), len(freq), len(directions)))\n",
    "print(np.shape(efth_array))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3ce758f-471e-4865-bfe6-d43350ff4ecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(time)):\n",
    "   efth_array[i,:,:] = estimate_directional_distribution( drifter.a1[i], drifter.b1[i], \\\n",
    "                                                         drifter.a2[i], drifter.b2[i], directions, method='mem' )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b99490d9-ce52-49ca-ab58-2575f7c41d00",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.shape(efth_array)\n",
    "efth_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcea1468-450b-4e2f-8ed9-d65b38de6dd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.pcolormesh(np.squeeze(efth_array[20]))"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}