PythonicDISORT 0.9.0__py3-none-any.whl → 0.9.1__py3-none-any.whl

PythonicDISORT/_assemble_intensity_and_fluxes.py

@@ -20,12 +20,12 @@ def _assemble_intensity_and_fluxes(
     there_is_beam_source,               # Is there a beam source?
     b_pos, b_neg,                       # Dirichlet BCs
     b_pos_is_scalar, b_neg_is_scalar,   # Is each Dirichlet BCs scalar (isotropic)?
-    s_poly_coeffs,                      # Polynomial coefficients of isotropic source
     Nscoeffs,                           # Number of isotropic source polynomial coefficients
+    s_poly_coeffs,                      # Polynomial coefficients of isotropic source
     there_is_iso_source,                # Is there an isotropic source?
     scale_tau,                          # Delta-scale factor for tau
     only_flux,                          # Only compute fluxes?
-    NLayers_to_use_sparse,              # Number of layers above or equal which to use sparse matrices
+    use_banded_solver_NLayers,          # Number of layers above or equal which to use `scipy.linalg.solve_banded`
     _autograd_compatible,               # Should the output functions be compatible with autograd?
 ):  
     """Assembles the solution functions: intensity `u`, upward flux `flux_up`, downward flux `flux_down` 
@@ -66,13 +66,13 @@ def _assemble_intensity_and_fluxes(
     | `b_pos`                        | `NQuad/2 x NFourier` or scalar     |
     | `b_neg`                        | `NQuad/2 x NFourier` or scalar     |
     | `b_pos_is_scalar`              | boolean                            |
-    | `b_neg_is_scalar`              | boolean                            |              
+    | `b_neg_is_scalar`              | boolean                            |        
+    | `Nscoeffs`                     | scalar                             |    
     | `s_poly_coeffs`                | `NLayers x Nscoeffs` or `Nscoeffs` |
-    | `Nscoeffs`                     | scalar                             |
     | `there_is_iso_source`          | boolean                            |
     | `scale_tau`                    | `NLayers`                          |
     | `only_flux`                    | boolean                            |
-    | `NLayers_to_use_sparse`        | scalar                             |
+    | `use_banded_solver_NLayers`    | scalar                             |
     | `_autograd_compatible`         | boolean                            |
     
     Notable internal variables of _assemble_intensity_and_fluxes
@@ -146,10 +146,10 @@ def _assemble_intensity_and_fluxes(
         there_is_beam_source,
         b_pos, b_neg,
         b_pos_is_scalar, b_neg_is_scalar,
-        s_poly_coeffs,
         Nscoeffs,
+        s_poly_coeffs,
         there_is_iso_source,
-        NLayers_to_use_sparse,
+        use_banded_solver_NLayers,
     )
     
     G_collect_0 = G_collect[0, :, :, :]
@@ -198,26 +198,28 @@ def _assemble_intensity_and_fluxes(
             
             if is_antiderivative_wrt_tau:
                 if there_is_beam_source:
-                    um = np.einsum(
-                        "mtij, mtj -> mti",
-                        GC_collect[:, l, :, :],
-                        np.exp(exponent) / (scale_tau[None, l, None] * K_collect[:, l, :]),
-                        optimize=True,
-                    ) + B_collect[:, l, :] * np.exp(-scaled_tau[None, :, None] / mu0) / (
-                        -scale_tau[None, l, None] / mu0
+                    um = (
+                        np.einsum(
+                            "mtij, mtj -> mti",
+                            GC_collect[:, l, :, :],
+                            np.exp(exponent) / (scale_tau[None, :, None] * K_collect)[:, l, :],
+                            optimize=True,
+                        )
+                        + (B_collect / (-scale_tau / mu0)[None, :, None])[:, l, :]
+                        * np.exp(-scaled_tau / mu0)[None, :, None]
                     )
                 else:
                     um = np.einsum(
                         "mtij, mtj -> mti",
                         GC_collect[:, l, :, :],
-                        np.exp(exponent) / (scale_tau[None, l, None] * K_collect[:, l, :]),
+                        np.exp(exponent) / (scale_tau[None, l, None] * K_collect)[:, l, :],
                         optimize=True,
                     )
             else:
                 if there_is_beam_source:
                     um = np.einsum(
                         "mtij, mtj -> mti", GC_collect[:, l, :, :], np.exp(exponent), optimize=True
-                    ) + B_collect[:, l, :] * np.exp(-scaled_tau[None, :, None] / mu0)
+                    ) + B_collect[:, l, :] * np.exp(-scaled_tau / mu0)[None, :, None]
                 else:
                     um = np.einsum(
                         "mtij, mtj -> mti", GC_collect[:, l, :, :], np.exp(exponent), optimize=True
@@ -225,13 +227,13 @@ def _assemble_intensity_and_fluxes(
             
             # Contribution from particular solution for isotropic internal sources
             if there_is_iso_source:
+                l_uniq, l_inv = np.unique(l, return_inverse=True)
                 _mathscr_v_contribution = _mathscr_v(
-                    tau, l,
-                    s_poly_coeffs,
-                    Nscoeffs,
-                    G_collect_0,
-                    K_collect_0,
-                    G_inv_collect_0,
+                    tau, l_inv, Nscoeffs,
+                    s_poly_coeffs[l_uniq],
+                    G_collect_0[l_uniq],
+                    K_collect_0[l_uniq],
+                    G_inv_collect_0[l_uniq],
                     mu_arr,
                     is_antiderivative_wrt_tau,
                     _autograd_compatible
@@ -262,26 +264,28 @@ def _assemble_intensity_and_fluxes(
                 )
                 if is_antiderivative_wrt_tau:
                     if there_is_beam_source:
-                        ulast = np.einsum(
-                            "tij, tj -> it",
-                            GC_collect[-1, l, :, :],
-                            np.exp(exponent) / (scale_tau[l, None] * K_collect[-1, l, :]),
-                            optimize=True,
-                        ) + B_collect[-1, l, :].T * np.exp(-scaled_tau[None, :] / mu0) / (
-                            -scale_tau[None, l] / mu0
+                        ulast = (
+                            np.einsum(
+                                "tij, tj -> it",
+                                GC_collect[-1, l, :, :],
+                                np.exp(exponent) / (scale_tau[:, None] * K_collect[-1, :, :])[l, :],
+                                optimize=True,
+                            )
+                            + (B_collect.T[-1, :, :] / (-scale_tau / mu0)[None, :])[:, l]
+                            * np.exp(-scaled_tau / mu0)[None, :]
                         )
                     else:
                         ulast = np.einsum(
                             "tij, tj -> it",
                             GC_collect[-1, l, :, :],
-                            np.exp(exponent) / (scale_tau[l, None] * K_collect[-1, l, :]),
+                            np.exp(exponent) / (scale_tau[:, None] * K_collect[-1, :, :])[l, :],
                             optimize=True,
                         )
                 else:
                     if there_is_beam_source:
                         ulast = np.einsum(
                             "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
-                        ) + B_collect[-1, l, :].T * np.exp(-scaled_tau[None, :] / mu0)
+                        ) + B_collect.T[-1, :, l] * np.exp(-scaled_tau / mu0)[None, :]
                     else:
                         ulast = np.einsum(
                             "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
@@ -333,7 +337,10 @@ def _assemble_intensity_and_fluxes(
             # in performing the reclassification of delta-scaled actinic flux
             if _return_act_dscale_for_reclass: 
                 if is_antiderivative_wrt_tau:
-                    act_dscale_reclassification = I0 * np.exp(-scaled_tau / mu0) / (-scale_tau[l] / mu0) - I0 * np.exp(-tau / mu0) * -mu0
+                    act_dscale_reclassification = (
+                        I0 * np.exp(-scaled_tau / mu0) / (-scale_tau[l] / mu0)
+                        - I0 * np.exp(-tau / mu0) * -mu0
+                    )
                 else:
                     act_dscale_reclassification = I0 * np.exp(-scaled_tau / mu0) - I0 * np.exp(-tau / mu0)
         else:
@@ -351,26 +358,28 @@ def _assemble_intensity_and_fluxes(
         
         if is_antiderivative_wrt_tau:
             if there_is_beam_source:
-                u0 = np.einsum(
-                    "tij, tj -> it",
-                    GC_collect_0[l, :, :],
-                    np.exp(exponent) / (scale_tau[l, None] * K_collect_0[l, :]),
-                    optimize=True,
-                ) + B_collect_0[l, :].T * np.exp(-scaled_tau[None, :] / mu0) / (
-                    -scale_tau[None, l] / mu0
+                u0 = (
+                    np.einsum(
+                        "tij, tj -> it",
+                        GC_collect_0[l, :, :],
+                        np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :],
+                        optimize=True,
+                    )
+                    + (B_collect_0.T / (-scale_tau / mu0)[None, :])[:, l]
+                    * np.exp(-scaled_tau / mu0)[None, :]
                 )
             else:
                 u0 = np.einsum(
                     "tij, tj -> it",
                     GC_collect_0[l, :, :],
-                    np.exp(exponent) / (scale_tau[l, None] * K_collect_0[l, :]),
+                    np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :],
                     optimize=True,
                 )
         else:
             if there_is_beam_source:
                 u0 = np.einsum(
                     "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
-                ) + B_collect_0[l, :].T * np.exp(-scaled_tau[None, :] / mu0)
+                ) + B_collect_0.T[:, l] * np.exp(-scaled_tau / mu0)[None, :]
             else:
                 u0 = np.einsum(
                     "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
@@ -378,13 +387,13 @@ def _assemble_intensity_and_fluxes(
         
         # Contribution from particular solution for isotropic internal sources
         if there_is_iso_source:
+            l_uniq, l_inv = np.unique(l, return_inverse=True)
             _mathscr_v_contribution = _mathscr_v(
-                tau, l,
-                s_poly_coeffs,
-                Nscoeffs,
-                G_collect_0,
-                K_collect_0,
-                G_inv_collect_0,
+                tau, l_inv, Nscoeffs,
+                s_poly_coeffs[l_uniq],
+                G_collect_0[l_uniq],
+                K_collect_0[l_uniq],
+                G_inv_collect_0[l_uniq],
                 mu_arr,
                 is_antiderivative_wrt_tau,
                 _autograd_compatible
@@ -431,27 +440,27 @@ def _assemble_intensity_and_fluxes(
             scaled_tau = tau
 
         if there_is_iso_source:
+            l_uniq, l_inv = np.unique(l, return_inverse=True)
             _mathscr_v_contribution = _mathscr_v(
-                tau, l,
-                s_poly_coeffs,
-                Nscoeffs,
-                G_collect_0,
-                K_collect_0,
-                G_inv_collect_0,
+                tau, l_inv, Nscoeffs,
+                s_poly_coeffs[l_uniq],
+                G_collect_0[l_uniq, :N, :],
+                K_collect_0[l_uniq],
+                G_inv_collect_0[l_uniq],
                 mu_arr,
                 is_antiderivative_wrt_tau,
                 _autograd_compatible
-            )[:N, :]
+            )
         else:
             _mathscr_v_contribution = 0
         
         if there_is_beam_source:
             if is_antiderivative_wrt_tau:
                 direct_beam_contribution = (
-                    B_pos[:, l] * np.exp(-scaled_tau[None, :] / mu0) / (-scale_tau[None, l] / mu0)
+                    (B_pos / (-scale_tau / mu0)[None, :])[:, l] * np.exp(-scaled_tau / mu0)[None, :] 
                 )
             else:
-                direct_beam_contribution = B_pos[:, l] * np.exp(-scaled_tau[None, :] / mu0)
+                direct_beam_contribution = B_pos[:, l] * np.exp(-scaled_tau / mu0)[None, :]
         else:
             direct_beam_contribution = 0
 
@@ -468,7 +477,7 @@ def _assemble_intensity_and_fluxes(
                 np.einsum(
                     "tij, tj -> it",
                     GC_pos[l, :, :],
-                    np.exp(exponent) / (scale_tau[l, None] * K_collect_0[l, :]),
+                    np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :],
                 )
                 + direct_beam_contribution
                 + _mathscr_v_contribution
@@ -510,27 +519,27 @@ def _assemble_intensity_and_fluxes(
             scaled_tau = tau
 
         if there_is_iso_source:
+            l_uniq, l_inv = np.unique(l, return_inverse=True)
             _mathscr_v_contribution = _mathscr_v(
-                tau, l,
-                s_poly_coeffs,
-                Nscoeffs,
-                G_collect_0,
-                K_collect_0,
-                G_inv_collect_0,
+                tau, l_inv, Nscoeffs,
+                s_poly_coeffs[l_uniq],
+                G_collect_0[l_uniq, N:, :],
+                K_collect_0[l_uniq],
+                G_inv_collect_0[l_uniq],
                 mu_arr,
                 is_antiderivative_wrt_tau,
                 _autograd_compatible
-            )[N:, :]
+            )
         else:
             _mathscr_v_contribution = 0
         
         if there_is_beam_source:
             if is_antiderivative_wrt_tau:
-                direct_beam_contribution = B_neg[:, l] * np.exp(-scaled_tau[None, :] / mu0) / (-scale_tau[None, l] / mu0)
+                direct_beam_contribution = (B_neg / (-scale_tau / mu0)[None, l])[:, l] * np.exp(-scaled_tau / mu0)[None, :]
                 direct_beam = I0 * mu0 * np.exp(-tau / mu0) * -mu0
-                direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0) / (-scale_tau[None, l] / mu0)
+                direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0) / (-scale_tau / mu0)[None, l]
             else:
-                direct_beam_contribution = B_neg[:, l] * np.exp(-scaled_tau[None, :] / mu0)
+                direct_beam_contribution = B_neg[:, l] * np.exp(-scaled_tau / mu0)[None, :]
                 direct_beam = I0 * mu0 * np.exp(-tau / mu0)
                 direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0)
         else:
@@ -548,7 +557,7 @@ def _assemble_intensity_and_fluxes(
         
         if is_antiderivative_wrt_tau:
             u0_neg = (
-                np.einsum("tij, tj -> it", GC_neg[l, :, :], np.exp(exponent) / (scale_tau[l, None] * K_collect_0[l, :]))
+                np.einsum("tij, tj -> it", GC_neg[l, :, :], np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :])
                 + direct_beam_contribution
                 + _mathscr_v_contribution
             )

PythonicDISORT/_solve_for_coeffs.py

@@ -1,5 +1,5 @@
 from PythonicDISORT.subroutines import _mathscr_v
-from PythonicDISORT.subroutines import _nd_slice_to_indexes
+from PythonicDISORT.subroutines import to_diag_ordered_form
 
 import numpy as np
 import scipy as sc
@@ -23,10 +23,10 @@ def _solve_for_coeffs(
     there_is_beam_source,                   # Is there a beam source?
     b_pos, b_neg,                           # Dirichlet BCs
     b_pos_is_scalar, b_neg_is_scalar,       # Is each Dirichlet BCs scalar (isotropic)?
-    s_poly_coeffs,                          # Polynomial coefficients of isotropic source           
     Nscoeffs,                               # Number of isotropic source polynomial coefficients
+    s_poly_coeffs,                          # Polynomial coefficients of isotropic source           
     there_is_iso_source,                    # Is there an isotropic source?
-    NLayers_to_use_sparse,                  # Number of layers above or equal which to use sparse matrices
+    use_banded_solver_NLayers,              # Number of layers above or equal which to use `scipy.linalg.solve_banded`
 ):
     """
     Uses the boundary conditions to solve for the unknown coefficients 
@@ -62,11 +62,11 @@ def _solve_for_coeffs(
     | `b_pos`                        | `NQuad/2 x NFourier` or scalar         |
     | `b_neg`                        | `NQuad/2 x NFourier` or scalar         |
     | `b_pos_is_scalar`              | boolean                                |
-    | `b_neg_is_scalar`              | boolean                                |              
-    | `s_poly_coeffs`                | `NLayers x Nscoeffs` or `Nscoeffs`     |
+    | `b_neg_is_scalar`              | boolean                                |
     | `Nscoeffs`                     | scalar                                 |
+    | `s_poly_coeffs`                | `NLayers x Nscoeffs` or `Nscoeffs`     |
     | `there_is_iso_source`          | boolean                                |
-    | `NLayers_to_use_sparse`        | scalar                                 |
+    | `use_banded_solver_NLayers`    | scalar                                 |
     
     Notable internal variables of _solve_for_coeffs
     |   Variable   |                 Shape                |
@@ -77,13 +77,16 @@ def _solve_for_coeffs(
     ################################## Solve for coefficients of homogeneous solution ##########################################
     ############# Refer to Section 3.6.2 (single-layer) and 4 (multi-layer) of the Comprehensive Documentation #################
     
+    dim = NLayers * NQuad
     GC_collect = np.empty((NFourier, NLayers, NQuad, NQuad))
-    use_sparse_framework = NLayers >= NLayers_to_use_sparse
+    use_banded_solver = (NLayers >= use_banded_solver_NLayers)
+    if use_banded_solver:
+        sym_offset = 3 * N - 1
     
     # The following loops can easily be parallelized, but the speed-up is unlikely to be worth the overhead
     for m in range(NFourier):
         m_equals_0 = (m == 0)
-        BDRF_bool = m < NBDRF
+        BDRF_bool = (NBDRF > m) 
         
         G_collect_m = G_collect[m, :, :, :]
         K_collect_m = K_collect[m, :, :]
@@ -122,12 +125,12 @@ def _solve_for_coeffs(
             b_neg_m = b_neg[:, m]
         
         # _mathscr_v_contribution
-        if there_is_iso_source and m_equals_0:
+        if m_equals_0 and there_is_iso_source:
             _mathscr_v_contribution_top = -_mathscr_v(
                         np.array([0]), 
                         np.array([0]),
-                        s_poly_coeffs[0:1, :],
                         Nscoeffs,
+                        s_poly_coeffs[0:1, :],
                         G_collect_m[0:1, N:, :],
                         K_collect_m[0:1, :],
                         G_inv_collect_0[0:1, :, :],
@@ -140,22 +143,22 @@ def _solve_for_coeffs(
                 _mathscr_v_contribution_middle = (
                     _mathscr_v(
                         tau_arr[:-1],
-                        indices + 1,
-                        s_poly_coeffs,
+                        indices,
                         Nscoeffs,
-                        G_collect_m,
-                        K_collect_m,
-                        G_inv_collect_0,
+                        s_poly_coeffs[indices + 1],
+                        G_collect_m[indices + 1],
+                        K_collect_m[indices + 1],
+                        G_inv_collect_0[indices + 1],
                         mu_arr,
                     )
                     - _mathscr_v(
                         tau_arr[:-1],
                         indices,
-                        s_poly_coeffs,
                         Nscoeffs,
-                        G_collect_m,
-                        K_collect_m,
-                        G_inv_collect_0,
+                        s_poly_coeffs[indices],
+                        G_collect_m[indices],
+                        K_collect_m[indices],
+                        G_inv_collect_0[indices],
                         mu_arr,
                     )
                 ).ravel(order='F')
@@ -163,8 +166,8 @@ def _solve_for_coeffs(
             _mathscr_v_contribution_bottom = -_mathscr_v(
                     tau_arr[-1:], 
                     np.array([0]),
-                    s_poly_coeffs[-1:, :],
                     Nscoeffs,
+                    s_poly_coeffs[-1:, :],
                     G_collect_m[-1:, :N, :],
                     K_collect_m[-1:, :],
                     G_inv_collect_0[-1:, :, :],
@@ -177,8 +180,8 @@ def _solve_for_coeffs(
                     @ _mathscr_v(
                         tau_arr[-1:],
                         np.array([0]),
-                        s_poly_coeffs[-1:, :],
                         Nscoeffs,
+                        s_poly_coeffs[-1:, :],
                         G_collect_m[-1:, N:, :],
                         K_collect_m[-1:, :],
                         G_inv_collect_0[-1:, :, :],
@@ -197,41 +200,45 @@ def _solve_for_coeffs(
             _mathscr_v_contribution = 0
         
         if there_is_beam_source:
-            if BDRF_bool:
-                BDRF_RHS_contribution = mathscr_X_pos + R @ B_collect_m[-1, N:]
-            else:
-                BDRF_RHS_contribution = 0
-            
             RHS_middle = np.array([])
             if is_atmos_multilayered:
-                RHS_middle = np.array(
-                    [
-                        (B_collect_m[l + 1, :] - B_collect_m[l, :])
-                        * np.exp(-mu0 * scaled_tau_arr_with_0[l + 1])
-                        for l in range(NLayers - 1)
-                    ]
+                l_range = np.arange(1, NLayers)
+                RHS_middle = (
+                    (B_collect_m[l_range, :] - B_collect_m[l_range - 1, :])
+                    * np.exp(-mu0 * scaled_tau_arr_with_0)[l_range, None]
                 ).ravel()
-            
-            RHS = (
-                np.concatenate(
-                    [
-                        b_neg_m - B_collect_m[0, N:],
-                        RHS_middle,
-                        b_pos_m
-                        + (BDRF_RHS_contribution - B_collect_m[-1, :N])
-                        * np.exp(-scaled_tau_arr_with_0[-1] / mu0),
-                    ]
+                
+            if BDRF_bool:
+                RHS = (
+                    np.concatenate(
+                        [
+                            b_neg_m - B_collect_m[0, N:],
+                            RHS_middle,
+                            b_pos_m
+                            + (mathscr_X_pos + R @ B_collect_m[-1, N:] - B_collect_m[-1, :N])
+                            * np.exp(-scaled_tau_arr_with_0[-1] / mu0),
+                        ]
+                    )
+                    + _mathscr_v_contribution
+                )
+            else:
+                RHS = (
+                    np.concatenate(
+                        [
+                            b_neg_m - B_collect_m[0, N:],
+                            RHS_middle,
+                            b_pos_m
+                            - B_collect_m[-1, :N] * np.exp(-scaled_tau_arr_with_0[-1] / mu0),
+                        ]
+                    )
+                    + _mathscr_v_contribution
                 )
-                + _mathscr_v_contribution
-            )
         else:
             RHS_middle = np.zeros(NQuad * (NLayers - 1))
             RHS = np.concatenate([b_neg_m, RHS_middle, b_pos_m]) + _mathscr_v_contribution
         # --------------------------------------------------------------------------------------------------------------------------
         
         # Assemble LHS (much of this code is replicated in Section 4 of the Comprehensive Documentation)
-        dim = NLayers * NQuad
-
         G_0_nn = G_collect_m[0, N:, :N]
         G_0_np = G_collect_m[0, N:, N:]
         G_L_pn = G_collect_m[-1, :N, :N]
@@ -241,147 +248,62 @@ def _solve_for_coeffs(
         E_Lm1L = np.exp(
             K_collect_m[-1, :N] * (scaled_tau_arr_with_0[-1] - scaled_tau_arr_with_0[-2])
         )
-        if m < NBDRF:
-            BDRF_LHS_contribution_neg = R @ G_L_nn
-            BDRF_LHS_contribution_pos = R @ G_L_np
-        else:
-            BDRF_LHS_contribution_neg = 0
-            BDRF_LHS_contribution_pos = 0
         
-        if use_sparse_framework:
-            N_sq = N**2
-            density = 4 * N_sq * (2 * NLayers - 1)
-            row_indices = np.empty(density)
-            col_indices = np.empty(density)
-            block_rows = np.empty(density)
+        LHS = np.zeros((dim, dim))
+        # BCs for the entire atmosphere
+        LHS[:N, :N] = G_0_nn
+        LHS[:N, N : NQuad] = (
+            G_0_np
+            * np.exp(K_collect_m[0, :N] * scaled_tau_arr_with_0[1])[None, :]
+        )
+        if BDRF_bool:
+            LHS[-N:, -NQuad : -N] = (G_L_pn - R @ G_L_nn) * E_Lm1L[None, :]
+            LHS[-N:, -N:] = G_L_pp - R @ G_L_np
+        else:
+            LHS[-N:, -NQuad : -N] = G_L_pn * E_Lm1L[None, :]
+            LHS[-N:, -N:] = G_L_pp
+
+        # Interlayer / continuity BCs
+        for l in range(NLayers - 1):
+            G_l_pn = G_collect_m[l, :N, :N]
+            G_l_nn = G_collect_m[l, N:, :N]
+            G_l_ap = G_collect_m[l, :, N:]
+            G_lp1_an = G_collect_m[l + 1, :, :N]
+            G_lp1_pp = G_collect_m[l + 1, :N, N:]
+            G_lp1_np = G_collect_m[l + 1, N:, N:]
+            scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
+            scaled_tau_arr_l = scaled_tau_arr_with_0[l + 1]
+            scaled_tau_arr_lp1 = scaled_tau_arr_with_0[l + 2]
+            # Postive eigenvalues
+            K_l_pos = K_collect_m[l, N:]
+            K_lp1_pos = K_collect_m[l + 1, N:]
+            E_lm1l = np.exp(K_l_pos * (scaled_tau_arr_lm1 - scaled_tau_arr_l))
+            E_llp1 = np.exp(K_lp1_pos * (scaled_tau_arr_l - scaled_tau_arr_lp1))
             
-            block_row_len = 2 * NQuad**2
-            block_row_ind = np.arange(block_row_len).reshape(NQuad, 2 * NQuad)
-            block_row_sort = np.empty(block_row_len, dtype='int')
+            start_row = N + l * NQuad
+            start_col = l * NQuad
+            LHS[start_row : N + start_row, start_col : N + start_col] = G_l_pn * E_lm1l[None, :]
+            LHS[N + start_row : 2 * N + start_row, start_col : N + start_col] = G_l_nn * E_lm1l[None, :]
+            LHS[start_row : 2 * N + start_row, N + start_col : 2 * N + start_col] = G_l_ap
+            LHS[start_row : 2 * N + start_row, 2 * N + start_col : 3 * N + start_col] = -G_lp1_an
+            LHS[start_row : N + start_row, 3 * N + start_col : 4 * N + start_col] = -G_lp1_pp * E_llp1[None, :]
+            LHS[N + start_row : 2 * N + start_row, 3 * N + start_col : 4 * N + start_col] = -G_lp1_np * E_llp1[None, :]
             
-            block_row_sort[:N_sq] = block_row_ind[:N, :N].ravel()
-            block_row_sort[N_sq : 2 * N_sq] = block_row_ind[N : 2 * N, :N].ravel()
-            block_row_sort[2 * N_sq : 4 * N_sq] = block_row_ind[:2 * N, N : 2 * N].ravel()
-            block_row_sort[4 * N_sq : 6 * N_sq] = block_row_ind[:2 * N, 2 * N : 3 * N].ravel()
-            block_row_sort[6 * N_sq : 7 * N_sq] = block_row_ind[:N, 3 * N : 4 * N].ravel()
-            block_row_sort[7 * N_sq : 8 * N_sq] = block_row_ind[N : 2 * N, 3 * N : 4 * N].ravel()
-
-            # BCs for the entire atmosphere
-            row_indices[:N_sq], col_indices[:N_sq] = _nd_slice_to_indexes(np.s_[:N, :N])
-            block_rows[:N_sq] = G_0_nn.ravel()
-            (
-                row_indices[N_sq : 2 * N_sq],
-                col_indices[N_sq : 2 * N_sq],
-            ) = _nd_slice_to_indexes(np.s_[:N, N:NQuad])
-            block_rows[N_sq : 2 * N_sq] = (
-                G_0_np * np.exp(K_collect_m[0, :N] * scaled_tau_arr_with_0[1])[None, :]
-            ).ravel()
-            (
-                row_indices[2 * N_sq : 3 * N_sq],
-                col_indices[2 * N_sq : 3 * N_sq],
-            ) = _nd_slice_to_indexes(np.s_[dim - N : dim, dim - NQuad : dim - N])
-            block_rows[2 * N_sq : 3 * N_sq] = (
-                (G_L_pn - BDRF_LHS_contribution_neg) * E_Lm1L[None, :]
-            ).ravel()
-            (
-                row_indices[3 * N_sq : 4 * N_sq],
-                col_indices[3 * N_sq : 4 * N_sq],
-            ) = _nd_slice_to_indexes(np.s_[dim - N : dim, dim - N : dim])
-            block_rows[3 * N_sq : 4 * N_sq] = (
-                G_L_pp - BDRF_LHS_contribution_pos
-            ).ravel()
-
-            # Interlayer / continuity BCs
-            for l in range(NLayers - 1):
-                G_l_pn = G_collect_m[l, :N, :N]
-                G_l_nn = G_collect_m[l, N:, :N]
-                G_l_ap = G_collect_m[l, :, N:]
-                G_lp1_an = G_collect_m[l + 1, :, :N]
-                G_lp1_pp = G_collect_m[l + 1, :N, N:]
-                G_lp1_np = G_collect_m[l + 1, N:, N:]
-                scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
-                scaled_tau_arr_l = scaled_tau_arr_with_0[l + 1]
-                scaled_tau_arr_lp1 = scaled_tau_arr_with_0[l + 2]
-                # Postive eigenvalues
-                K_l_pos = K_collect_m[l, N:]
-                K_lp1_pos = K_collect_m[l + 1, N:]
-                E_lm1l = np.exp(K_l_pos * (scaled_tau_arr_lm1 - scaled_tau_arr_l))
-                E_llp1 = np.exp(K_lp1_pos * (scaled_tau_arr_l - scaled_tau_arr_lp1))
-                
-                row_indices_l, col_indices_l = _nd_slice_to_indexes(
-                    np.s_[N + l * NQuad : N + (l + 1) * NQuad, l * NQuad : (l + 2) * NQuad]
-                )
-                row_indices[
-                    4 * N_sq + l * block_row_len : 4 * N_sq + (l + 1) * block_row_len
-                ] = row_indices_l[block_row_sort]
-                col_indices[
-                    4 * N_sq + l * block_row_len : 4 * N_sq + (l + 1) * block_row_len
-                ] = col_indices_l[block_row_sort]
-                
-                start_ind = 4 * N_sq + l * block_row_len
-                block_rows[start_ind : N_sq + start_ind] = (G_l_pn * E_lm1l[None, :]).ravel()
-                block_rows[start_ind + N_sq : 2 * N_sq + start_ind] = (G_l_nn * E_lm1l[None, :]).ravel()
-                block_rows[start_ind + 2 * N_sq : 4 * N_sq + start_ind] = G_l_ap.ravel()
-                block_rows[start_ind + 4 * N_sq : 6 * N_sq + start_ind] = -G_lp1_an.ravel()
-                block_rows[start_ind + 6 * N_sq : 7 * N_sq + start_ind] = -(G_lp1_pp * E_llp1[None, :]).ravel()
-                block_rows[start_ind + 7 * N_sq : 8 * N_sq + start_ind] = -(G_lp1_np * E_llp1[None, :]).ravel()
-    
-            LHS = sc.sparse.coo_matrix(
-                (
-                    block_rows,
-                    (
-                        row_indices,
-                        col_indices,
-                    ),
-                ),
-                shape=(dim, dim),
-            )
-        
-            # Solve the system
-            C_m = sc.sparse.linalg.spsolve(LHS.tocsr(), RHS).reshape(NLayers, NQuad)
+        # --------------------------------------------------------------------------------------------------------------------------
         
-        else:            
-            LHS = np.zeros((dim, dim))
-
-            # BCs for the entire atmosphere
-            LHS[:N, :N] = G_0_nn
-            LHS[:N, N : NQuad] = (
-                G_0_np
-                * np.exp(K_collect_m[0, :N] * scaled_tau_arr_with_0[1])[None, :]
+        # Solve the system
+        if use_banded_solver:
+            C_m = sc.linalg.solve_banded(
+                (sym_offset, sym_offset),
+                to_diag_ordered_form(LHS, sym_offset),
+                RHS,
+                True,
+                True,
+                False,
             )
-            LHS[-N:, -NQuad : -N] = (G_L_pn - BDRF_LHS_contribution_neg) * E_Lm1L[None, :]
-            LHS[-N:, -N:] = G_L_pp - BDRF_LHS_contribution_pos
-
-            # Interlayer / continuity BCs
-            for l in range(NLayers - 1):
-                G_l_pn = G_collect_m[l, :N, :N]
-                G_l_nn = G_collect_m[l, N:, :N]
-                G_l_ap = G_collect_m[l, :, N:]
-                G_lp1_an = G_collect_m[l + 1, :, :N]
-                G_lp1_pp = G_collect_m[l + 1, :N, N:]
-                G_lp1_np = G_collect_m[l + 1, N:, N:]
-                scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
-                scaled_tau_arr_l = scaled_tau_arr_with_0[l + 1]
-                scaled_tau_arr_lp1 = scaled_tau_arr_with_0[l + 2]
-                # Postive eigenvalues
-                K_l_pos = K_collect_m[l, N:]
-                K_lp1_pos = K_collect_m[l + 1, N:]
-                E_lm1l = np.exp(K_l_pos * (scaled_tau_arr_lm1 - scaled_tau_arr_l))
-                E_llp1 = np.exp(K_lp1_pos * (scaled_tau_arr_l - scaled_tau_arr_lp1))
-                
-                start_row = N + l * NQuad
-                start_col = l * NQuad
-                LHS[start_row : N + start_row, start_col : N + start_col] = G_l_pn * E_lm1l[None, :]
-                LHS[N + start_row : 2 * N + start_row, start_col : N + start_col] = G_l_nn * E_lm1l[None, :]
-                LHS[start_row : 2 * N + start_row, N + start_col : 2 * N + start_col] = G_l_ap
-                LHS[start_row : 2 * N + start_row, 2 * N + start_col : 3 * N + start_col] = -G_lp1_an
-                LHS[start_row : N + start_row, 3 * N + start_col : 4 * N + start_col] = -G_lp1_pp * E_llp1[None, :]
-                LHS[N + start_row : 2 * N + start_row, 3 * N + start_col : 4 * N + start_col] = -G_lp1_np * E_llp1[None, :]
-            
-            # Solve the system
-            C_m = np.linalg.solve(LHS, RHS).reshape(NLayers, NQuad)
-        # --------------------------------------------------------------------------------------------------------------------------
+        else:
+            C_m = np.linalg.solve(LHS, RHS)
 
-        GC_collect[m, :, :, :] = G_collect_m * C_m[:, None, :]
+        GC_collect[m, :, :, :] = G_collect_m * C_m.reshape(NLayers, NQuad)[:, None, :]
         
     return GC_collect

PythonicDISORT/_solve_for_gen_and_part_sols.py

@@ -53,7 +53,6 @@ def _solve_for_gen_and_part_sols(
     | `K_collect`         | `NFourier*NLayers x NQuad`             | Reshaped to NFourier x NLayers x NQuad
     | `alpha_arr`         | `NFourier*NLayers x NQuad/2 x NQuad/2` | Reshaped to NFourier x NLayers x NQuad/2 x NQuad/2
     | `beta_arr`          | `NFourier*NLayers x NQuad/2 x NQuad/2` | Reshaped to NFourier x NLayers x NQuad/2 x NQuad/2
-    | `X_tilde_arr`       | `NFourier*NLayers x NQuad`             | Reshaped to NFourier x NLayers x NQuad
     | `B_collect`         | `NFourier*NLayers x NQuad` or `None`   | Reshaped to NFourier x NLayers x NQuad
     | `eigenvecs_GpG_arr` | `NFourier*NLayers x NQuad/2 x NQuad    |
     | `eigenvecs_GmG_arr` | `NFourier*NLayers x NQuad/2 x NQuad    |
@@ -72,7 +71,6 @@ def _solve_for_gen_and_part_sols(
     K_collect = np.empty((NFourier * NLayers, NQuad))
     alpha_arr = np.empty((NFourier * NLayers, N, N))
     beta_arr = np.empty((NFourier * NLayers, N, N))
-    X_tilde_arr = np.empty((NFourier * NLayers, NQuad))
     no_shortcut_indices = []
     no_shortcut_indices_0 = []
     
@@ -103,21 +101,35 @@ def _solve_for_gen_and_part_sols(
         # Loop over NLayers atmospheric layers
         # --------------------------------------------------------------------------------------------------------------------------          
         for l in range(NLayers):
-            # More setup
             weighted_asso_Leg_coeffs_l = weighted_scaled_Leg_coeffs[l, :][ells] * fac
-            scaled_omega_l = scaled_omega_arr[l]
             
-            # We take precautions against overflow and underflow (this shouldn't happen though)
-            if np.any(weighted_asso_Leg_coeffs_l > 0) and np.all(
-                np.isfinite(asso_leg_term_pos)
+            if np.any(weighted_asso_Leg_coeffs_l > 0) and np.all( # Due to `fac` or the coeffs being zero
+                np.isfinite(asso_leg_term_pos) # We take precautions against overflow and underflow (this shouldn't happen though)
             ):  
-                # Generate D and X (phase function terms)
+                scaled_omega_l = scaled_omega_arr[l]
+                
+                # Generate D
                 # --------------------------------------------------------------------------------------------------------------------------
                 D_temp = weighted_asso_Leg_coeffs_l[None, :] * asso_leg_term_pos.T
                 D_pos = (scaled_omega_l / 2) * D_temp @ asso_leg_term_pos
                 D_neg = (scaled_omega_l / 2) * D_temp @ asso_leg_term_neg
+                
+                # --------------------------------------------------------------------------------------------------------------------------
+                
 
+                # Assemble the coefficient matrix and additional terms
+                # --------------------------------------------------------------------------------------------------------------------------
+                DW = D_pos * W[None, :]
+                np.fill_diagonal(DW, np.diag(DW) - 1)
+                alpha = M_inv[:, None] * DW
+                beta = M_inv[:, None] * D_neg * W[None, :]
+                
+                # --------------------------------------------------------------------------------------------------------------------------
+                
+                # Particular solution for the direct beam source (refer to Section 3.6.1 of the Comprehensive Documentation)
+                # --------------------------------------------------------------------------------------------------------------------------
                 if there_is_beam_source:
+                    # Generate X
                     X_temp = (
                         (scaled_omega_l * I0 * (2 - (m == 0)) / (4 * pi))
                         * weighted_asso_Leg_coeffs_l
@@ -125,32 +137,38 @@ def _solve_for_gen_and_part_sols(
                     )
                     X_pos = X_temp @ asso_leg_term_pos
                     X_neg = X_temp @ asso_leg_term_neg
-                # --------------------------------------------------------------------------------------------------------------------------
-
-                # Assemble the coefficient matrix and additional terms
-                # --------------------------------------------------------------------------------------------------------------------------
-                alpha = M_inv[:, None] * (D_pos * W[None, :] - np.eye(N))
-                beta = M_inv[:, None] * D_neg * W[None, :]
-                if there_is_beam_source:
-                    X_tilde_arr[ind, :] = np.concatenate([-M_inv * X_pos, M_inv * X_neg])
+                    X_tilde = np.concatenate([-M_inv * X_pos, M_inv * X_neg])
+                    
+                    A = np.concatenate(
+                        [
+                            np.concatenate([-alpha, -beta], axis=1),
+                            np.concatenate([beta, alpha], axis=1),
+                        ],
+                        axis=0,
+                    )
+                    np.fill_diagonal(A, np.diag(A) + 1 / mu0)
+                    B_collect[ind, :] = -np.linalg.solve(A, X_tilde) # We moved the minus sign out
+                    
                 # --------------------------------------------------------------------------------------------------------------------------
                 
+                
                 alpha_arr[ind, :, :] = alpha
                 beta_arr[ind, :, :] = beta
                 no_shortcut_indices.append(ind)
-                if there_is_iso_source and m_equals_0:
+                if m_equals_0 and there_is_iso_source:
                     no_shortcut_indices_0.append(l)
                 
             else:
                 # This is a shortcut to the diagonalization results
                 G = np.zeros((NQuad, NQuad))
-                G[N:, :N] = np.eye(N)
-                G[:N, N:] = np.eye(N)
+                np.fill_diagonal(G[N:, :N], 1)
+                np.fill_diagonal(G[:N, N:], 1)
                 
                 G_collect[ind, :, :] = G
                 K_collect[ind, :] = -1 / mu_arr
-                if there_is_iso_source and m_equals_0:
+                if m_equals_0 and there_is_iso_source:
                     G_inv_collect_0[l, :, :] = G
+                    
             ind += 1
                 
     if len(no_shortcut_indices) > 0:
@@ -184,26 +202,16 @@ def _solve_for_gen_and_part_sols(
             ),
             axis=1,
         )
-        G_inv_arr = np.linalg.inv(G_arr)
         
         G_collect[no_shortcut_indices, :, :] = G_arr
         K_collect[no_shortcut_indices, :] = K_arr
         if len(no_shortcut_indices_0) > 0:
-            G_inv_collect_0[no_shortcut_indices_0, :, :] = G_inv_arr[: len(no_shortcut_indices_0), :, :]
-        # --------------------------------------------------------------------------------------------------------------------------
-
-        # Particular solution for the direct beam source (refer to Section 3.6.1 of the Comprehensive Documentation)
-        # --------------------------------------------------------------------------------------------------------------------------
-        if there_is_beam_source:
-            X_tilde_arr = X_tilde_arr[: len(no_shortcut_indices), :]
-            B_collect[no_shortcut_indices, :] = np.einsum(
-                "lij, ljk, lk -> li",
-                -G_arr / (1 / mu0 + K_arr)[:, None, :],
-                G_inv_arr,
-                X_tilde_arr,
-                optimize=True,
+            G_inv_collect_0[no_shortcut_indices_0, :, :] = np.linalg.inv(
+                G_collect[: len(no_shortcut_indices_0), :, :]
             )
+
         # --------------------------------------------------------------------------------------------------------------------------
+
         
     if there_is_beam_source:
         if there_is_iso_source:

PythonicDISORT/pydisort.py

@@ -22,7 +22,7 @@ def pydisort(
     NT_cor=False,
     BDRF_Fourier_modes=[],
     s_poly_coeffs=np.array([[]]),
-    NLayers_to_use_sparse=13,
+    use_banded_solver_NLayers=13,
     _autograd_compatible = False,
 ):
     """Solves the 1D RTE for the fluxes, and optionally intensity,
@@ -75,8 +75,8 @@ def pydisort(
     s_poly_coeffs : optional, array
         Polynomial coefficients of isotropic internal sources.
         Arrange coefficients from lowest order term to highest.
-    NLayers_to_use_sparse : optional, int
-        At or above how many atmospheric layers should SciPy's sparse matrix framework be used?
+    use_banded_solver_NLayers : optional, int
+        At or above how many atmospheric layers should `scipy.linalg.solve_banded` be used?
 
     Returns
     -------
@@ -115,25 +115,25 @@ def pydisort(
     
     """
     Arguments of pydisort
-    |         Variable        |            Type / Shape            |
-    | ----------------------- | ---------------------------------- |
-    | `tau_arr`               | `NLayers`                          |
-    | `omega_arr`             | `NLayers`                          |
-    | `NQuad`                 | scalar                             |
-    | `Leg_coeffs_all`        | `NLeg_all`                         |
-    | `mu0`                   | scalar                             |
-    | `I0`                    | scalar                             |
-    | `phi0`                  | scalar                             |
-    | `NLeg`                  | scalar                             |
-    | `NFourier`              | scalar                             |
-    | `b_pos`                 | `NQuad/2 x NFourier` or scalar     |
-    | `b_neg`                 | `NQuad/2 x NFourier` or scalar     |              
-    | `only_flux`             | boolean                            |
-    | `f_arr`                 | `NLayers` or scalar                |
-    | `NT_cor`                | boolean                            |
-    | `BDRF_Fourier_modes`    | `NBDRF`                            |
-    | `s_poly_coeffs`         | `NLayers x Nscoeffs` or `Nscoeffs` |
-    | `NLayers_to_use_sparse` | scalar                             |
+    |          Variable           |            Type / Shape            |
+    | --------------------------- | ---------------------------------- |
+    | `tau_arr`                   | `NLayers`                          |
+    | `omega_arr`                 | `NLayers`                          |
+    | `NQuad`                     | scalar                             |
+    | `Leg_coeffs_all`            | `NLeg_all`                         |
+    | `mu0`                       | scalar                             |
+    | `I0`                        | scalar                             |
+    | `phi0`                      | scalar                             |
+    | `NLeg`                      | scalar                             |
+    | `NFourier`                  | scalar                             |
+    | `b_pos`                     | `NQuad/2 x NFourier` or scalar     |
+    | `b_neg`                     | `NQuad/2 x NFourier` or scalar     |              
+    | `only_flux`                 | boolean                            |
+    | `f_arr`                     | `NLayers` or scalar                |
+    | `NT_cor`                    | boolean                            |
+    | `BDRF_Fourier_modes`        | `NBDRF`                            |
+    | `s_poly_coeffs`             | `NLayers x Nscoeffs` or `Nscoeffs` |
+    | `use_banded_solver_NLayers` | scalar                             |
     
     Notable internal variables of pydisort
     |          Variable            |     Type / Shape     |
@@ -190,7 +190,7 @@ def pydisort(
     NLayers = len(tau_arr)
     b_pos_is_scalar = False
     b_neg_is_scalar = False
-    thickness_arr = np.concatenate([[tau_arr[0]], np.diff(tau_arr)])
+    thickness_arr = np.append(tau_arr[0], np.diff(tau_arr))
     NLeg_all = np.shape(Leg_coeffs_all)[1]
     N = NQuad // 2
     there_is_beam_source = I0 > 0
@@ -243,8 +243,8 @@ def pydisort(
         assert np.shape(b_neg) == (N, NFourier)
     # The fractional scattering must be between 0 and 1
     assert np.all(0 <= f_arr) and np.all(f_arr <= 1)
-    # The minimum threshold is the minimum numbers of layers: 1
-    assert(NLayers_to_use_sparse >= 1)
+    # The minimum threshold is 3 else the matrix will not be banded
+    assert(use_banded_solver_NLayers >= 3)
     # --------------------------------------------------------------------------------------------------------------------------
     
     # Some more setup
@@ -277,9 +277,7 @@ def pydisort(
     if np.any(f_arr > 0):
         scale_tau = 1 - omega_arr * f_arr
         scaled_thickness_arr = scale_tau * thickness_arr
-        scaled_tau_arr_with_0 = np.array(
-            list(map(lambda l: np.sum(scaled_thickness_arr[:l]), range(NLayers + 1)))
-        )
+        scaled_tau_arr_with_0 = np.append(0, np.cumsum(scaled_thickness_arr))
         weighted_scaled_Leg_coeffs = ((Leg_coeffs - f_arr[:, None]) / (1 - f_arr[:, None])) * (
             2 * np.arange(NLeg) + 1
         )[None, :]
@@ -287,7 +285,7 @@ def pydisort(
     else:
         # This is a shortcut to the same results
         scale_tau = np.ones(NLayers)
-        scaled_tau_arr_with_0 = np.concatenate([[0], tau_arr])
+        scaled_tau_arr_with_0 = np.append(0, tau_arr)
         weighted_scaled_Leg_coeffs = Leg_coeffs * (2 * np.arange(NLeg) + 1)[None, :]
         scaled_omega_arr = omega_arr
     # --------------------------------------------------------------------------------------------------------------------------
@@ -313,12 +311,12 @@ def pydisort(
             there_is_beam_source,
             b_pos, b_neg,
             b_pos_is_scalar, b_neg_is_scalar,
-            s_poly_coeffs,
             Nscoeffs,
+            s_poly_coeffs,
             there_is_iso_source,
             scale_tau,
             only_flux,
-            NLayers_to_use_sparse,
+            use_banded_solver_NLayers,
             _autograd_compatible,
         )
         
@@ -371,14 +369,14 @@ def pydisort(
             
             if is_antiderivative_wrt_tau:
                 TMS_correction_pos = (
-                    np.exp(-scaled_tau / mu0)[None, :] / (-scale_tau[None, l] / mu0)
+                    np.exp(-scaled_tau / mu0)[None, :] / (-scale_tau / mu0)[None, l]
                     - np.exp(
                         (scaled_tau - scaled_tau_arr_l)[None, :] / mu_arr_pos[:, None]
                         - scaled_tau_arr_l[None, :] / mu0
                     ) / (scale_tau[None, l] / mu_arr_pos[:, None])
                 )
                 TMS_correction_neg = (
-                    np.exp(-scaled_tau / mu0)[None, :] / (-scale_tau[None, l] / mu0)
+                    np.exp(-scaled_tau / mu0)[None, :] / (-scale_tau / mu0)[None, l]
                     - np.exp(
                         (scaled_tau_arr_lm1 - scaled_tau)[None, :] / mu_arr_pos[:, None]
                         - scaled_tau_arr_lm1[None, :] / mu0
@@ -501,7 +499,7 @@ def pydisort(
                                         - scaled_tau_arr_with_0_l_repeat_neg / mu0
                                     )
                                 )
-                                / (-scale_tau_repeat_pos / mu_arr_pos[:, None])
+                                / (-scale_tau_repeat_neg / mu_arr_pos[:, None])
                             )[:, :, None]
                         )
                         contribution_from_other_layers_neg = np.sum(
@@ -637,11 +635,11 @@ def pydisort(
             there_is_beam_source,
             b_pos, b_neg,
             b_pos_is_scalar, b_neg_is_scalar,
-            s_poly_coeffs,
             Nscoeffs,
+            s_poly_coeffs,
             there_is_iso_source,
             scale_tau,
             only_flux,
-            NLayers_to_use_sparse,
+            use_banded_solver_NLayers,
             _autograd_compatible,
         )

PythonicDISORT/subroutines.py

@@ -136,7 +136,7 @@ def Clenshaw_Curtis_quad(Nphi, c=0, d=(2 * pi)):
     Nphi_pos = Nphi // 2
     phi_arr_pos = np.cos(pi * np.arange(Nphi_pos) / Nphi)
     phi_arr = np.concatenate([-phi_arr_pos, [0], np.flip(phi_arr_pos)])
-    diff = np.concatenate([[2], 2 / (1 - 4 * np.arange(1, Nphi_pos + 1) ** 2)])
+    diff = np.append(2, 2 / (1 - 4 * np.arange(1, Nphi_pos + 1) ** 2))
     weights_phi_pos = sc.fft.idct(diff, type=1)
     weights_phi_pos[0] /= 2
     full_weights_phi = np.hstack((weights_phi_pos, np.flip(weights_phi_pos[:-1])))
@@ -320,12 +320,47 @@ def interpolate(mu_arr_pos, u):
             return np.squeeze(results)[()]
 
     return u_interpol
-                
+          
 
+def to_diag_ordered_form(A, sym_offset):
+    """
+    Convert a matrix A to the diagonal ordered form required by `scipy.linalg.solve_banded`.
+    We assume that the matrix has the same number of super- and sub-diagonals.
+
+    Parameters
+    ----------
+    A : 2darray
+        The square matrix to be converted.
+    sym_offset : int
+        The number of super- or sub-diagonals (assumed to be equal)
+
+    Returns
+    -------
+    2darray
+        The diagonal ordered form matrix as required by solve_banded.
+    """
+    n = A.shape[0]
+    indices = np.arange(n)
+
+    return np.concatenate(
+        [
+            A[
+                indices - np.arange(sym_offset, -1, -1)[:, None],
+                indices[None, :],
+            ],
+            A[
+                indices - np.arange(n - 1, n - sym_offset - 1, -1)[:, None],
+                indices[None, :],
+            ],
+        ],
+        axis=0,
+    )
+
+    
 def _mathscr_v(tau,                             # Input optical depths
                 l,                              # Layer index of each input optical depth
-                s_poly_coeffs,                  # Polynomial coefficients of isotropic source
                 Nscoeffs,                       # Number of isotropic source polynomial coefficients
+                s_poly_coeffs,                  # Polynomial coefficients of isotropic source
                 G,                              # Eigenvector matrices
                 K,                              # Eigenvalues
                 G_inv,                          # Inverse of eigenvector matrix
@@ -344,6 +379,7 @@ def _mathscr_v(tau,                             # Input optical depths
     | --------------------------- | ------------------------------------- |
     | `tau`                       | `Ntau`                                |
     | `l`                         | `Ntau`                                |
+    | `Nscoeffs`                  | scalar                                |
     | `s_poly_coeffs`             | `NLayers x Nscoeffs` or `Nscoeffs`    |
     | `G`                         | `NLayers<= x NQuad x NQuad`           |
     | `K`                         | `NLayers<= x NQuad`                   |
@@ -409,30 +445,20 @@ def _mathscr_v(tau,                             # Input optical depths
             ),
         )
     
-    powers = np.flip(np.arange(Nscoeffs))[None, :]
+    powers = np.arange(Nscoeffs - 1, -1, -1)[None, :]
     if is_antiderivative_wrt_tau:
         tau_poly = tau[:, None] ** (powers + 1) / (powers + 1)
     else:
         tau_poly = tau[:, None] ** powers
-        
+    
     return np.einsum(
-        "tik, tc, ctk, tkj, j -> it",
+        "tik, tc, ctk -> it",
         G[l, :, :],
         tau_poly,
-        mathscr_v_coeffs[:, l, :],
-        G_inv[l, :, :],
-        1 / mu_arr,
+        (mathscr_v_coeffs * (G_inv @ (1 / mu_arr))[None, :, :])[:, l, :],
         optimize=True,
     )
     
-def _nd_slice_to_indexes(nd_slice):
-    """Code taken from https://stackoverflow.com/questions/64097025/how-to-convert-n-d-slice-to-indexes-in-numpy.
-    This function is used to efficiently construct sparse COO matrices.
-    
-    """
-    grid = np.mgrid[{tuple: nd_slice, slice: (nd_slice,)}[type(nd_slice)]]
-    return tuple(grid[i].ravel() for i in range(grid.shape[0]))
-    
 
 def _compare(results, mu_to_compare, reorder_mu, flux_up, flux_down, u):
     """Performs pointwise comparisons between results from Stamnes' DISORT,

{PythonicDISORT-0.9.0.dist-info → PythonicDISORT-0.9.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PythonicDISORT
-Version: 0.9.0
+Version: 0.9.1
 Summary: Discrete Ordinates Solver for the (1D) Radiative Transfer Equation in a single or multi-layer atmosphere.
 Author-email: Dion HO Jia Xu <dh3065@columbia.edu>
 Project-URL: Homepage, https://github.com/LDEO-CREW/Pythonic-DISORT

PythonicDISORT-0.9.1.dist-info/RECORD

@@ -0,0 +1,11 @@
+PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
+PythonicDISORT/_assemble_intensity_and_fluxes.py,sha256=j3w78Y8844HlI1z8JFmOgNs4ogDsma8VySbI8W2-m-8,27868
+PythonicDISORT/_solve_for_coeffs.py,sha256=UccDeCrr0jBCKbK73Eq9SwWEkSkjVyAVheew6nYcWdw,15368
+PythonicDISORT/_solve_for_gen_and_part_sols.py,sha256=MNNg5f4S5XY0D6btH3D-9p-eye9M_rT7GUlaHnys7Ts,13138
+PythonicDISORT/pydisort.py,sha256=SNpTpEcRPHNZqnVL86xV5AzPPNu7AVgmpXhV6a5Rqh8,33204
+PythonicDISORT/subroutines.py,sha256=on0VpAtUc8b7ysaQqy8VetSnTMhsE_yEDAwyRHX2IuU,20049
+PythonicDISORT-0.9.1.dist-info/LICENSE.md,sha256=q1i1QVh0OLzK0fj6kaNRNWsM_KZca66JwWgWwbuCPlA,1092
+PythonicDISORT-0.9.1.dist-info/METADATA,sha256=SLwWEGhdEqHtpHNN14-Y0Kd2tvIYByXubdqzpr1nqKc,7102
+PythonicDISORT-0.9.1.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
+PythonicDISORT-0.9.1.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
+PythonicDISORT-0.9.1.dist-info/RECORD,,

{PythonicDISORT-0.9.0.dist-info → PythonicDISORT-0.9.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (72.1.0)
+Generator: setuptools (74.1.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

PythonicDISORT-0.9.0.dist-info/RECORD

@@ -1,11 +0,0 @@
-PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
-PythonicDISORT/_assemble_intensity_and_fluxes.py,sha256=SFX5SMUhKNt_MU-0t3kv5lLXxc1EzVK2EIDLnggoIEw,27206
-PythonicDISORT/_solve_for_coeffs.py,sha256=ZGThwif73zAlhzXYdVzV2nJxdxRioq4MIFsni3_5-Hc,19594
-PythonicDISORT/_solve_for_gen_and_part_sols.py,sha256=9ScaxY7SH3lYaE2UBnEIsmsDeXdjcf74_QWaeqyOIRs,12973
-PythonicDISORT/pydisort.py,sha256=HmAdvbDHQoHQy1xXToJmT3X5MGRudMKoI0L7zAaOd-8,33188
-PythonicDISORT/subroutines.py,sha256=2cJLQc46ctAiqY_qARlsxQZPYVcBOwkXf0-diCtnGXo,19446
-PythonicDISORT-0.9.0.dist-info/LICENSE.md,sha256=q1i1QVh0OLzK0fj6kaNRNWsM_KZca66JwWgWwbuCPlA,1092
-PythonicDISORT-0.9.0.dist-info/METADATA,sha256=dLzNm79C3rWzVMzC2GIboGpMloxs4Fiy3RpQgOMIMP4,7102
-PythonicDISORT-0.9.0.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-PythonicDISORT-0.9.0.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
-PythonicDISORT-0.9.0.dist-info/RECORD,,