{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f276fb94",
   "metadata": {},
   "source": [
    "# Daily Noise Source Inversion Analysis\n",
    "\n",
    "This jupyter notebook accompanies the website and contains code to analyse and plot the results of a noise source inversion. To run this notebook, please download all files for one day, unzip it, and assign the path to the \"project_path\" variable. The plots created in this notebook will be saved in a new folder within the project directory. They basically reproduce all plots in the \"output_plots\" directory.\n",
    "\n",
    "The packages needed for this are:\n",
    "\n",
    "- numpy\n",
    "- matplotlib\n",
    "- cartopy\n",
    "- pandas\n",
    "- h5py\n",
    "- pyyaml\n",
    "- cmasher (optional)\n",
    "- pprintpp (optional)\n",
    "\n",
    "These can be installed using (for example) anaconda:\n",
    "\n",
    "https://www.anaconda.com/products/individual\n",
    "\n",
    "If you already have conda installed (and are in your chosen environment) you can run this command to install all packages:\n",
    "\n",
    "<pre style='padding-left:20px;'>\n",
    "conda install numpy matplotlib cartopy pandas h5py pyyaml cmasher pprintpp --yes\n",
    "</pre>\n",
    "\n",
    "\n",
    "Feel free to contact us if you have any questions or if there are any problems:\n",
    "\n",
    "jonas.igel@erdw.ethz.ch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "773c2aad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import necessary packages\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy as cart\n",
    "import cartopy.feature as cfeature\n",
    "import matplotlib.tri as tri\n",
    "from matplotlib import colors\n",
    "import os\n",
    "from glob import glob\n",
    "from pandas import read_csv\n",
    "import h5py\n",
    "import yaml\n",
    "\n",
    "try:\n",
    "    import pprint\n",
    "    pp = pprint.PrettyPrinter(indent=4,sort_dicts=False)\n",
    "    pp_import = True\n",
    "except:\n",
    "    pp_import = False\n",
    "    \n",
    "try:\n",
    "    import cmasher as cmr\n",
    "    cmap = cmr.cosmic    # CMasher\n",
    "    cmap = plt.get_cmap('cmr.cosmic')   # MPL\n",
    "    cmash_import = True\n",
    "except:\n",
    "    cmash_import = False\n",
    "    cmap=plt.get_cmap('Blues_r')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbc15c3a",
   "metadata": {},
   "source": [
    "### Change the project_path parameter to the path to your downloaded inversion folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53af60b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print all the info?\n",
    "print_info = True\n",
    "\n",
    "# provide path to project (unzipped inversion folder)\n",
    "project_path = !!!!!YOUR PATH HERE!!!!!\n",
    "\n",
    "# check path\n",
    "if not os.path.isdir(project_path):\n",
    "    print(\"!!!!!\"*20+\"\\n\")\n",
    "    print(\"Please change the project_path parameter to your path to the inversion folder \\n\")\n",
    "    print(\"!!!!!\"*20+\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27cc6d96",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create paths\n",
    "project_name = os.path.basename(os.path.normpath(project_path))\n",
    "print(f\"Project name: {project_name} \\n\")\n",
    "\n",
    "output_path = project_path\n",
    "plot_path = os.path.join(project_path,'notebook_plots')\n",
    "\n",
    "if not os.path.exists(plot_path):\n",
    "    os.makedirs(plot_path)\n",
    "    \n",
    "    \n",
    "print(\"=\"*100+'\\n')\n",
    "print(f\"All plots will be saved here: {plot_path} \\n\")\n",
    "print(\"=\"*100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f77ceecd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# first read in config file, runtime file, stationlist, sourcegrid\n",
    "inv_config_file = os.path.join(output_path,\"inversion_config.yml\")\n",
    "runtime_file = os.path.join(output_path,\"runtime.txt\")\n",
    "station_file = os.path.join(output_path,\"stationlist.csv\")\n",
    "sourcegrid_file = os.path.join(output_path,\"sourcegrid.npy\")\n",
    "sourcegrid_voronoi_file = os.path.join(output_path,\"sourcegrid_voronoi.npy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33cfcecc",
   "metadata": {},
   "source": [
    "## Inversion parameters\n",
    "\n",
    "Let's first have a look at the inversion config files and the runtime file.\n",
    "The config files contain all the parameters that were used to run the inversion. \n",
    "The runtime file has information about the number of correlations used and how long each part of the inversion took."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a119eb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check all parameters used\n",
    "with open(inv_config_file) as f:\n",
    "    inv_config = yaml.safe_load(f)\n",
    "    \n",
    "if pp_import and print_info:    \n",
    "    pp.pprint(inv_config)\n",
    "elif print_info:\n",
    "    print(inv_config)\n",
    "    \n",
    "# get some info for plotting\n",
    "only_ocean = inv_config['svp_grid_config']['svp_only_ocean']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a220397",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get information from \n",
    "runtime_dict = {}\n",
    "file = open(runtime_file)\n",
    "\n",
    "for line in file:\n",
    "    key,value = line.split(':')\n",
    "    runtime_dict[key.replace(' ','_')] = value.strip()\n",
    "\n",
    "if pp_import and print_info:    \n",
    "    pp.pprint(runtime_dict)\n",
    "elif print_info:\n",
    "    print(runtime_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c7fd344",
   "metadata": {},
   "source": [
    "## Sourcegrid and station locations\n",
    "\n",
    "We now import the stationlist and source grid and visualise them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f2d0f84",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read in stationlist and sourcegrid\n",
    "stationlist = read_csv(station_file,keep_default_na=False)\n",
    "stat_lat = stationlist['lat']\n",
    "stat_lon = stationlist['lon']\n",
    "\n",
    "sourcegrid = np.load(sourcegrid_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96b73919",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot stationlist and sourcegrid\n",
    "plt.figure(figsize=(50,20))\n",
    "ax = plt.axes(projection=ccrs.Robinson(central_longitude=0))\n",
    "ax.set_global()\n",
    "\n",
    "if only_ocean:\n",
    "    ax.add_feature(cfeature.NaturalEarthFeature('cultural', 'admin_0_countries', '50m', edgecolor='black', facecolor=cfeature.COLORS['land']),zorder=2)\n",
    "else:\n",
    "    ax.coastlines(color='k',linewidth=1,zorder=2)\n",
    "    ax.coastlines(color='w',linewidth=2,zorder=1)\n",
    "\n",
    "plt.scatter(sourcegrid[0],sourcegrid[1],s=20,c='k',transform=ccrs.PlateCarree(),zorder=3)\n",
    "plt.title(f'Sourcegrid for {project_name}',fontsize=50,pad=30)\n",
    "plt.scatter(stat_lon,stat_lat,s=150,c='lawngreen',marker='^',edgecolor='k',linewidth=1,label='Stations',transform=ccrs.PlateCarree(),zorder=3)\n",
    "plt.legend(fontsize=25,loc=3)\n",
    "plt.savefig(os.path.join(plot_path,f'sourcegrid.png'),bbox_inches='tight',facecolor='white')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ef19e06",
   "metadata": {},
   "source": [
    "## Voronoi cells\n",
    "\n",
    "The voronoi cells are necessary to scale each noise sources as we use spatially variable grids for the inversion. For the global inversions the grid is homogeneous so the voronoi cells are all the same size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f9e53a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the voronoi cells which are used to scale the noise sources\n",
    "# triangulating the grid for this\n",
    "sourcegrid_voronoi = np.load(sourcegrid_voronoi_file)\n",
    "grid_tri = tri.Triangulation(sourcegrid_voronoi[0],sourcegrid_voronoi[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4a37652",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(50,20))\n",
    "ax = plt.axes(projection=ccrs.Robinson(central_longitude=0))\n",
    "ax.set_global()\n",
    "if only_ocean:\n",
    "    ax.add_feature(cfeature.NaturalEarthFeature('cultural', 'admin_0_countries', '50m', edgecolor='black', facecolor=cfeature.COLORS['land']),zorder=2)\n",
    "else:\n",
    "    ax.coastlines(color='k',linewidth=1,zorder=2)\n",
    "    ax.coastlines(color='w',linewidth=2,zorder=1)\n",
    "    \n",
    "#plt.scatter(sourcegrid[0],sourcegrid[1],s=20,c=sourcegrid_voronoi[2],cmap=cmap,transform=ccrs.PlateCarree(),zorder=3)\n",
    "plt.tripcolor(grid_tri,sourcegrid_voronoi[2],cmap=cmap,vmin=0,linewidth=0.0,edgecolor='none',zorder=1,transform=ccrs.Geodetic())\n",
    "cbar = plt.colorbar(pad=0.01)\n",
    "cbar.ax.tick_params(labelsize=30) \n",
    "cbar.set_label('Voronoi cell size',rotation=270,labelpad=60,fontsize=40)\n",
    "\n",
    "plt.title(f'Voronoi cells for {project_name}',fontsize=50,pad=30)\n",
    "plt.scatter(stat_lon,stat_lat,s=150,c='lawngreen',marker='^',edgecolor='k',linewidth=1,label='Stations',transform=ccrs.PlateCarree(),zorder=3)\n",
    "plt.legend(fontsize=25,loc=3)\n",
    "plt.savefig(os.path.join(plot_path,f'sourcegrid_voronoi.png'),bbox_inches='tight',facecolor='white')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e47441b",
   "metadata": {},
   "source": [
    "## Station ray coverage\n",
    "\n",
    "The file 'used_obs_corr_list.csv' contains all the station pairs that were actually used for the inversion. We use this list to create a station ray coverage plot. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1380625c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create dictionary with stations and their antipoles\n",
    "station_dict = dict()\n",
    "station_anti_dict = dict()\n",
    "\n",
    "for sta_i in stationlist.iterrows():\n",
    "    sta = sta_i[1]\n",
    "    station_dict.update({f\"{sta['net']}.{sta['sta']}\":[sta['lat'],sta['lon']]})\n",
    "\n",
    "    if sta['lon'] < 0:\n",
    "        station_anti_dict.update({f\"{sta['net']}.{sta['sta']}\":[-sta['lat'],sta['lon']+180]})\n",
    "    elif sta['lon'] >= 0:\n",
    "        station_anti_dict.update({f\"{sta['net']}.{sta['sta']}\":[-sta['lat'],sta['lon']-180]})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14bdf15e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read in list of station pairs used\n",
    "used_obs_corr = read_csv(os.path.join(output_path,'used_obs_corr_list.csv'),header=None)\n",
    "\n",
    "stat_pair_list = []\n",
    "\n",
    "for i in used_obs_corr.iterrows():\n",
    "    sta_1 = f\"{i[1][0].split('--')[0].split('.')[0]}.{i[1][0].split('--')[0].split('.')[1]}\"\n",
    "    sta_2 = f\"{i[1][0].split('--')[1].split('.')[0]}.{i[1][0].split('--')[1].split('.')[1]}\"\n",
    "    \n",
    "    stat_pair_list.append(f\"{sta_1}--{sta_2}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2586b823",
   "metadata": {},
   "outputs": [],
   "source": [
    "stat_pair_dict = {}\n",
    "n_rays = 0\n",
    "\n",
    "plt.figure(figsize=(50,20))\n",
    "ax = plt.axes(projection=ccrs.Robinson())\n",
    "ax.add_feature(cfeature.NaturalEarthFeature('cultural', 'admin_0_countries', '50m', edgecolor='black', facecolor=cfeature.COLORS['land']),zorder=1)\n",
    "#ax.coastlines()\n",
    "ax.set_global()\n",
    "\n",
    "for stat_pair in stat_pair_list:\n",
    "    i = stat_pair.split('--')[0]\n",
    "    j = stat_pair.split('--')[1]\n",
    "\n",
    "\n",
    "    stat_pair_dict.update({i:[station_dict[i][0],station_dict[i][1]]})\n",
    "    stat_pair_dict.update({j:[station_dict[j][0],station_dict[j][1]]})\n",
    "\n",
    "    plt.plot([station_dict[i][1],station_anti_dict[j][1]],[station_dict[i][0],station_anti_dict[j][0]], color='k',zorder=2, transform=ccrs.Geodetic(),alpha=4/np.size(list(station_dict.keys())))\n",
    "    plt.plot([station_anti_dict[i][1],station_dict[j][1]],[station_anti_dict[i][0],station_dict[j][0]], color='k',zorder=2, transform=ccrs.Geodetic(),alpha=4/np.size(list(station_dict.keys())))\n",
    "\n",
    "    n_rays += 1\n",
    "\n",
    "stat_lat = np.asarray(list(stat_pair_dict.values())).T[0]\n",
    "stat_lon = np.asarray(list(stat_pair_dict.values())).T[1]\n",
    "\n",
    "plt.scatter(stat_lon,stat_lat,marker='^',s=250,c='k',edgecolors='w',linewidths=2,transform=ccrs.PlateCarree(),zorder=3)\n",
    "\n",
    "plt.title(f\"Ray coverage for {project_name} with {n_rays} rays\",pad=30,fontsize=30)\n",
    "plt.savefig(os.path.join(plot_path,f'kernel_ray_coverage.png'),bbox_inches='tight',facecolor='white')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef5e6c86",
   "metadata": {},
   "source": [
    "## Inversion misfit\n",
    "\n",
    "The main goal of the inversion is to decrease the misfit between the measurement of the observed and synthetic cross-correlations. Here we plot the misfit of each iteration. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c943a9c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot misfit\n",
    "# misfit\n",
    "steps_avail_path = [os.path.join(output_path,i) for i in os.listdir(output_path) if i.startswith('iteration') and not os.path.isfile(os.path.join(output_path,i))]\n",
    "\n",
    "misfit_step = []\n",
    "misfit_dict = dict()\n",
    "\n",
    "for j in steps_avail_path:\n",
    "\n",
    "    measr_file_paths_var = [os.path.join(j,i) for i in os.listdir(j) if i.endswith('measurement.csv')]\n",
    "\n",
    "    #print(measr_file_paths_var)\n",
    "    if measr_file_paths_var == []:\n",
    "        print(f'No measurement for {os.path.basename(j)}')\n",
    "\n",
    "    else:\n",
    "        i = measr_file_paths_var[0]\n",
    "        step_nr_var = int(i.split('/')[-2].split('_')[1])\n",
    "        measr_step_var = read_csv(i,keep_default_na=False)\n",
    "\n",
    "        l2_norm_all = np.asarray(measr_step_var['l2_norm'])\n",
    "        l2_norm = l2_norm_all[~np.isnan(l2_norm_all)]\n",
    "        mf_step_var = np.mean(l2_norm)\n",
    "\n",
    "        misfit_step.append([step_nr_var,mf_step_var])  \n",
    "        misfit_dict.update({step_nr_var:mf_step_var})\n",
    "\n",
    "\n",
    "misfit = [i[1] for i in misfit_step]\n",
    "step = [i[0] for i in misfit_step]\n",
    "\n",
    "step,misfit = zip(*sorted(zip(step,misfit)))\n",
    "\n",
    "# misfit reduction\n",
    "mf_reduc = (1-(misfit[-1]/misfit[0]))*100\n",
    "\n",
    "#### MISFIT PLOTS\n",
    "plt.figure(figsize=(15,8))\n",
    "plt.plot(step,misfit,'k',marker='o',markerfacecolor='b',markersize=10)\n",
    "plt.grid()\n",
    "plt.xlabel('Iterations',fontsize=15)\n",
    "plt.xticks(fontsize=15)\n",
    "plt.yticks(fontsize=15)\n",
    "plt.title(f'Misfit for \"{project_name}\" reduced by {np.around(mf_reduc,2)}%',fontsize=25,pad=20)\n",
    "plt.savefig(os.path.join(plot_path,'misfit_vs_iterations.png'),bbox_inches='tight',facecolor='white')\n",
    "plt.show()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "410e6dcf",
   "metadata": {},
   "source": [
    "## Steplength test\n",
    "\n",
    "A steplength test is performed before each update of the noise source distribution. This uses half of the correlations to save computational cost and calculates the misfit for several steplengths. We then fit a exponential and take the minimum of it as our optimal steplength.\n",
    "\n",
    "Here we plot the misfits of the steplength tests. You can pick an iteration. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30beb5f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pick an iteration\n",
    "iteration_nr = 0\n",
    "\n",
    "\n",
    "### plot steplength tests\n",
    "steplength_files = sorted([os.path.join(i,'misfit_step_test.npy') for i in steps_avail_path if os.path.isfile(os.path.join(i,'misfit_step_test.npy'))])\n",
    "steplength_fit_files = sorted([os.path.join(i,'misfit_step_test_fit.npy') for i in steps_avail_path if os.path.isfile(os.path.join(i,'misfit_step_test_fit.npy'))])\n",
    "\n",
    "step = steplength_files[iteration_nr].split('/')[-2].split('_')[1]\n",
    "\n",
    "mf_step = np.asarray(np.load(steplength_files[iteration_nr],allow_pickle=True)[0])\n",
    "mf_final_step = np.load(steplength_files[iteration_nr],allow_pickle=True)[1]\n",
    "mf_step_fit = np.load(steplength_fit_files[iteration_nr],allow_pickle=True)\n",
    "\n",
    "step_m = [i[0] for i in mf_step]\n",
    "mf_m = [i[1] for i in mf_step]\n",
    "\n",
    "plt.figure(figsize=(15,8))\n",
    "plt.scatter(step_m,mf_m,c='r')\n",
    "#plt.scatter(mf_step[:,0],mf_step[:,1],c='r')\n",
    "plt.scatter(mf_final_step[0],mf_final_step[1],c='r',s=100,marker='x',label='Predicted misfit')\n",
    "plt.scatter(mf_final_step[0],misfit_dict[int(step)+1],c='b',s=100,marker='x',label='Actual misfit')\n",
    "plt.plot(mf_step_fit[0],mf_step_fit[1],c='b',label='Fitted exponential')\n",
    "plt.grid()\n",
    "plt.title(f'Steplength test for iteration {step} with final step {np.around(mf_final_step[0],2)}. Predicted misfit: {np.around(mf_final_step[1],2)}. Actual misfit: {np.around(misfit_dict[int(step)+1],2)}')\n",
    "plt.legend()\n",
    "plt.savefig(os.path.join(plot_path,f'iteration_{step}_0_slt.png'),bbox_inches='tight',facecolor='white')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaafc0c2",
   "metadata": {},
   "source": [
    "# Inversion models\n",
    "\n",
    "Now we plot the actual inversion models. First we check the available keys in the hdf5 model files.\n",
    "\n",
    "They are: ['coordinates', 'frequencies', 'model', 'spectral_basis', 'surface_areas']\n",
    "\n",
    "The frequencies and spectral_basis keys contain information about the frequency content of the noise sources that were used to model the cross-correlations. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe355df4",
   "metadata": {},
   "outputs": [],
   "source": [
    "inversionmodel_paths = sorted(glob(os.path.join(output_path,\"models/iteration*.h5\")))\n",
    "\n",
    "# load the first file and check content\n",
    "inversionmodel_0 = h5py.File(inversionmodel_paths[0],'r')\n",
    "\n",
    "print(list(inversionmodel_0.keys()))\n",
    "\n",
    "# To access the model and grid:\n",
    "# inversionmodel_0['model']\n",
    "# inversionmodel_0['coordinates']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fce78d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot frequency spectrum used for the forward modelling of the sources\n",
    "inversionmodel_0_freq = inversionmodel_0['frequencies'][()]\n",
    "inversionmodel_0_spec = inversionmodel_0['spectral_basis'][()]\n",
    "\n",
    "plt.figure(figsize=(20,10))\n",
    "plt.plot(inversionmodel_0_freq,inversionmodel_0_spec[0],linewidth=2,c='k')\n",
    "plt.xlabel(\"Frequency [Hz]\",fontsize=20)\n",
    "plt.ylabel(\"Normalised Power\",fontsize=20)\n",
    "plt.xticks(fontsize=15)\n",
    "plt.yticks(fontsize=15)\n",
    "plt.grid()\n",
    "plt.title(f'Source frequency spectrum for \"{project_name}\"',fontsize=30,pad=20)\n",
    "plt.savefig(os.path.join(plot_path,f'frequency_spectrum.png'),bbox_inches='tight',facecolor='white')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c0595d4",
   "metadata": {},
   "source": [
    "### Plot inversion models\n",
    "\n",
    "And now you can pick an iteration and plot the noise source model. If you set triangulation to True the grid will be triangulated to create a smoother plot. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "776fb39d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the sourcemodel\n",
    "iteration_nr = 0\n",
    "triangulation = True\n",
    "\n",
    "#print(step)\n",
    "source_distr_file = h5py.File(inversionmodel_paths[iteration_nr],'r')\n",
    "source_grid = np.asarray(source_distr_file['coordinates'])\n",
    "source_distr = np.asarray(source_distr_file['model']).T[0]\n",
    "source_distr_norm = source_distr/np.max(np.abs(source_distr))\n",
    "\n",
    "plt.figure(figsize=(50,20))\n",
    "ax = plt.axes(projection=ccrs.Robinson(central_longitude=0))\n",
    "ax.set_global()\n",
    "\n",
    "if only_ocean:\n",
    "    ax.add_feature(cfeature.NaturalEarthFeature('cultural', 'admin_0_countries', '50m', edgecolor='black', facecolor=cfeature.COLORS['land']),zorder=2)\n",
    "else:\n",
    "    ax.coastlines(color='k',linewidth=1,zorder=2)\n",
    "    ax.coastlines(color='w',linewidth=2,zorder=1)\n",
    "\n",
    "if triangulation:\n",
    "    triangles = tri.Triangulation(source_grid[0],source_grid[1])\n",
    "\n",
    "    if cmash_import:\n",
    "        plt.tripcolor(triangles,source_distr_norm,cmap=cmap,linewidth=0.0,edgecolor='none',vmin=0,zorder=1,transform=ccrs.Geodetic())\n",
    "    else:\n",
    "        plt.tripcolor(triangles,source_distr_norm,cmap=plt.get_cmap('Blues_r'),linewidth=0.0,edgecolor='none',vmin=0,zorder=1,transform=ccrs.Geodetic())\n",
    "\n",
    "else:\n",
    "\n",
    "    if cmash_import:\n",
    "        plt.scatter(source_grid[0],source_grid[1],s=20,c=source_distr_norm,vmin=0,transform=ccrs.PlateCarree(),cmap=cmap,zorder=3)\n",
    "    else:\n",
    "        plt.scatter(source_grid[0],source_grid[1],s=20,c=source_distr_norm,vmin=0,transform=ccrs.PlateCarree(),cmap=plt.get_cmap('Blues_r'),zorder=3)\n",
    "\n",
    "\n",
    "cbar = plt.colorbar(pad=0.01)\n",
    "cbar.ax.tick_params(labelsize=30) \n",
    "cbar.set_label('Power Spectral Density',rotation=270,labelpad=40,fontsize=40)\n",
    "\n",
    "plt.title(f'Noise distribution for {project_name}: iteration {iteration_nr}',fontsize=50,pad=30)\n",
    "\n",
    "plt.scatter(stat_lon,stat_lat,s=150,c='lawngreen',marker='^',edgecolor='k',linewidth=1,transform=ccrs.PlateCarree(),zorder=4)\n",
    "plt.savefig(os.path.join(plot_path,f'iteration_{iteration_nr}_1_noise_distribution.png'),bbox_inches='tight',facecolor='white')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11733a39",
   "metadata": {},
   "source": [
    "### Plot the gradient (smoothed and unsmoothed)\n",
    "\n",
    "And finally we can plot the gradient (i.e. sum of all sensitivity kernels) which is used to update the noise source distribution in each iteration. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54dfd444",
   "metadata": {},
   "outputs": [],
   "source": [
    "iteration_nr = 0\n",
    "triangulation = True\n",
    "\n",
    "sourcegrid = np.load(sourcegrid_file)\n",
    "grid_tri = tri.Triangulation(sourcegrid_voronoi[0],sourcegrid_voronoi[1])\n",
    "\n",
    "grad_all_smooth = os.path.join(output_path,f'iteration_{iteration_nr}','grad_all_smooth.npy')\n",
    "grad_all = os.path.join(output_path,f'iteration_{iteration_nr}','grad_all.npy')\n",
    "smooth_param = os.path.join(output_path,f'iteration_{iteration_nr}',f'smoothing_iter_{iteration_nr}.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd13d26c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot unsmoothed gradient\n",
    "grad = np.load(grad_all,allow_pickle=True)[0]\n",
    "\n",
    "v = np.max(np.abs(grad))\n",
    "\n",
    "plt.figure(figsize=(50,20))\n",
    "ax = plt.axes(projection=ccrs.Robinson(central_longitude=0))\n",
    "ax.set_global()\n",
    "\n",
    "if only_ocean:\n",
    "    ax.add_feature(cfeature.NaturalEarthFeature('cultural', 'admin_0_countries', '50m', edgecolor='black', facecolor=cfeature.COLORS['land']),zorder=2)\n",
    "else:\n",
    "    ax.coastlines()\n",
    "\n",
    "if triangulation:\n",
    "    plt.tripcolor(grid_tri,grad,cmap=plt.get_cmap('seismic'),linewidth=0.0,edgecolor='none',vmin=-v,vmax=v,zorder=1,transform=ccrs.PlateCarree())\n",
    "else:\n",
    "    plt.scatter(sourcegrid[0],sourcegrid[1],s=20,c=grad,transform=ccrs.PlateCarree(),cmap='seismic',vmin=-v,vmax=v,zorder=3)\n",
    "\n",
    "cbar = plt.colorbar(pad=0.01)\n",
    "cbar.formatter.set_powerlimits((0, 0))\n",
    "cbar.ax.tick_params(labelsize=30) \n",
    "cbar.set_label('Gradient',rotation=270,labelpad=40,fontsize=40)\n",
    "\n",
    "plt.title(f'Gradient for {project_name}: iteration {step}',fontsize=50, pad=30)\n",
    "\n",
    "plt.scatter(stat_lon,stat_lat,s=150,c='lawngreen',marker='^',edgecolor='k',linewidth=1,transform=ccrs.PlateCarree(),zorder=3)\n",
    "plt.savefig(os.path.join(plot_path,f'iteration_{step}_2_gradient.png'),bbox_inches='tight',facecolor='white')\n",
    "plt.show() \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a04be94",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot smoothed gradient\n",
    "grad = np.load(grad_all_smooth,allow_pickle=True)[0]\n",
    "\n",
    "v = np.max(np.abs(grad))\n",
    "\n",
    "plt.figure(figsize=(50,20))\n",
    "ax = plt.axes(projection=ccrs.Robinson(central_longitude=0))\n",
    "ax.set_global()\n",
    "\n",
    "if only_ocean:\n",
    "    ax.add_feature(cfeature.NaturalEarthFeature('cultural', 'admin_0_countries', '50m', edgecolor='black', facecolor=cfeature.COLORS['land']),zorder=2)\n",
    "else:\n",
    "    ax.coastlines()\n",
    "\n",
    "if triangulation:\n",
    "    plt.tripcolor(grid_tri,grad,cmap=plt.get_cmap('seismic'),linewidth=0.0,edgecolor='none',vmin=-v,vmax=v,zorder=1,transform=ccrs.PlateCarree())\n",
    "else:\n",
    "    plt.scatter(sourcegrid[0],sourcegrid[1],s=20,c=grad,transform=ccrs.PlateCarree(),cmap='seismic',vmin=-v,vmax=v,zorder=3)\n",
    "\n",
    "cbar = plt.colorbar(pad=0.01)\n",
    "cbar.formatter.set_powerlimits((0, 0))\n",
    "cbar.ax.tick_params(labelsize=30) \n",
    "cbar.set_label('Smoothed gradient',rotation=270,labelpad=40,fontsize=40)\n",
    "\n",
    "plt.title(f'Smoothed gradient for {project_name}: iteration {step} with {np.load(smooth_param)}° smoothing',fontsize=50, pad=30)\n",
    "\n",
    "plt.scatter(stat_lon,stat_lat,s=150,c='lawngreen',marker='^',edgecolor='k',linewidth=1,transform=ccrs.PlateCarree(),zorder=3)\n",
    "plt.savefig(os.path.join(plot_path,f'iteration_{step}_3_gradient_smooth.png'),bbox_inches='tight',facecolor='white')\n",
    "plt.show() "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:noisi_inv] *",
   "language": "python",
   "name": "conda-env-noisi_inv-py"
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}