PythonicDISORT 1.6__py3-none-any.whl → 1.8__py3-none-any.whl

PythonicDISORT/_assemble_intensity_and_fluxes.py

@@ -29,7 +29,8 @@ def _assemble_intensity_and_fluxes(
     only_flux,                          # Only compute fluxes?
     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?
-):  
+    cache_asso_leg,                     # Assoc. Legendre table cache mode: None, "no_mu0" or "mu0"
+):
     """Assembles the solution functions: intensity `u`, upward flux `flux_up`, downward flux `flux_down` 
     from the previously solved eigenpairs and coefficients from the boundary conditions. 
     Returns these solution functions.
@@ -78,7 +79,8 @@ def _assemble_intensity_and_fluxes(
     | `only_flux`                    | boolean                                      |
     | `use_banded_solver_NLayers`    | scalar                                       |
     | `autograd_compatible`          | boolean                                      |
-    
+    | `cache_asso_leg`               | None, "no_mu0" or "mu0"                      |
+
     Notable internal variables of _assemble_intensity_and_fluxes
     |     Variable      |             Type / Shape               |
     | ----------------- | -------------------------------------- |
@@ -117,6 +119,7 @@ def _assemble_intensity_and_fluxes(
         mu0, I0_div_4pi,
         there_is_beam_source,
         there_is_iso_source,
+        cache_asso_leg,
     )
     if there_is_beam_source and there_is_iso_source:
         G_collect, K_collect, B_collect, G_inv_collect_0 = outputs
@@ -161,8 +164,66 @@ def _assemble_intensity_and_fluxes(
     G_collect_0 = G_collect[0, :, :, :]
     K_collect_0 = K_collect[0, :, :]
     GC_collect_0 = GC_collect[0, :, :, :]
-    # --------------------------------------------------------------------------------------------------------------------------   
-        
+    # --------------------------------------------------------------------------------------------------------------------------
+
+    # Value-keyed cache for the most recent `tau` evaluation grid, shared across the output
+    # functions `u`, `u0`, `flux_up` and `flux_down` (evaluating several outputs on the same
+    # grid, e.g. `flux_up(tau); flux_down(tau)`, then recomputes nothing).
+    # Stores a copy of `tau` (compared by value like the cache in `subroutines.interpolate`)
+    # with the derived layer indices and delta-scaled `tau`, and, lazily, the exponentials of
+    # the homogeneous-solution exponents for the 0th mode ("exp0") and for all modes ("expm").
+    # The cache is bypassed when `autograd_compatible` as it would break gradient tracing.
+    # --------------------------------------------------------------------------------------------------------------------------
+    is_delta_scaled = np.any(scale_tau != 1)
+    _tau_cache = {"tau": None, "l": None, "scaled_tau": None, "exp0": None, "expm": None}
+
+    def _process_tau(tau):
+        """Layer indices and delta-scaled `tau` for the evaluation grid `tau`, cached."""
+        if (not autograd_compatible and _tau_cache["tau"] is not None
+                and np.array_equal(tau, _tau_cache["tau"])):
+            return _tau_cache["l"], _tau_cache["scaled_tau"]
+        # Atmospheric layer indices: smallest l with tau <= tau_arr[l]
+        l = np.searchsorted(tau_arr, tau, side="left")
+        if is_delta_scaled:
+            scaled_tau = scaled_tau_arr_with_0[1:][l] - (tau_arr[l] - tau) * scale_tau[l]
+        else:
+            scaled_tau = tau
+        if not autograd_compatible:
+            _tau_cache.update(
+                tau=np.array(tau, copy=True), l=l, scaled_tau=scaled_tau, exp0=None, expm=None
+            )
+        return l, scaled_tau
+
+    def _exp_exponent(l, scaled_tau, all_modes):
+        """exp of the homogeneous-solution exponents for the (cached) evaluation grid.
+        Returns shape (Ntau, NQuad) for the 0th mode, or (NFourier, Ntau, NQuad) for all modes."""
+        key = "expm" if all_modes else "exp0"
+        if not autograd_compatible and _tau_cache[key] is not None:
+            return _tau_cache[key]
+        scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
+        scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
+        if all_modes:
+            exponent = np.concatenate(
+                [
+                    K_collect[:, l, :N] * (scaled_tau - scaled_tau_arr_lm1)[None, :, None],
+                    K_collect[:, l, N:] * (scaled_tau - scaled_tau_arr_l)[None, :, None],
+                ],
+                axis=-1,
+            )
+        else:
+            exponent = np.concatenate(
+                [
+                    K_collect_0[l, :N] * (scaled_tau - scaled_tau_arr_lm1)[:, None],
+                    K_collect_0[l, N:] * (scaled_tau - scaled_tau_arr_l)[:, None],
+                ],
+                axis=-1,
+            )
+        result = np.exp(exponent)
+        if not autograd_compatible:
+            _tau_cache[key] = result
+        return result
+    # --------------------------------------------------------------------------------------------------------------------------
+
     if not only_flux:
  
         # Construct the intensity function (refer to section 3.7 of the Comprehensive Documentation)
@@ -181,27 +242,10 @@ def _assemble_intensity_and_fluxes(
             if np.any(tau < 0) or np.any(tau > tau_arr[-1]):
                 raise ValueError("tau input outside the tau range specified for the atmosphere (check `tau_arr`).")
             
-            # Atmospheric layer indices
-            l = np.argmax(tau[:, None] <= tau_arr[None, :], axis=1)
-            scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
-            scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
+            # Atmospheric layer indices, delta-M scaling and exponentials (cached)
+            l, scaled_tau = _process_tau(tau)
+            exp_exponent = _exp_exponent(l, scaled_tau, all_modes=True)
 
-            # Delta-M scaling
-            if np.any(scale_tau != np.ones(NLayers)):
-                tau_dist_from_top = tau_arr[l] - tau
-                scaled_tau_dist_from_top = tau_dist_from_top * scale_tau[l]
-                scaled_tau = scaled_tau_arr_l - scaled_tau_dist_from_top
-            else:
-                scaled_tau = tau
-            
-            exponent = np.concatenate(
-                [
-                    K_collect[:, l, :N] * (scaled_tau - scaled_tau_arr_lm1)[None, :, None],
-                    K_collect[:, l, N:] * (scaled_tau - scaled_tau_arr_l)[None, :, None],
-                ],
-                axis=-1,
-            )
-            
             direct_beam_contribution = 0
             if is_antiderivative_wrt_tau:
                 if there_is_beam_source:
@@ -213,7 +257,7 @@ def _assemble_intensity_and_fluxes(
                     np.einsum(
                         "mtij, mtj -> mit",
                         GC_collect[:, l, :, :],
-                        np.exp(exponent) / (scale_tau[None, :, None] * K_collect)[:, l, :],
+                        exp_exponent / (scale_tau[None, :, None] * K_collect)[:, l, :],
                         optimize=True,
                     )
                     + direct_beam_contribution
@@ -227,7 +271,7 @@ def _assemble_intensity_and_fluxes(
 
                 um = (
                     np.einsum(
-                        "mtij, mtj -> mit", GC_collect[:, l, :, :], np.exp(exponent), optimize=True
+                        "mtij, mtj -> mit", GC_collect[:, l, :, :], exp_exponent, optimize=True
                     )
                     + direct_beam_contribution
                 )
@@ -236,7 +280,8 @@ def _assemble_intensity_and_fluxes(
             if there_is_iso_source:
                 l_uniq, l_inv = np.unique(l, return_inverse=True)
                 _mathscr_v_contribution = _mathscr_v(
-                    scaled_tau, scale_tau, l_inv, 
+                    scaled_tau, scale_tau[l_uniq], # `scale_tau` is only required to determine the antiderivative
+                    l_inv, 
                     Nscoeffs,
                     scaled_s_poly_coeffs[l_uniq],
                     G_collect_0[l_uniq],
@@ -253,49 +298,52 @@ def _assemble_intensity_and_fluxes(
                 else:
                     um[0, :, :] += _mathscr_v_contribution
                 
-            u = np.sum(
-                um[:, :, :, None]
-                * np.cos(np.arange(NFourier)[:, None] * (phi0 - phi)[None, :])[:, None, None, :],
-                axis=0,
+            # Contract over the Fourier modes directly: materializing the
+            # (NFourier, NQuad, Ntau, Nphi) product only to sum over its first axis
+            # is substantially slower and heavier for large NFourier * Nphi
+            u = np.einsum(
+                "mit, mp -> itp",
+                um,
+                np.cos(np.arange(NFourier)[:, None] * (phi0 - phi)[None, :]),
+                optimize=True,
             )
             
             outputs = (rescale_factor * np.squeeze(u),)
             output_is_tuple = False
 
             if return_Fourier_error:
-                exponent = np.concatenate(
-                    [
-                        K_collect[-1, l, :N] * (scaled_tau - scaled_tau_arr_lm1)[:, None],
-                        K_collect[-1, l, N:] * (scaled_tau - scaled_tau_arr_l)[:, None],
-                    ],
-                    axis=-1,
-                )
-                
+                # The last Fourier mode's exponentials are the last slice of the cached tensor
+                exp_exponent_last = exp_exponent[-1]
+
                 direct_beam_contribution = 0
                 if is_antiderivative_wrt_tau:
                     if there_is_beam_source:
+                        # `B_collect[-1].T` has shape (NQuad, NLayers): the last Fourier mode's
+                        # particular-solution coefficients (`B_collect.T[-1]` would wrongly
+                        # index the last quadrature node)
                         direct_beam_contribution = (
-                            B_collect.T[-1, :, :] / (-scale_tau / mu0)[None, :]
+                            B_collect[-1].T / (-scale_tau / mu0)[None, :]
                         )[:, l] * np.exp(-scaled_tau / mu0)[None, :]
 
                     ulast = (
                         np.einsum(
                             "tij, tj -> it",
                             GC_collect[-1, l, :, :],
-                            np.exp(exponent) / (scale_tau[:, None] * K_collect[-1, :, :])[l, :],
+                            exp_exponent_last / (scale_tau[:, None] * K_collect[-1, :, :])[l, :],
                             optimize=True,
                         )
                         + direct_beam_contribution
                     )
                 else:
                     if there_is_beam_source:
+                        # See the shape note in the `is_antiderivative_wrt_tau` branch above
                         direct_beam_contribution = (
-                            B_collect.T[-1, :, l] * np.exp(-scaled_tau / mu0)[None, :]
+                            B_collect[-1].T[:, l] * np.exp(-scaled_tau / mu0)[None, :]
                         )
 
                     ulast = (
                         np.einsum(
-                            "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
+                            "tij, tj -> it", GC_collect[-1, l, :, :], exp_exponent_last, optimize=True
                         )
                         + direct_beam_contribution
                     )
@@ -338,41 +386,25 @@ def _assemble_intensity_and_fluxes(
         if np.any(tau < 0) or np.any(tau > tau_arr[-1]):
             raise ValueError("tau input outside the tau range given for the atmosphere (check `tau_arr`).")
         
-        # Atmospheric layer indices
-        l = np.argmax(tau[:, None] <= tau_arr[None, :], axis=1)
-        scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
-        scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
+        # Atmospheric layer indices, delta-M scaling and exponentials (cached)
+        l, scaled_tau = _process_tau(tau)
+        exp_exponent = _exp_exponent(l, scaled_tau, all_modes=False)
 
-        # Delta-M scaling
-        if np.any(scale_tau != np.ones(NLayers)):
-            tau_dist_from_top = tau_arr[l] - tau
-            scaled_tau_dist_from_top = tau_dist_from_top * scale_tau[l]
-            scaled_tau = scaled_tau_arr_l - scaled_tau_dist_from_top
-            # The following complements the function `subroutines.generate_diff_act_flux_funcs`
-            # 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
-                    )
-                else:
-                    act_dscale_reclassification = I0 * np.exp(-scaled_tau / mu0) - I0 * np.exp(-tau / mu0)
-        else:
-            scaled_tau = tau
-            if _return_act_dscale_for_reclass:
+        # The following complements the function `subroutines.generate_diff_act_flux_funcs`
+        # in performing the reclassification of delta-scaled actinic flux
+        if _return_act_dscale_for_reclass:
+            if not is_delta_scaled:
                 act_dscale_reclassification = 0
+            elif is_antiderivative_wrt_tau:
+                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)
 
-        exponent = np.concatenate(
-            [
-                K_collect_0[l, :N] * (scaled_tau - scaled_tau_arr_lm1)[:, None],
-                K_collect_0[l, N:] * (scaled_tau - scaled_tau_arr_l)[:, None],
-            ],
-            axis=-1,
-        )
-        
         direct_beam_contribution = 0
-        if is_antiderivative_wrt_tau and there_is_beam_source:
+        if is_antiderivative_wrt_tau:
             if there_is_beam_source:
                 direct_beam_contribution = (B_collect_0.T / (-scale_tau / mu0)[None, :])[
                     :, l
@@ -382,7 +414,7 @@ def _assemble_intensity_and_fluxes(
                 np.einsum(
                     "tij, tj -> it",
                     GC_collect_0[l, :, :],
-                    np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :],
+                    exp_exponent / (scale_tau[:, None] * K_collect_0)[l, :],
                     optimize=True,
                 )
                 + direct_beam_contribution
@@ -395,7 +427,7 @@ def _assemble_intensity_and_fluxes(
 
             u0 = (
                 np.einsum(
-                    "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
+                    "tij, tj -> it", GC_collect_0[l, :, :], exp_exponent, optimize=True
                 )
                 + direct_beam_contribution
             )
@@ -404,7 +436,8 @@ def _assemble_intensity_and_fluxes(
         if there_is_iso_source:
             l_uniq, l_inv = np.unique(l, return_inverse=True)
             _mathscr_v_contribution = _mathscr_v(
-                scaled_tau, scale_tau, l_inv,
+                scaled_tau, scale_tau[l_uniq], # `scale_tau` is only required to determine the antiderivative
+                l_inv,
                 Nscoeffs,
                 scaled_s_poly_coeffs[l_uniq],
                 G_collect_0[l_uniq],
@@ -447,23 +480,15 @@ def _assemble_intensity_and_fluxes(
         if np.any(tau < 0) or np.any(tau > tau_arr[-1]):
             raise ValueError("tau input outside the tau range given for the atmosphere (check `tau_arr`).")
         
-        # Atmospheric layer indices
-        l = np.argmax(tau[:, None] <= tau_arr[None, :], axis=1)
-        scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
-        scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
-
-        # Delta-M scaling
-        if np.any(scale_tau != np.ones(NLayers)):
-            tau_dist_from_top = tau_arr[l] - tau
-            scaled_tau_dist_from_top = tau_dist_from_top * scale_tau[l]
-            scaled_tau = scaled_tau_arr_l - scaled_tau_dist_from_top
-        else:
-            scaled_tau = tau
+        # Atmospheric layer indices, delta-M scaling and exponentials (cached)
+        l, scaled_tau = _process_tau(tau)
+        exp_exponent = _exp_exponent(l, scaled_tau, all_modes=False)
 
         if there_is_iso_source:
             l_uniq, l_inv = np.unique(l, return_inverse=True)
             _mathscr_v_contribution = _mathscr_v(
-                scaled_tau, scale_tau, l_inv,
+                scaled_tau, scale_tau[l_uniq], # `scale_tau` is only required to determine the antiderivative
+                l_inv,
                 Nscoeffs,
                 scaled_s_poly_coeffs[l_uniq],
                 G_collect_0[l_uniq, :N, :],
@@ -475,14 +500,6 @@ def _assemble_intensity_and_fluxes(
         else:
             _mathscr_v_contribution = 0
         
-        exponent = np.concatenate(
-            [
-                K_collect_0[l, :N] * (scaled_tau - scaled_tau_arr_lm1)[:, None],
-                K_collect_0[l, N:] * (scaled_tau - scaled_tau_arr_l)[:, None],
-            ],
-            axis=-1,
-        )
-        
         direct_beam_contribution = 0
         if is_antiderivative_wrt_tau:
             if there_is_beam_source:
@@ -494,7 +511,7 @@ def _assemble_intensity_and_fluxes(
                 np.einsum(
                     "tij, tj -> it",
                     GC_pos[l, :, :],
-                    np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :],
+                    exp_exponent / (scale_tau[:, None] * K_collect_0)[l, :],
                 )
                 + direct_beam_contribution
                 + _mathscr_v_contribution
@@ -504,7 +521,7 @@ def _assemble_intensity_and_fluxes(
                 direct_beam_contribution = B_pos[:, l] * np.exp(-scaled_tau / mu0)[None, :]
                 
             u0_pos = (
-                np.einsum("tij, tj -> it", GC_pos[l, :, :], np.exp(exponent))
+                np.einsum("tij, tj -> it", GC_pos[l, :, :], exp_exponent)
                 + direct_beam_contribution
                 + _mathscr_v_contribution
             )
@@ -529,23 +546,15 @@ def _assemble_intensity_and_fluxes(
         if np.any(tau < 0) or np.any(tau > tau_arr[-1]):
             raise ValueError("tau input outside the tau range given for the atmosphere (check `tau_arr`).")
 
-        # Atmospheric layer indices
-        l = np.argmax(tau[:, None] <= tau_arr[None, :], axis=1)
-        scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
-        scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
-
-        # Delta-M scaling
-        if np.any(scale_tau != np.ones(NLayers)):
-            tau_dist_from_top = tau_arr[l] - tau
-            scaled_tau_dist_from_top = tau_dist_from_top * scale_tau[l]
-            scaled_tau = scaled_tau_arr_l - scaled_tau_dist_from_top
-        else:
-            scaled_tau = tau
+        # Atmospheric layer indices, delta-M scaling and exponentials (cached)
+        l, scaled_tau = _process_tau(tau)
+        exp_exponent = _exp_exponent(l, scaled_tau, all_modes=False)
 
         if there_is_iso_source:
             l_uniq, l_inv = np.unique(l, return_inverse=True)
             _mathscr_v_contribution = _mathscr_v(
-                scaled_tau, scale_tau, l_inv,
+                scaled_tau, scale_tau[l_uniq], # `scale_tau` is only required to determine the antiderivative
+                l_inv,
                 Nscoeffs,
                 scaled_s_poly_coeffs[l_uniq],
                 G_collect_0[l_uniq, N:, :],
@@ -557,25 +566,17 @@ def _assemble_intensity_and_fluxes(
         else:
             _mathscr_v_contribution = 0
             
-        exponent = np.concatenate(
-            [
-                K_collect_0[l, :N] * (scaled_tau - scaled_tau_arr_lm1)[:, None],
-                K_collect_0[l, N:] * (scaled_tau - scaled_tau_arr_l)[:, None],
-            ],
-            axis=-1,
-        )
-        
         direct_beam_contribution = 0
         direct_beam = 0
         direct_beam_scaled = 0
         if is_antiderivative_wrt_tau:
             if there_is_beam_source:
-                direct_beam_contribution = (B_neg / (-scale_tau / mu0)[None, l])[:, l] * np.exp(-scaled_tau / mu0)[None, :]
+                direct_beam_contribution = (B_neg / (-scale_tau / mu0)[None, :])[:, 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 / mu0)[None, l]
+                direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0) / (-scale_tau[l] / mu0)
             
             u0_neg = (
-                np.einsum("tij, tj -> it", GC_neg[l, :, :], np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :])
+                np.einsum("tij, tj -> it", GC_neg[l, :, :], exp_exponent / (scale_tau[:, None] * K_collect_0)[l, :])
                 + direct_beam_contribution
                 + _mathscr_v_contribution
             )
@@ -586,7 +587,7 @@ def _assemble_intensity_and_fluxes(
                 direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0)
             
             u0_neg = (
-                np.einsum("tij, tj -> it", GC_neg[l, :, :], np.exp(exponent))
+                np.einsum("tij, tj -> it", GC_neg[l, :, :], exp_exponent)
                 + direct_beam_contribution
                 + _mathscr_v_contribution
             )

PythonicDISORT/_solve_for_coeffs.py

@@ -79,7 +79,8 @@ def _solve_for_coeffs(
     use_banded_solver = (NLayers >= use_banded_solver_NLayers)
     dim = NLayers * NQuad
     C_collect = np.empty((NFourier, dim))
-    
+    thickness = np.diff(scaled_tau_arr_with_0)                                                       # (NLayers,)
+
     RHS = np.empty(dim)
     i_bot = dim - N # Equals `j_bot_right`
     j_bot_left = dim - NQuad
@@ -89,23 +90,52 @@ def _solve_for_coeffs(
         LHS_dof = np.zeros((6 * N - 1, dim), order="F") # Diagonal ordered form
         s0, s1 = LHS_dof.strides
         col_stride = s1 - s0  # Stride when moving +1 col and -1 row
-        
+
         def _view_slanted(i0, j0, nrows):
-            """
-            Return a view into `LHS_dof` using slanted mapping.
-            The function uses `numpy.lib.stride_tricks.as_strided` which is performant but dangerous, 
-            which is why this function is not exposed in `PythonicDISORT.subroutines`.
+            """Single slanted view into `LHS_dof` anchored at conceptual LHS coords (i0, j0).
+            Returns a 2D view of shape (nrows, N). Used for top/bottom boundary-condition rows.
+
+            Uses `numpy.lib.stride_tricks.as_strided` (zero-copy but unsafe), which is why
+            this helper is not exposed in `PythonicDISORT.subroutines`.
             See https://numpy.org/devdocs/reference/generated/numpy.lib.stride_tricks.as_strided.html.
-            
+
             Assume the following exist in the outer scope:
-                `LHS_dof`, `Nsupsubdiags`, `s0`, `col_stride`
+                `LHS_dof`, `Nsupsubdiags`, `N`, `s0`, `col_stride`.
             """
             r0 = Nsupsubdiags + i0 - j0
             base = LHS_dof[r0 : r0 + nrows, j0 : j0 + N]
             return as_strided(base, shape=(nrows, N), strides=(s0, col_stride))
-            
+
+        def _view_slanted_batched(i0, j0, nrows):
+            """Batched (NLayers-1, nrows, N) slanted view across interlayer continuity rows.
+            Anchored at conceptual LHS coords (i0, j0) for layer 0; layer l advances both
+            i and j by `NQuad` along the diagonal.
+
+            Same `(i0, j0)` convention as `_view_slanted` (slant arithmetic done internally).
+            Assume the following exist in the outer scope:
+                `LHS_dof`, `Nsupsubdiags`, `N`, `NQuad`, `NLayers`, `s0`, `s1`, `col_stride`.
+            """
+            r0 = Nsupsubdiags + i0 - j0 # matches `_view_slanted`'s internal slant
+            base = LHS_dof[r0 : r0 + nrows, j0 : j0 + N + (NLayers - 2) * NQuad]
+            return as_strided(base, shape=(NLayers - 1, nrows, N),
+                              strides=(NQuad * s1, s0, col_stride))
+
     else:
         LHS = np.zeros((dim, dim), order="F")
+        sR, sC = LHS.strides
+        # Stacks for a single batched `np.linalg.solve` over all Fourier modes after the
+        # assembly loop; this removes the per-mode LAPACK dispatch overhead. The banded
+        # path keeps its per-mode loop as `scipy.linalg.solve_banded` cannot batch.
+        LHS_batch = np.empty((NFourier, dim, dim))
+        RHS_batch = np.empty((NFourier, dim))
+        def _dense_slab(r0, c0, nrows, ncols):
+            """Batched (NLayers-1, nrows, ncols) view across interlayer blocks of `LHS`.
+            Anchored at LHS coords (r0, c0) for layer 0; layer l advances both row and column
+            anchors by `NQuad` along the diagonal. Each slab has shape (nrows, ncols)."""
+            base = LHS[r0 : r0 + nrows + (NLayers - 2) * NQuad,                                       
+                       c0 : c0 + ncols + (NLayers - 2) * NQuad]                                       
+            return as_strided(base, shape=(NLayers - 1, nrows, ncols),                                
+                              strides=(NQuad * (sR + sC), sR, sC))                                    
     
     # The following loops can easily be parallelized, but the speed-up is unlikely to be worth the overhead
     for m in range(NFourier):
@@ -256,13 +286,13 @@ def _solve_for_coeffs(
         # --------------------------------------------------------------------------------------------------------------------------
         
         
-        # Assemble LHS and solve (much of this code is replicated in section 4 of the Comprehensive Documentation)
+        # Assemble LHS and solve (the old version of this code is replicated in section 4 of the Comprehensive Documentation)
         # --------------------------------------------------------------------------------------------------------------------------
         G_0 = G_collect_m[0]
         G_0_nn = G_0[N:, :N]
         G_0_np = G_0[N:, N:]
         E_01 = np.exp(K_collect_m[0, :N] * scaled_tau_arr_with_0[1])[None, :]
-        
+
         G_L = G_collect_m[-1]
         G_L_pn = G_L[:N, :N]
         G_L_nn = G_L[N:, :N]
@@ -271,6 +301,9 @@ 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])
         )[None, :]
+        # Per-layer thickness exponentials, computed once per Fourier mode and reused below.
+        # If looping over `l in range(NLayers)`, the line below would be `E_pos_full[l] = exp(-thickness[l] * K_pos[l])`
+        E_pos_full = np.exp(-thickness[:, None] * K_collect_m[:, N:])                                # (NLayers, N)
         
         ################################## Use `sc.linalg.solve_banded` ##################################
         if use_banded_solver:       
@@ -292,44 +325,27 @@ def _solve_for_coeffs(
                 _view_slanted(i_bot, j_bot_left, N)[:] = G_L_pn * E_Lm1L
                 _view_slanted(i_bot, i_bot, N)[:] = G_L_pp
 
-            # ---------------- Interlayer / continuity BCs ----------------
-            for l in range(NLayers - 1):
-                G_l = G_collect_m[l]
-                G_l_pn = G_l[:N, :N]      # (N,N)
-                G_l_nn = G_l[N:, :N]      # (N,N)
-                G_l_ap = G_l[:, N:]       # (2N,N)
-
-                G_lp1 = G_collect_m[l + 1]
-                G_lp1_an = G_lp1[:, :N]   # (2N,N)
-                G_lp1_pp = G_lp1[:N, N:]  # (N,N)
-                G_lp1_np = G_lp1[N:, N:]  # (N,N)
-
-                tau_lm1 = scaled_tau_arr_with_0[l]
-                tau_l   = scaled_tau_arr_with_0[l + 1]
-                tau_lp1 = scaled_tau_arr_with_0[l + 2]
-
-                E_lm1l = np.exp(K_collect_m[l,     N:] * (tau_lm1 - tau_l))[None, :]    # (N,)
-                E_llp1 = np.exp(K_collect_m[l + 1, N:] * (tau_l   - tau_lp1))[None, :]  # (N,)
+            # ---------------- Interlayer / continuity BCs (batched across layers) ----------------
+            E_lm1l = E_pos_full[:-1, None, :]                                                        # (NLayers-1, 1, N)
+            E_llp1 = E_pos_full[1:,  None, :]                                                        # (NLayers-1, 1, N)
+            G_lo = G_collect_m[:-1]                                                                  # (NLayers-1, NQuad, NQuad)
+            G_hi = G_collect_m[1:]                                                                   # (NLayers-1, NQuad, NQuad)
 
-                start_row = N + l * NQuad
-                start_col = l * NQuad
-
-                _view_slanted(start_row,     start_col,     N)[:] = G_l_pn * E_lm1l
-                _view_slanted(start_row + N, start_col,     N)[:] = G_l_nn * E_lm1l
-                _view_slanted(start_row,     start_col + N, NQuad)[:] = G_l_ap
-
-                _view_slanted(start_row,     start_col + NQuad, NQuad)[:] = -G_lp1_an
-                _view_slanted(start_row,     start_col + 3 * N, N)[:] = -G_lp1_pp * E_llp1
-                _view_slanted(start_row + N, start_col + 3 * N, N)[:] = -G_lp1_np * E_llp1
+            _view_slanted_batched(N,       0,       N)[:]    = G_lo[:, :N, :N] * E_lm1l            
+            _view_slanted_batched(2 * N,   0,       N)[:]    = G_lo[:, N:, :N] * E_lm1l            
+            _view_slanted_batched(N,       N,       NQuad)[:] = G_lo[:, :,  N:]                    
+            _view_slanted_batched(N,       NQuad,   NQuad)[:] = -G_hi[:, :,  :N]                   
+            _view_slanted_batched(N,       3 * N,   N)[:]    = -G_hi[:, :N, N:] * E_llp1           
+            _view_slanted_batched(2 * N,   3 * N,   N)[:]    = -G_hi[:, N:, N:] * E_llp1           
         
             # Solve the system
-            C_m = sc.linalg.solve_banded(
+            C_collect[m, :] = sc.linalg.solve_banded(
                 (Nsupsubdiags, Nsupsubdiags),
                 LHS_dof,
                 RHS,
                 False,
                 True,
-                False,
+                True,
             )
         
         ##################################################################################################
@@ -348,43 +364,28 @@ def _solve_for_coeffs(
                 LHS[-N:, -NQuad : -N] = G_L_pn * E_Lm1L
                 LHS[-N:, -N:] = G_L_pp
 
-            # ---------------- Interlayer / continuity BCs ----------------
-            for l in range(NLayers - 1):
-                G_l = G_collect_m[l]
-                G_l_pn = G_l[:N, :N]
-                G_l_nn = G_l[N:, :N]
-                G_l_ap = G_l[:, N:]
-                
-                G_lp1 = G_collect_m[l + 1]
-                G_lp1_an = G_lp1[:, :N]
-                G_lp1_pp = G_lp1[:N, N:]
-                G_lp1_np = G_lp1[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
-                E_lm1l = np.exp(K_collect_m[l, N:] * (scaled_tau_arr_lm1 - scaled_tau_arr_l))
-                E_llp1 = np.exp(K_collect_m[l + 1, N:] * (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
-                LHS[N + start_row : NQuad + start_row, start_col : N + start_col] = G_l_nn * E_lm1l
-                LHS[start_row : NQuad + start_row, N + start_col : NQuad + start_col] = G_l_ap
-                
-                LHS[start_row : NQuad + start_row, NQuad + 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
-                LHS[N + start_row : NQuad + start_row, 3 * N + start_col : 4 * N + start_col] = -G_lp1_np * E_llp1
-                
-            # Solve the system
-            C_m = np.linalg.solve(LHS, RHS)
+            # ---------------- Interlayer / continuity BCs (batched across layers) ----------------
+            E_lm1l = E_pos_full[:-1, None, :]                                                        # (NLayers-1, 1, N)
+            E_llp1 = E_pos_full[1:,  None, :]                                                        # (NLayers-1, 1, N)
+            G_lo = G_collect_m[:-1]                                                                  
+            G_hi = G_collect_m[1:]                                                                   
+
+            _dense_slab(N,       0,       N,     N)[:]    = G_lo[:, :N, :N] * E_lm1l                 
+            _dense_slab(2 * N,   0,       N,     N)[:]    = G_lo[:, N:, :N] * E_lm1l                 
+            _dense_slab(N,       N,       NQuad, N)[:]    = G_lo[:, :,  N:]                          
+            _dense_slab(N,       NQuad,   NQuad, N)[:]    = -G_hi[:, :,  :N]                         
+            _dense_slab(N,       3 * N,   N,     N)[:]    = -G_hi[:, :N, N:] * E_llp1                
+            _dense_slab(2 * N,   3 * N,   N,     N)[:]    = -G_hi[:, N:, N:] * E_llp1                
                 
-            ################################################################################################## 
+            # Defer to the single batched solve over all Fourier modes below
+            LHS_batch[m] = LHS
+            RHS_batch[m] = RHS
+
+            ##################################################################################################
             # --------------------------------------------------------------------------------------------------------------------------
-        
-        C_collect[m, :] = C_m
-        
+
+    if not use_banded_solver:
+        # One batched call for all Fourier modes
+        C_collect = np.linalg.solve(LHS_batch, RHS_batch[:, :, None])[:, :, 0]
+
     return G_collect * C_collect.reshape((NFourier, NLayers, 1, NQuad))