Bbonev/gradient analysis notebooks (#17)

* Adding notebooks for gradient analysis

* Updated the Readme

README.md

@@ -183,14 +183,16 @@ coeffs = sht(signal)
 
 To enable scalable model-parallelism, `torch-harmonics` implements a distributed variant of the SHT located in `torch_harmonics.distributed`.
 
-Detailed usage of torch-harmonics is demonstrated in a series of notebooks:
+Detailed usage of torch-harmonics, alongside helpful analysis provided in a series of notebooks:
 
 1. [Getting started](./notebooks/getting_started.ipynb)
 2. [Quadrature](./notebooks/quadrature.ipynb)
 3. [Visualizing the spherical harmonics](./notebooks/plot_spherical_harmonics.ipynb)
-4. [Solving the Helmholtz equation](./notebooks/helmholtz.ipynb)
-5. [Solving the shallow water equations](./notebooks/shallow_water_equations.ipynb)
-6. [Training Spherical Fourier Neural Operators](./notebooks/train_sfno.ipynb)
+4. [Spectral fitting vs. SHT](./notebooks/gradient_analysis.ipynb)
+5. [Conditioning of the Gramian](./notebooks/conditioning_sht.ipynb)
+6. [Solving the Helmholtz equation](./notebooks/helmholtz.ipynb)
+7. [Solving the shallow water equations](./notebooks/shallow_water_equations.ipynb)
+8. [Training Spherical Fourier Neural Operators](./notebooks/train_sfno.ipynb)
 
 ## Contributors
 

notebooks/plotting.py

@@ -33,7 +33,15 @@ import numpy as np
 import matplotlib.pyplot as plt
 import cartopy.crs as ccrs
 
-def plot_sphere(data, fig=None, cmap="RdBu", title=None, central_latitude=20, central_longitude=20, lon=None, lat=None):
+def plot_sphere(data,
+                fig=None,
+                cmap="RdBu",
+                title=None,
+                colorbar=False,
+                central_latitude=20,
+                central_longitude=20,
+                lon=None,
+                lat=None):
     if fig == None:
         fig = plt.figure()
 
@@ -55,12 +63,20 @@ def plot_sphere(data, fig=None, cmap="RdBu", title=None, central_latitude=20, ce
     # contour data over the map.
     im = ax.pcolormesh(Lon, Lat, data, cmap=cmap, transform=ccrs.PlateCarree(), antialiased=False)
     # ax.add_feature(cartopy.feature.COASTLINE, edgecolor='white', facecolor='none', linewidth=1.5)
-    # plt.colorbar(im)
+    if colorbar:
+        plt.colorbar(im)
     plt.title(title, y=1.05)
 
     return im
 
-def plot_data(data, fig=None, projection=None, cmap="RdBu", title=None, lon=None, lat=None):
+def plot_data(data,
+              fig=None,
+              projection=None,
+              cmap="RdBu",
+              title=None,
+              colorbar=False,
+              lon=None,
+              lat=None):
     if fig == None:
         fig = plt.figure()
     
@@ -75,7 +91,8 @@ def plot_data(data, fig=None, projection=None, cmap="RdBu", title=None, lon=None
     fig = plt.figure(figsize=(10, 5))
     ax = fig.add_subplot(1, 1, 1, projection=projection)
     im = ax.pcolormesh(Lon, Lat, data, cmap=cmap)
-    # plt.colorbar(im)
+    if colorbar:
+        plt.colorbar(im)
     plt.title(title, y=1.05)
 
     ax.grid(True)