reof.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rotated eofs \n",
    "=====\n",
    "\n",
    "Author: [Ray Bell](https://github.com/raybellwaves)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from IPython.display import Image\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "import cartopy.crs as ccrs\n",
    "import holoviews as hv\n",
    "import geoviews as gv\n",
    "import geoviews.feature as gf\n",
    "hv.notebook_extension()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# See pre_proc/README for notes on creating this file.\n",
    "da = xr.open_dataarray('NCEP_z500_monthly_1950-2000_NH.nc')\n",
    "print(da)\n",
    "\n",
    "# 3-month rolling mean to create seasonal averages e.g. JFM, FMA, MAM, ...\n",
    "# centre=True gives the time coordinate for the middle of the window\n",
    "sm = da.rolling(time=3, center=True).mean().dropna('time')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose a season of interest (0 = JFM, ..., 11 = DJF)\n",
    "season_dim = ['JFM', 'FMA', 'MAM', 'AMJ', 'MJJ', 'JJA', 'JAS', 'ASO', 'SON', 'OND', 'NDJ', 'DJF']\n",
    "season = 0\n",
    "\n",
    "da_sea = sm.isel(time=slice(season, -1, 12))\n",
    "# Make sure you've grabbed all seasons e.g. JFM should give you YYYY-02-DDYHH...\n",
    "print(da_sea.coords['time'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate anomalies\n",
    "#da_sea_anom = da_sea - da_sea.mean(dim=('time'))\n",
    "# Calculate standarzied anomalies\n",
    "da_sea_anom = (da_sea - da_sea.mean(dim=('time'))) / (da_sea - da_sea.mean(dim=('time'))).std(dim=('time'))\n",
    "print(da_sea_anom)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Weight data as square root of cos(lat)\n",
    "data = da_sea_anom.values\n",
    "wgts = np.expand_dims(np.sqrt(np.cos(np.deg2rad(da_sea_anom.coords['lat'].values)).clip(0., 1.)), axis=1)\n",
    "# Project the weights to 2d\n",
    "weights = np.broadcast_arrays(data[0:1], wgts)[1][0] # The [1][0] indexes (2, 1, 71, 360) to (71, 360)\n",
    "data_weighted = data * weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reshape the data to be (time, space)\n",
    "records = len(da_sea_anom.coords['time'])\n",
    "originalshape = data_weighted.shape[1:]\n",
    "channels = np.prod(originalshape)\n",
    "data_weighted_flat = data_weighted.reshape([records, channels])\n",
    "print(np.shape(data_weighted_flat))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the singular value decomposition (https://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.svd.html)\n",
    "# Principal component, eigenvalue, eof\n",
    "A, Lh, E = np.linalg.svd(data_weighted_flat, full_matrices=False)\n",
    "print(np.shape(A))\n",
    "print(np.shape(Lh))\n",
    "print(np.shape(E))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct the eigenvalues and normalize by N-1\n",
    "L = (Lh * Lh) / (float(records - 1))\n",
    "print(np.max(L))\n",
    "\n",
    "# Remove the scaling on the principal component time-series that is\n",
    "# implicitily introduced by using SVD instead of eigen-decomposition.\n",
    "# The PCs may be re-scaled later if required.\n",
    "P = A * Lh\n",
    "print(np.shape(P))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use 10 eofs in rotation\n",
    "neofs = 10\n",
    "\n",
    "eofs = E[0:neofs, :].copy()\n",
    "eofsall = E.copy()\n",
    "print(np.min(eofs), np.max(eofs))\n",
    "\n",
    "# Keep a copy without normalizing for plotting\n",
    "eofs2d = eofs.reshape((neofs,) + originalshape)\n",
    "# Put eofs in an xr.DataArray for easy of plotting\n",
    "eofs2d_da = xr.DataArray(eofs2d, coords=[range(eofs2d.shape[0]), da_sea_anom.coords['lat'], da_sea_anom.coords['lon']],\n",
    "                         dims=['mode', 'lat', 'lon'], name='eofs')\n",
    "\n",
    "# Normalize by square root of eigenvalues\n",
    "eofs = eofs / np.sqrt(L[0:neofs])[:, np.newaxis]\n",
    "print(np.shape(eofs))\n",
    "print(np.min(eofs), np.max(eofs))\n",
    "eofsall = E / np.sqrt(L)[:, np.newaxis]\n",
    "\n",
    "pcs = P[0:neofs, :].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# variance explained by each mode\n",
    "sub_varexpl = (L[0:neofs] / L.sum()) * 100\n",
    "print(sub_varexpl)\n",
    "print(np.sum(sub_varexpl))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Have a look at all modes of variability using [geoviews](https://github.com/ioam/geoviews). Just move the slider."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%opts Image [projection=ccrs.Orthographic(0, 90) colorbar=False fig_size=200] (cmap='RdBu_r') Overlay [xaxis=None yaxis=None]\n",
    "dataset = gv.Dataset(eofs2d_da, kdims=['mode', 'lon', 'lat'])\n",
    "dataset.to(gv.Image, ['lon', 'lat']) * gf.coastline()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('top 4 eofs computed in NCL for season '+season_dim[season])\n",
    "Image(filename = 'NCL/'+season_dim[season]+'_eof.png', width=1000, height=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Based on a matlab function in the MATLAB folder\n",
    "# But doesn't work...\n",
    "\n",
    "#tol = 1e-10\n",
    "#it_max = 1000\n",
    "#p, nc = np.shape(pcs)\n",
    "\n",
    "# Normalize\n",
    "#scale = np.expand_dims(np.sqrt(np.diag(np.dot(pcs, pcs.T))), axis=1)\n",
    "#scale2 = np.repeat(scale, nc, axis=1)\n",
    "\n",
    "#pcs2 = pcs / scale\n",
    "\n",
    "#TT = np.eye(nc)\n",
    "#d = 0;\n",
    "\n",
    "#for i in range(it_max):\n",
    "#    z = np.dot(pcs2,TT)\n",
    "#    zdiag = np.diag(np.squeeze(np.dot(np.ones((1,p)), (z ** 2)))) / p\n",
    "#    ztmp = np.dot(z,zdiag)\n",
    "#    ztmp2 = (z ** 3) - ztmp\n",
    "#    B = np.dot(pcs2.T,ztmp2)\n",
    "\n",
    "#    U, S, V = np.linalg.svd(B)\n",
    "#    TT = np.dot(U, V)\n",
    "\n",
    "#    d2 = d\n",
    "#    d = np.sum(S)\n",
    "#    if d < (d2 * (1 + tol)): break\n",
    "\n",
    "#rpcs = np.dot(pcs2, TT)\n",
    "#rpcs = rpcs * scale\n",
    "\n",
    "#reofs = np.dot(eofsall.T, TT).T\n",
    "\n",
    "#reof2d = reofs.reshape((records,) + originalshape)\n",
    "#reof2d = reof2d[0:10,:,:]\n",
    "#reofs2d_da = xr.DataArray(reof2d,\n",
    "#                          coords=[range(10),\n",
    "#                                  da_sea_anom.coords['lat'],\n",
    "#                                  da_sea_anom.coords['lon']],\n",
    "#                          dims=['mode', 'lat', 'lon'], name='eofs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rotate the eofs\n",
    "# This is based on https://github.com/ajdawson/eofs/blob/experimental-rotation/lib/eofs/experimental/rotation/kernels.py\n",
    "# and https://github.com/bmcmenamin/fa_kit/blob/master/fa_kit/rotation.py\n",
    "# Not quite like the NCL output. i.e. there are spatial differences\n",
    "\n",
    "itermax = 1000\n",
    "gamma = 1.0 # For varimax\n",
    "\n",
    "# Apply Kaiser row normalization.\n",
    "scale = np.sqrt((eofs ** 2).sum(axis=0))\n",
    "# In some cases such as NCEP data the scaling gives 0.\n",
    "# This workaround replaces the 0's with a very small number to allow division\n",
    "scale[np.where(scale == 0)] = 1e-15\n",
    "\n",
    "eofs /= scale\n",
    "print(np.min(eofs), np.max(eofs))\n",
    "\n",
    "# Define the initial values of the rotation matrix and the convergence monitor.\n",
    "rot_mat = np.eye(neofs, dtype=eofs.dtype)\n",
    "delta = 0.\n",
    "\n",
    "# Iteratively compute the rotation matrix.\n",
    "for i in range(itermax):\n",
    "    z = np.dot(eofs.T, rot_mat)\n",
    "    zdiag = np.diag(np.mean(z ** 2, axis=0) * gamma) # Check axis here\n",
    "    \n",
    "    b = np.dot(eofs, z ** 3 - np.dot(z, zdiag))\n",
    "    \n",
    "    svd_u, svd_s, svd_v = np.linalg.svd(b)\n",
    "    \n",
    "    rot_mat = np.dot(svd_u, svd_v)\n",
    "    \n",
    "    delta_new = np.sum(svd_s)\n",
    "    if delta_new < delta: break\n",
    "    delta = delta_new\n",
    "        \n",
    "# Apply the rotation to the input EOFs.\n",
    "reofs = np.dot(eofs.T, rot_mat).T\n",
    "\n",
    "# de-normalize\n",
    "reofs *= scale\n",
    "print(np.min(reofs), np.max(reofs))\n",
    "\n",
    "# Show the variances of the modes (in descending order)\n",
    "divisor = (reofs ** 2).sum(axis=0).sum()\n",
    "reof_varexpl = ((reofs ** 2).sum(axis=1) / divisor) *100\n",
    "reof_varexpl = np.flip(np.sort(reof_varexpl), axis=0)\n",
    "print('reof var explained:')\n",
    "print(reof_varexpl)\n",
    "\n",
    "# Put the reofs into descending order by ...\n",
    "# I'll admit i'm not sure what this is doing but it works (black magic)\n",
    "abs_min = np.abs(np.min(reofs, axis=1))\n",
    "abs_max = np.abs(np.max(reofs, axis=1))\n",
    "to_flip = abs_max < abs_min\n",
    "reofs[to_flip, :] *= -1\n",
    "\n",
    "# Normalize by square root of the eigenvalues\n",
    "LR = np.sqrt(np.sum(reofs**2, axis=1, keepdims=True))\n",
    "reofs = reofs / LR\n",
    "print(np.min(reofs), np.max(reofs))\n",
    "\n",
    "# Put to xr.DataArray for plotting\n",
    "reofs2d = reofs.reshape((neofs,) + originalshape)\n",
    "# Multiply by minus 1 to match the color of the NCL plots\n",
    "reofs2d *= -1\n",
    "print(np.min(reofs), np.max(reofs))\n",
    "reofs2d_da = xr.DataArray(reofs2d,\n",
    "                          coords=[range(reofs2d.shape[0]),\n",
    "                                  da_sea_anom.coords['lat'],\n",
    "                                  da_sea_anom.coords['lon']],\n",
    "                          dims=['mode', 'lat', 'lon'], name='eofs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the reofs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%opts Image [projection=ccrs.Orthographic(0, 90) colorbar=False fig_size=200] (cmap='RdBu_r') Overlay [xaxis=None yaxis=None]\n",
    "dataset = gv.Dataset(reofs2d_da, kdims=['mode', 'lon', 'lat'])\n",
    "dataset.to(gv.Image, ['lon', 'lat']) * gf.coastline()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('top 4 reofs computed in NCL for season '+season_dim[season])\n",
    "Image(filename = 'NCL/'+season_dim[season]+'_rot_eof.png', width=1000, height=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the location of minimum and maximum for reof1 and reof2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reof1 = reofs2d_da.isel(mode=0)\n",
    "ix = reof1.where(reof1==reof1.max(), drop=True)\n",
    "latmax = ix.coords['lat'].values; lonmax = ix.coords['lon'].values\n",
    "ix = reof1.where(reof1==reof1.min(), drop=True)\n",
    "latmin = ix.coords['lat'].values; lonmin = ix.coords['lon'].values\n",
    "print(latmax, lonmax)\n",
    "print(latmin, lonmin)\n",
    "print('season number is (see third cell):', season)\n",
    "\n",
    "ax = plt.axes(projection=ccrs.Orthographic(-90, 90))\n",
    "ax.coastlines()\n",
    "ax.set_global()\n",
    "plt.plot(lonmax, latmax, color='blue', marker='x', transform=ccrs.PlateCarree())\n",
    "plt.plot(lonmin, latmin, color='red', marker='x', transform=ccrs.PlateCarree())\n",
    "plt.show()\n",
    "\n",
    "reof2 = reofs2d_da.isel(mode=1)\n",
    "ix = reof2.where(reof2==reof2.max(), drop=True)\n",
    "latmax = ix.coords['lat'].values; lonmax = ix.coords['lon'].values\n",
    "ix = reof2.where(reof2==reof2.min(), drop=True)\n",
    "latmin = ix.coords['lat'].values; lonmin = ix.coords['lon'].values\n",
    "print(latmax, lonmax)\n",
    "print(latmin, lonmin)\n",
    "\n",
    "ax = plt.axes(projection=ccrs.Orthographic(-90, 90))\n",
    "ax.coastlines()\n",
    "ax.set_global()\n",
    "plt.plot(lonmax, latmax, color='red', marker='x', transform=ccrs.PlateCarree())\n",
    "plt.plot(lonmin, latmin, color='blue', marker='x', transform=ccrs.PlateCarree())\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the CPC NAO definition for season = 11 (DJF), season = 2 (MAM), season = 5 (JJA), season = 8 (SON):\n",
    "![title](CPC/nao_correlation_map.gif)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}