changing defuault grid to equiangular in all SHT routines for consistency (#59)

Changelog.md

@@ -2,6 +2,10 @@
 
 ## Versioning
 
+### v0.7.3
+
+* Changing default grid in all SHT routines to `equiangular`
+
 ### v0.7.2
 
 * Added resampling modules for convenience

torch_harmonics/distributed/distributed_sht.py

@@ -53,7 +53,7 @@ class DistributedRealSHT(nn.Module):
     [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
     """
 
-    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="lobatto", norm="ortho", csphase=True):
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
         """
         Distribtued SHT layer. Expects the last 3 dimensions of the input tensor to be channels, latitude, longitude.
 
@@ -178,7 +178,7 @@ class DistributedInverseRealSHT(nn.Module):
     [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
     """
 
-    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="lobatto", norm="ortho", csphase=True):
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
 
         super().__init__()
 
@@ -288,7 +288,7 @@ class DistributedRealVectorSHT(nn.Module):
     [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
     """
 
-    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="lobatto", norm="ortho", csphase=True):
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
         """
         Initializes the vector SHT Layer, precomputing the necessary quadrature weights
 
@@ -432,7 +432,7 @@ class DistributedInverseRealVectorSHT(nn.Module):
     [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
     [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
     """
-    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="lobatto", norm="ortho", csphase=True):
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
 
         super().__init__()
 

torch_harmonics/examples/pde_sphere.py

@@ -2,7 +2,7 @@
 
 # SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
-# 
+#
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are met:
 #
@@ -43,7 +43,7 @@ class SphereSolver(nn.Module):
     - Allen-Cahn eq
     """
 
-    def __init__(self, nlat, nlon, dt, lmax=None, mmax=None, grid='legendre-gauss', radius=1.0, coeff=0.001):
+    def __init__(self, nlat, nlon, dt, lmax=None, mmax=None, grid="equiangular", radius=1.0, coeff=0.001):
         super().__init__()
 
         # time stepping param
@@ -96,7 +96,7 @@ class SphereSolver(nn.Module):
 
     def grid2spec(self, u):
         """spectral coefficients from spatial data"""
-        
+
         return self.sht(u)
 
     def spec2grid(self, uspec):
@@ -116,7 +116,7 @@ class SphereSolver(nn.Module):
             dudtspec = uspec + (1. + 2.j)*self.coeff*self.lap*uspec - (1. + 2.j)*u3spec
         else:
             NotImplementedError
-        
+
         return dudtspec
 
     def randspec(self):

torch_harmonics/examples/shallow_water_equations.py

@@ -2,7 +2,7 @@
 
 # SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
 # SPDX-License-Identifier: BSD-3-Clause
-# 
+#
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are met:
 #
@@ -43,7 +43,7 @@ class ShallowWaterSolver(nn.Module):
     SWE solver class. Interface inspired bu pyspharm and SHTns
     """
 
-    def __init__(self, nlat, nlon, dt, lmax=None, mmax=None, grid='legendre-gauss', radius=6.37122E6, \
+    def __init__(self, nlat, nlon, dt, lmax=None, mmax=None, grid="equiangular", radius=6.37122E6, \
                  omega=7.292E-5, gravity=9.80616, havg=10.e3, hamp=120.):
         super().__init__()
 
@@ -249,7 +249,7 @@ class ShallowWaterSolver(nn.Module):
         # initial geopotential
         uspec = torch.zeros(3, self.lmax, self.mmax, dtype=ctype, device=self.lap.device)
         uspec[:, :llimit, :mlimit] = torch.sqrt(torch.tensor(4 * torch.pi / llimit / (llimit+1), device=device, dtype=ctype)) * torch.randn_like(uspec[:, :llimit, :mlimit])
-        
+
         uspec[0] = self.gravity * self.hamp * uspec[0]
         uspec[0, 0, 0] += torch.sqrt(torch.tensor(4 * torch.pi, device=device, dtype=ctype)) * self.havg * self.gravity
         uspec[1:] = mach * uspec[1:] * torch.sqrt(self.gravity * self.havg) / self.radius
@@ -276,7 +276,7 @@ class ShallowWaterSolver(nn.Module):
         # uspec = torch.zeros(3, self.lmax, self.mmax, dtype=phispec.dtype, device=device)
         # uspec[0] = phispec
         # uspec[1:] = vrtdivspec
-        
+
         return torch.tril(uspec)
 
     def timestep(self, uspec: torch.Tensor, nsteps: int) -> torch.Tensor:
@@ -293,7 +293,7 @@ class ShallowWaterSolver(nn.Module):
 
         for iter in range(nsteps):
             dudtspec[inew] = self.dudtspec(uspec)
-            
+
             # update vort,div,phiv with third-order adams-bashforth.
             # forward euler, then 2nd-order adams-bashforth time steps to start.
             if iter == 0:
@@ -311,11 +311,11 @@ class ShallowWaterSolver(nn.Module):
             inew = (inew - 1) % 3
             inow = (inow - 1) % 3
             iold = (iold - 1) % 3
-        
+
         return uspec
 
     def integrate_grid(self, ugrid, dimensionless=False, polar_opt=0):
-        dlon = 2 * torch.pi / self.nlon 
+        dlon = 2 * torch.pi / self.nlon
         radius = 1 if dimensionless else self.radius
         if polar_opt > 0:
             out = torch.sum(ugrid[..., polar_opt:-polar_opt, :] * self.quad_weights[polar_opt:-polar_opt] * dlon * radius**2, dim=(-2, -1))

torch_harmonics/sht.py

@@ -48,7 +48,7 @@ class RealSHT(nn.Module):
     [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
     """
 
-    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="lobatto", norm="ortho", csphase=True):
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
         r"""
         Initializes the SHT Layer, precomputing the necessary quadrature weights
 
@@ -141,7 +141,7 @@ class InverseRealSHT(nn.Module):
     [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
     """
 
-    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="lobatto", norm="ortho", csphase=True):
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
 
         super().__init__()
 
@@ -214,7 +214,7 @@ class RealVectorSHT(nn.Module):
     [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
     """
 
-    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="lobatto", norm="ortho", csphase=True):
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
         r"""
         Initializes the vector SHT Layer, precomputing the necessary quadrature weights
 
@@ -321,7 +321,7 @@ class InverseRealVectorSHT(nn.Module):
     [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
     [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
     """
-    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="lobatto", norm="ortho", csphase=True):
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
 
         super().__init__()