0.2.5

Improvements:
Specify einsum path to optimize the einsum computation, which increases the computation by 30%.
Added the option to input lmaxgrid in SH_deriv_store and Displacement_strains, which will control the output gridsizes for the SH derivatives and strain.

Corrections:
Minor corrections to B1986_nmax where the drho_corr to the tangential displacement was wrongly incrementally increased in the case where the density contrast is located in the surface topo or crust.
Fixed default values of base_drho and top_drho to be consistent in all functions.
Corrected a typo in SH_deriv and SH_deriv_store that caused a blow up of the derivatives at the pole.

Displacement_strain_planet/B1986_nmax.py

@@ -94,30 +94,30 @@ def corr_nmax_drho(
 
     # Density contrast in the relief correction.
     if density_var:
-        FA_lm_drho, D = pysh.gravmag.CilmPlusRhoHDH(
-            shape_grid, 1.0, mass, rho_grid, lmax=lmax
+        MS_lm_drho, D = pysh.gravmag.CilmPlusRhoHDH(
+            shape_grid, 1, mass, rho_grid, lmax=lmax
         )
-        MS_lm_drho = MS_lm_nmax
+        MS_lm_drho_cst = MS_lm_nmax
         if filter_in is not None:
             for l in range(1, lmax + 1):
+                MS_lm_drho_cst[:, l, : l + 1] /= filter_in[l]
                 MS_lm_drho[:, l, : l + 1] /= filter_in[l]
-                FA_lm_drho[:, l, : l + 1] /= filter_in[l]
         elif filter is not None:
             for l in range(1, lmax + 1):
-                MS_lm_drho[:, l, : l + 1] /= DownContFilter(
+                MS_lm_drho_cst[:, l, : l + 1] /= DownContFilter(
                     l, filter_half, R, R - c, type=filter
                 )
-                FA_lm_drho[:, l, : l + 1] /= DownContFilter(
+                MS_lm_drho[:, l, : l + 1] /= DownContFilter(
                     l, filter_half, R, R - c, type=filter
                 )
         if nmax == 1:
-            MS_lm_drho *= D ** 2
+            MS_lm_drho_cst *= D ** 2
 
         # Divide because the thin-shell code multiplies by
         # density contrast, to correct for finite-amplitude.
         # Here we also correct for density variations, so the
         # correction is already scaled by the density contrast.
-        delta_MS_FA = R * (FA_lm_drho - MS_lm_drho + FA_lm_nmax - MS_lm_nmax) / drho
+        delta_MS_FA = R * (MS_lm_drho - MS_lm_drho_cst + FA_lm_nmax - MS_lm_nmax) / drho
     else:
         delta_MS_FA = R * (FA_lm_nmax - MS_lm_nmax)
 
@@ -891,8 +891,8 @@ def Thin_shell_matrix_nmax(
     add_array=None,
     quiet=True,
     remove_equation=None,
-    base_drho=150e3,
-    top_drho=50e3,
+    base_drho=50e3,
+    top_drho=0,
     delta_max=5,
     iter_max=250,
     delta_out=500e3,
@@ -1008,9 +1008,9 @@ def Thin_shell_matrix_nmax(
     remove_equation : string, optional, default = None
         String of the equation to be removed. This must be either
         'G_lm', 'Gc_lm', 'w_lm', 'omega_lm', or 'q_lm'.
-    base_drho : float, optional, default = 150e3
+    base_drho : float, optional, default = 50e3
         Lower depth for the of the density contrast.
-    top_drho : float, optional, default = 50e3
+    top_drho : float, optional, default = 0
         Upper depth for the of the density contrast.
     delta_max : float, optional, default = 5
         The algorithm will continue to iterate until the maximum
@@ -1247,7 +1247,7 @@ def Thin_shell_matrix_nmax(
 
                 v1v = v / (1.0 - v)
                 if density_var_H:
-                    drho_corr += v1v * drhom_lm_o * g0 * Te * H_lm_o / R
+                    drho_corr = v1v * drhom_lm_o * g0 * Te * H_lm_o / R
                     drho_H = rhol
                     if drhom_lm is not None and sum_drho != 0:
                         H_drho_grid = rho_grid
@@ -1259,7 +1259,10 @@ def Thin_shell_matrix_nmax(
                         * (1.0 + (((R - c) / R) ** 3 - 1.0) * rhoc / rhobar)
                         / ((R - c) / R) ** 2
                     )
-                    drho_corr += v1v * drhom_lm_o * gmoho * (Te - c) * dc_lm_o / R
+                    if density_var_H:
+                        drho_corr += v1v * drhom_lm_o * gmoho * (Te - c) * dc_lm_o / R
+                    else:
+                        drho_corr = v1v * drhom_lm_o * gmoho * (Te - c) * dc_lm_o / R
                     drho_wdc = rhom - rhoc
                     if drhom_lm is not None and sum_drho != 0:
                         wdc_drho_grid = rhom - rho_grid

Displacement_strain_planet/Displacement_strain.py

@@ -65,7 +65,6 @@ def SH_deriv(theta, phi, lmax):
         p_theta, dp_theta = pysh.legendre.PlmBar_d1(lmax, cost)
         dp_theta *= -sint  # Derivative with respect to
         # theta.
-        p_theta /= sint
         costsint = cost / sint
         sintt = 1.0 / sint ** 2
     for l in range(lmax + 1):
@@ -94,7 +93,6 @@ def SH_deriv(theta, phi, lmax):
                 Y_lm_d2_thetaphi_a[1, l, m_abs] = dp_theta[index] * mcosmphi
                 y_lm[1, l, m_abs] = p_theta[index] * sinmphi
 
-        y_lm[:, l, : l + 1] *= sint
         if theta == 0 or theta == pi:
             Y_lm_d2_theta_a[:, l, : l + 1] = 0.0  # Not defined.
         else:
@@ -119,7 +117,7 @@ def SH_deriv(theta, phi, lmax):
 # ==== SH_deriv_store ====
 
 
-def SH_deriv_store(lmax, path, save=True, compressed=False):
+def SH_deriv_store(lmax, path, lmaxgrid=None, save=True, compressed=False):
     """
     Compute and store or load spherical harmonic derivatives
     over the entire sphere (first and second order).
@@ -146,26 +144,41 @@ def SH_deriv_store(lmax, path, save=True, compressed=False):
 
     Parameters
     ----------
-    path : string
-        Path to store or load spherical harmonic derivatives.
     lmax : int
         Maximum spherical harmonic degree to compute for the derivatives.
-        This parameter determines the resolution of the grid as
-        latitude(2 * lmax + 2) and longitude(2 * (2 * lmax + 2))
+    path : string
+        Path to store or load spherical harmonic derivatives.
+    lmaxgrid : int, optional, default = None
+        The maximum spherical harmonic degree resolvable by the grid,
+        latitude(2 * lmaxgrid + 2) and longitude(2 * (2 * lmaxgrid + 2)).
+        Should be higher or equal than lmax. If None, this parameter is set to lmax.
     save : boolean, optional, default = True
         If True, save the data at the given path location.
     compressed : boolean, optional, default = False
         If True, the data is saved in compressed .npz format instead of
         npy, which decreases the file size by about a factor 2. This is
         recommended when lmax > 75.
     """
-    n = 2 * lmax + 2
-    poly_file = "%s/Y_lmsd1d2_lmax%s.%s" % (path, lmax, "npz" if compressed else "npy")
+    if lmaxgrid is None:
+        lmaxgrid = lmax
+    elif lmaxgrid < lmax:
+        raise ValueError(
+            "lmaxgrid should be higher or equal than lmax, input is %s" % (lmaxgrid)
+            + " with lmax = %s." % (lmax)
+        )
+    n = 2 * lmaxgrid + 2
+    poly_file = "%s/Y_lmsd1d2_lmax%s_lmaxgrid%s.%s" % (
+        path,
+        lmax,
+        lmaxgrid,
+        "npz" if compressed else "npy",
+    )
 
     if Path(poly_file).exists() == 0:
         print(
             "Pre-compute SH derivatives, may take some"
-            + " time depending on lmax, input lmax is %s" % (lmax)
+            + " time depending on lmax and lmaxgrid, which are %s and %s."
+            % (lmax, lmaxgrid)
         )
         index_size = int((lmax + 1) * (lmax + 2) / 2)
         shape_save = (n, 2 * n, 2, lmax + 1, lmax + 1)
@@ -184,14 +197,13 @@ def SH_deriv_store(lmax, path, save=True, compressed=False):
             t_i += 1
             sint = np.sin(theta)
             cost = np.cos(theta)
-            if theta == 0 or theta == pi:
+            if theta == 0:
                 dp_theta = np.zeros((index_size))
                 p_theta = np.zeros((index_size))
             else:
                 p_theta, dp_theta = pysh.legendre.PlmBar_d1(lmax, cost)
                 dp_theta *= -sint  # Derivative with
                 # respect to theta.
-                p_theta /= sint
                 costsint = cost / sint
                 sintt = 1.0 / sint ** 2
 
@@ -225,10 +237,9 @@ def SH_deriv_store(lmax, path, save=True, compressed=False):
                         )
                         y_lm[:, 1, l, m_abs] = p_theta[index] * sinmphi
 
-                y_lm[:, :, l, : l + 1] *= sint
                 y_lm_save[t_i, :, :, l, : l + 1] = y_lm[:, :, l, : l + 1]
 
-                if theta == 0 or theta == pi:
+                if theta == 0:
                     Y_lm_d2_theta_a[t_i, :, :, l, : l + 1] = 0.0
                     # Not defined.
                 else:
@@ -239,6 +250,7 @@ def SH_deriv_store(lmax, path, save=True, compressed=False):
                         - Y_lm_d1_theta_a[t_i, :, :, l, : l + 1] * costsint
                         - sintt * Y_lm_d2_phi_a[t_i, :, :, l, : l + 1]
                     )
+
         if save:
             if compressed:
                 np.savez_compressed(
@@ -312,6 +324,7 @@ def Displacement_strains(
     colat_max=180,
     lon_min=0,
     lon_max=360,
+    lmaxgrid=None,
     only_deflec=False,
     Y_lm_d1_t=None,
     Y_lm_d1_p=None,
@@ -382,6 +395,10 @@ def Displacement_strains(
         Minimum longitude for grid computation of strains and stresses.
     lon_max : float, optional, default = 360
         Maximum longitude for grid computation of strains and stresses.
+    lmaxgrid : int, optional, default = None
+        The maximum spherical harmonic degree resolvable by the grid,
+        latitude(2 * lmaxgrid + 2) and longitude(2 * (2 * lmaxgrid + 2)).
+        If None, this parameter is set to lmax.
     only_deflec : bool, optional, default = False
         Output only the displacement grid for all latitude and longitudes.
     Y_lm_d1_t : array, float, size(2,lmax+1,lmax+1), optional, default = None
@@ -415,7 +432,14 @@ def Displacement_strains(
         A_lm = pysh.SHCoeffs.from_array(A_lm).pad(lmax=lmax).coeffs
         w_lm = pysh.SHCoeffs.from_array(w_lm).pad(lmax=lmax).coeffs
 
-    n = 2 * lmax + 2
+    if lmaxgrid is None:
+        lmaxgrid = lmax
+    elif lmaxgrid < lmax:
+        raise ValueError(
+            "lmaxgrid should be higher or equal than lmax, input is %s" % (lmaxgrid)
+            + " with lmax = %s." % (lmax)
+        )
+    n = 2 * lmaxgrid + 2
 
     if Y_lm_d1_p is not None:
         if quiet is False:
@@ -432,7 +456,7 @@ def Displacement_strains(
             Y_lm_d2_p,
             Y_lm_d2_tp,
             y_lm,
-        ) = SH_deriv_store(lmax, path)
+        ) = SH_deriv_store(lmax, path, lmaxgrid=lmaxgrid)
 
     # Some constants for the elastic model.
     Re = R - (0.5 * Te)
@@ -491,35 +515,40 @@ def Displacement_strains(
 
     ein_sum = "mijk,ijk->m"
     ein_sum_mul = "mijk,ijk,m->m"
-    w_deflec_ylm = R_m1 * np.einsum(ein_sum, y_lm, w_lm)
-    eps_theta[mask] = R_m1 * np.einsum(ein_sum, Y_lm_d2_t, A_lm) + w_deflec_ylm
+    path_sum = ["einsum_path", (0, 1)]  # Generated from np.einsum_path
+    path_mul = ["einsum_path", (0, 1), (0, 1)]  # Generated from np.einsum_path
+    w_deflec_ylm = R_m1 * np.einsum(ein_sum, y_lm, w_lm, optimize=path_sum)
+    eps_theta[mask] = (
+        R_m1 * np.einsum(ein_sum, Y_lm_d2_t, A_lm, optimize=path_sum) + w_deflec_ylm
+    )
     eps_phi[mask] = (
         R_m1
         * (
-            np.einsum(ein_sum_mul, Y_lm_d2_p, A_lm, csc2)
-            + np.einsum(ein_sum_mul, Y_lm_d1_t, A_lm, cot)
+            np.einsum(ein_sum_mul, Y_lm_d2_p, A_lm, csc2, optimize=path_mul)
+            + np.einsum(ein_sum_mul, Y_lm_d1_t, A_lm, cot, optimize=path_mul)
         )
         + w_deflec_ylm
     )
     omega[mask] = R_m1 * (
-        np.einsum(ein_sum_mul, Y_lm_d2_tp, A_lm, csc)
-        - np.einsum(ein_sum_mul, Y_lm_d1_p, A_lm, cotcsc)
+        np.einsum(ein_sum_mul, Y_lm_d2_tp, A_lm, csc, optimize=path_mul)
+        - np.einsum(ein_sum_mul, Y_lm_d1_p, A_lm, cotcsc, optimize=path_mul)
     )
 
     kappa_theta[mask] = (
-        n_Rm2 * np.einsum(ein_sum, Y_lm_d2_t, w_lm) + (-R_m1) * w_deflec_ylm
+        n_Rm2 * np.einsum(ein_sum, Y_lm_d2_t, w_lm, optimize=path_sum)
+        + (-R_m1) * w_deflec_ylm
     )
     kappa_phi[mask] = (
         n_Rm2
         * (
-            np.einsum(ein_sum_mul, Y_lm_d2_p, w_lm, csc2)
-            + np.einsum(ein_sum_mul, Y_lm_d1_t, w_lm, cot)
+            np.einsum(ein_sum_mul, Y_lm_d2_p, w_lm, csc2, optimize=path_mul)
+            + np.einsum(ein_sum_mul, Y_lm_d1_t, w_lm, cot, optimize=path_mul)
         )
         + (-R_m1) * w_deflec_ylm
     )
     tau[mask] = n_Rm2 * (
-        np.einsum(ein_sum_mul, Y_lm_d2_tp, w_lm, csc)
-        - np.einsum(ein_sum_mul, Y_lm_d1_p, w_lm, cotcsc)
+        np.einsum(ein_sum_mul, Y_lm_d2_tp, w_lm, csc, optimize=path_mul)
+        - np.einsum(ein_sum_mul, Y_lm_d1_p, w_lm, cotcsc, optimize=path_mul)
     )
 
     stress_theta = (