PythonicDISORT 1.5__tar.gz → 1.7__tar.gz

{pythonicdisort-1.5/src/PythonicDISORT.egg-info → pythonicdisort-1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PythonicDISORT
-Version: 1.5
+Version: 1.7
 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
@@ -41,7 +41,7 @@ You may contact me, Dion, through dh3065@columbia.edu.
 
 The **GitHub repository** is https://github.com/LDEO-CREW/Pythonic-DISORT.
 
-Accompanying Journal of Open Source Software paper: https://joss.theoj.org/papers/10.21105/joss.06442.
+Accompanying Journal of Open Source Software paper: https://joss.theoj.org/papers/10.21105/joss.06442. **Please cite this paper if you use PythonicDISORT**.
 
 # Documentation
 https://pythonic-disort.readthedocs.io/en/latest/
@@ -76,7 +76,7 @@ These notebooks double up as examples of how to use PythonicDISORT. The tests wh
 # Installation
 
 * From PyPI: `pip install PythonicDISORT`
-* From Conda-forge: (TODO: need to first publish on Conda-forge)
+* From Conda-Forge: `conda install -c conda-forge pythonicdisort` or `mamba install -c conda-forge pythonicdisort`
 * By cloning repository: `pip install .` in the `Pythonic-DISORT` directory; `pip install -r all_optional_dependencies.txt` to install all optional dependencies (see *Requirements to run PythonicDISORT*)
 
 ## Requirements to run PythonicDISORT

{pythonicdisort-1.5 → pythonicdisort-1.7}/README.md

@@ -19,7 +19,7 @@ You may contact me, Dion, through dh3065@columbia.edu.
 
 The **GitHub repository** is https://github.com/LDEO-CREW/Pythonic-DISORT.
 
-Accompanying Journal of Open Source Software paper: https://joss.theoj.org/papers/10.21105/joss.06442.
+Accompanying Journal of Open Source Software paper: https://joss.theoj.org/papers/10.21105/joss.06442. **Please cite this paper if you use PythonicDISORT**.
 
 # Documentation
 https://pythonic-disort.readthedocs.io/en/latest/
@@ -54,7 +54,7 @@ These notebooks double up as examples of how to use PythonicDISORT. The tests wh
 # Installation
 
 * From PyPI: `pip install PythonicDISORT`
-* From Conda-forge: (TODO: need to first publish on Conda-forge)
+* From Conda-Forge: `conda install -c conda-forge pythonicdisort` or `mamba install -c conda-forge pythonicdisort`
 * By cloning repository: `pip install .` in the `Pythonic-DISORT` directory; `pip install -r all_optional_dependencies.txt` to install all optional dependencies (see *Requirements to run PythonicDISORT*)
 
 ## Requirements to run PythonicDISORT

{pythonicdisort-1.5 → pythonicdisort-1.7}/pyproject.toml

@@ -8,7 +8,7 @@ dependencies = [
   "scipy>=1.8.0",
 ]
 name = "PythonicDISORT"
-version = "1.5"
+version = "1.7"
 authors = [
   { name="Dion HO Jia Xu", email="dh3065@columbia.edu" },
 ]

{pythonicdisort-1.5 → pythonicdisort-1.7}/src/PythonicDISORT/_assemble_intensity_and_fluxes.py

@@ -167,7 +167,7 @@ def _assemble_intensity_and_fluxes(
  
         # Construct the intensity function (refer to section 3.7 of the Comprehensive Documentation)
         # --------------------------------------------------------------------------------------------------------------------------
-        def u(tau, phi, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False, *, _return_l=False):
+        def u(tau, phi, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False):
             """
             Intensity function with arguments `(tau, phi)` of types `(array or float, array or float)`.
             Returns an ndarray with axes corresponding to variation with `mu, tau, phi` respectively.
@@ -182,12 +182,12 @@ def _assemble_intensity_and_fluxes(
                 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)
+            l = np.searchsorted(tau_arr, tau, side="left")
             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)):
+            if np.any(scale_tau != 1):
                 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
@@ -236,7 +236,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],
@@ -318,13 +319,6 @@ def _assemble_intensity_and_fluxes(
             if return_tau_arr:
                 outputs += (tau_arr,)
                 output_is_tuple = True
-            
-            if _return_l: # This indicates to the function that NT corrections are activated
-                output_is_tuple = True
-                outputs += (l,)
-                if there_is_iso_source:
-                    # The TMS correction will use these
-                    outputs += (l_uniq, l_inv)
                 
             return outputs if output_is_tuple else outputs[0]
         # --------------------------------------------------------------------------------------------------------------------------
@@ -346,12 +340,12 @@ def _assemble_intensity_and_fluxes(
             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)
+        l = np.searchsorted(tau_arr, tau, side="left")
         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)):
+        if np.any(scale_tau != 1):
             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
@@ -411,7 +405,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],
@@ -455,12 +450,12 @@ def _assemble_intensity_and_fluxes(
             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)
+        l = np.searchsorted(tau_arr, tau, side="left")
         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)):
+        if np.any(scale_tau != 1):
             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
@@ -470,7 +465,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, :N, :],
@@ -537,12 +533,12 @@ def _assemble_intensity_and_fluxes(
             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)
+        l = np.searchsorted(tau_arr, tau, side="left")
         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)):
+        if np.any(scale_tau != 1):
             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
@@ -552,7 +548,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, N:, :],
@@ -577,9 +574,9 @@ def _assemble_intensity_and_fluxes(
         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, :])

{pythonicdisort-1.5 → pythonicdisort-1.7}/src/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,47 @@ 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                                                                          
+        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):
@@ -262,7 +287,7 @@ def _solve_for_coeffs(
         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 +296,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,35 +320,18 @@ 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(
@@ -329,7 +340,7 @@ def _solve_for_coeffs(
                 RHS,
                 False,
                 True,
-                False,
+                True,
             )
         
         ##################################################################################################
@@ -348,36 +359,18 @@ 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
+            # ---------------- 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                
                 
             # Solve the system
             C_m = np.linalg.solve(LHS, RHS)

{pythonicdisort-1.5 → pythonicdisort-1.7}/src/PythonicDISORT/_solve_for_gen_and_part_sols.py

@@ -69,13 +69,14 @@ def _solve_for_gen_and_part_sols(
     G_trivial.ravel()[N : N + N * stride : stride] = 1
     G_trivial.ravel()[NQuad * N : NQuad * N + N * stride : stride] = 1
     
+   # Flat (NFourier * NLayers, ...) storage so a single integer index list can address
+   # every (m, l) pair across Fourier modes for the cross-mode batched `eig`.
     G_collect = np.empty((NFourier * NLayers, NQuad, NQuad))
     K_collect = np.empty((NFourier * NLayers, NQuad))
     A_arr = np.empty((NFourier * NLayers, NQuad, NQuad))
     no_shortcut_indices = []
     no_shortcut_indices_0 = []
-    ind = 0
-    
+
     if there_is_beam_source:
         mu_arr_pos_mu0 = np.append(mu_arr_pos, -mu0)
         
@@ -109,66 +110,45 @@ def _solve_for_gen_and_part_sols(
         fac_asso_leg_term_posT = poch[None, :] * asso_leg_term_pos.T
         # --------------------------------------------------------------------------------------------------------------------------
 
-        # Loop over NLayers atmospheric layers
-        # --------------------------------------------------------------------------------------------------------------------------          
-        for l in range(NLayers):
-            weighted_scaled_Leg_coeffs_l = weighted_scaled_Leg_coeffs[l, m:]
-            scaled_omega_l = scaled_omega_arr[l]
-            omega_times_Leg_coeffs = (scaled_omega_l / 2) * weighted_scaled_Leg_coeffs_l
-            
-            if np.any(np.abs(omega_times_Leg_coeffs) > 1e-8): # There are shortcuts if multiple scattering is insignificant
-                
-                # Generate D
-                # --------------------------------------------------------------------------------------------------------------------------
-                D_temp = omega_times_Leg_coeffs[None, :] * fac_asso_leg_term_posT
-                D_pos = D_temp @ asso_leg_term_pos
-                D_neg = D_temp @ asso_leg_term_neg
-                
-                # --------------------------------------------------------------------------------------------------------------------------
-                
+        # Vectorized over NLayers atmospheric layers
+        # --------------------------------------------------------------------------------------------------------------------------
+        omega_times_Leg = (scaled_omega_arr[:, None] / 2) * weighted_scaled_Leg_coeffs[:, m:]   # (NLayers, NLeg-m)
+        no_shortcut = np.any(np.abs(omega_times_Leg) > 1e-8, axis=1)                           # (NLayers,)
+        ns_ind = np.flatnonzero(no_shortcut)                                                        
+        sh_ind = np.flatnonzero(~no_shortcut)                                                       
+        ind_base = m * NLayers                                                                  
+
+        # Shortcut layers: trivial diagonalization (no multiple scattering)
+        if sh_ind.size:
+            G_collect[ind_base + sh_ind] = G_trivial
+            K_collect[ind_base + sh_ind, :N] = -M_inv
+            K_collect[ind_base + sh_ind, N:] = M_inv
+            if m_equals_0 and there_is_iso_source:
+                G_inv_collect_0[sh_ind] = G_trivial
+
+        # Non-shortcut layers: full computation, batched over NLayers
+        if ns_ind.size:
+            D_temp = omega_times_Leg[ns_ind][:, None, :] * fac_asso_leg_term_posT[None, :, :]     # (Nns, N, NLeg-m)
+            D_pos = D_temp @ asso_leg_term_pos                                                    # (Nns, N, N)
+            D_neg = D_temp @ asso_leg_term_neg                                                    # (Nns, N, N)
+            DW = D_pos * W[None, None, :]                                                         
+            DW.reshape(-1, N * N)[:, ::N + 1] -= 1                                                # batched diag-sub
+            inds = ind_base + ns_ind                                                             
+            A_arr[inds, N:, N:] = M_inv[None, :, None] * DW                                       # alpha
+            A_arr[inds, N:, :N] = M_inv[None, :, None] * D_neg * W[None, None, :]                 # beta
+            no_shortcut_indices.extend(inds.tolist())                                            
+            if m_equals_0 and there_is_iso_source:
+                no_shortcut_indices_0.extend(ns_ind.tolist())                                      
+            if there_is_beam_source:
+                X_temp = (                                                                        
+                    scalar_fac_asso_leg_term_mu0                                                  
+                    * scaled_omega_arr[ns_ind, None]                                                  
+                    * weighted_scaled_Leg_coeffs[ns_ind, m:]                                          
+                )                                                                                
+                # Signs flipped due to the extra minus sign
+                X_arr[inds, :N] =  M_inv * (X_temp @ asso_leg_term_pos)                           
+                X_arr[inds, N:] = -M_inv * (X_temp @ asso_leg_term_neg)                           
 
-                # Assemble the coefficient matrix and additional terms
-                # --------------------------------------------------------------------------------------------------------------------------
-                DW = D_pos * W[None, :]
-                DW.ravel()[::N + 1] -= 1
-                alpha = M_inv[:, None] * DW
-                beta = M_inv[:, None] * D_neg * W[None, :]
-                
-                # --------------------------------------------------------------------------------------------------------------------------
-                
-                # Setup to solve for particular solution w.r.t. direct beam source (see section 3.6.1 of the Comprehensive Documentation)
-                # --------------------------------------------------------------------------------------------------------------------------
-                if there_is_beam_source:
-                    # Generate X
-                    X_temp = (
-                        scalar_fac_asso_leg_term_mu0
-                        * scaled_omega_l 
-                        * weighted_scaled_Leg_coeffs_l
-                    )
-                    X_pos = X_temp @ asso_leg_term_pos
-                    X_neg = X_temp @ asso_leg_term_neg
-                    # Signs flipped due to the extra minus sign
-                    X_arr[ind, :N] = M_inv * X_pos
-                    X_arr[ind, N:] = -M_inv * X_neg
-                    
-                # --------------------------------------------------------------------------------------------------------------------------
-                
-                A_arr[ind, N:, :N] = beta
-                A_arr[ind, N:, N:] = alpha
-                no_shortcut_indices.append(ind)
-                if m_equals_0 and there_is_iso_source: # Keep the list empty if there is no isotropic source
-                    no_shortcut_indices_0.append(l)
-                
-            else:
-                # This is a shortcut to the diagonalization results
-                G_collect[ind, :, :] = G_trivial
-                K_collect[ind, :N] = -M_inv
-                K_collect[ind, N:] = M_inv
-                if m_equals_0 and there_is_iso_source:
-                    G_inv_collect_0[l, :, :] = G_trivial
-                    
-            ind += 1
-                
     if len(no_shortcut_indices) > 0:
         A_arr = A_arr[no_shortcut_indices, :, :]
         alpha_arr = A_arr[:, N:, N:]

{pythonicdisort-1.5 → pythonicdisort-1.7}/src/PythonicDISORT/pydisort.py

@@ -6,7 +6,7 @@ import warnings
 import numpy as np
 import scipy as sc
 from math import pi
-from numpy.polynomial.legendre import Legendre
+from numpy.polynomial.legendre import Legendre, legval
 from scipy import integrate
     
 
@@ -406,68 +406,59 @@ def pydisort(
         
         # TMS correction for the intensity (see section 3.7.2)
         # --------------------------------------------------------------------------------------------------------------------------
-        def TMS_correction(tau, phi, is_antiderivative_wrt_tau, l, l_uniq, l_inv):
+        def TMS_correction(tau, phi, is_antiderivative_wrt_tau):
             Ntau = len(tau)
             Nphi = len(phi)
-            # Atmospheric layer indices
-            if not there_is_iso_source:
-                l_uniq, l_inv = np.unique(l, return_inverse=True)
-            scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
+            # Atmospheric layer indices: smallest l with tau <= tau_arr[l].
+            l = np.searchsorted(tau_arr, tau, side="left")
+            scaled_tau_arr_l = scaled_tau_arr_with_0[l + 1]
             scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
 
             # Delta-M scaling
-            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
+            if np.any(scale_tau != 1):
+                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
             
             # mathscr_B
             # --------------------------------------------------------------------------------------------------------------------------
             nu = subroutines.atleast_2d_append(
                 subroutines.calculate_nu(mu_arr, phi, -mu0, phi0)
             )
-            p_true = np.concatenate(
-                [
-                    f(nu)[:, None, :]
-                    # Iterates over the 0th axis
-                    for f in map(Legendre, weighted_Leg_coeffs_all[l_uniq])
-                ],
-                axis=1,
-            ) # (NQuad, NLayers_uniq, Nphi)
-            p_trun = np.concatenate(
-                [
-                    f(nu)[:, None, :]
-                    # Iterates over the 0th axis
-                    for f in map(Legendre, weighted_scaled_Leg_coeffs[l_uniq])
-                ],
-                axis=1,
-            ) # (NQuad, NLayers_uniq, Nphi)
-            
+            # Vectorized Legendre evaluation: legval broadcasts the trailing-axis
+            # coefficient set, returning (NLayers, NQuad, Nphi). Move the layer axis
+            # to the middle to match the original (NQuad, NLayers, Nphi) layout.
+            # Always uses base numpy's legval (not autograd's): this matches the
+            # original `Legendre(...)(nu)` path, which was also pure numpy.
+            p_true = np.moveaxis(legval(nu, weighted_Leg_coeffs_all.T, tensor=True), 0, 1)
+            p_trun = np.moveaxis(legval(nu, weighted_scaled_Leg_coeffs.T, tensor=True), 0, 1)
             mathscr_B = (
-                (scaled_omega_arr[l_uniq] * I0_div_4pi)[None, :, None]
+                (scaled_omega_arr * I0_div_4pi)[None, :, None]
                 * (mu0 / (mu0 + mu_arr))[:, None, None]
-                * (p_true / (1 - f_arr)[None, l_uniq, None] - p_trun)
-            ) # (NQuad, NLayers_uniq, Nphi)
+                * (p_true / (1 - f_arr[None, :, None]) - p_trun)
+            )
             # --------------------------------------------------------------------------------------------------------------------------
             
             neg_scaled_tau_div_mu0 = -scaled_tau / mu0
             if is_antiderivative_wrt_tau:
 
-                scale_tau_l = scale_tau[l_uniq]
-                neg_scale_tau_l_div_mu0 = -scale_tau_l / mu0
-                scale_tau_l_div_mu_arr_pos = scale_tau_l[None, :] / mu_arr_pos[:, None]
+                neg_scale_tau_div_mu0 = -scale_tau / mu0
+                scale_tau_div_mu_arr_pos = scale_tau[None, :] / mu_arr_pos[:, None]
 
                 TMS_correction_pos = (
-                    np.exp(neg_scaled_tau_div_mu0) / neg_scale_tau_l_div_mu0[l_inv]
+                    np.exp(neg_scaled_tau_div_mu0) / neg_scale_tau_div_mu0[l]
                 )[None, :] - np.exp(
                     (scaled_tau - scaled_tau_arr_l)[None, :] / mu_arr_pos[:, None]
                     - scaled_tau_arr_l[None, :] / mu0
-                ) / scale_tau_l_div_mu_arr_pos[:, l_inv]
+                ) / scale_tau_div_mu_arr_pos[:, l]
                 TMS_correction_neg = (
-                    np.exp(neg_scaled_tau_div_mu0) / neg_scale_tau_l_div_mu0[l_inv]
+                    np.exp(neg_scaled_tau_div_mu0) / neg_scale_tau_div_mu0[l]
                 )[None, :] + np.exp(
                     (scaled_tau_arr_lm1 - scaled_tau)[None, :] / mu_arr_pos[:, None]
                     - scaled_tau_arr_lm1[None, :] / mu0
-                ) / scale_tau_l_div_mu_arr_pos[:, l_inv]
+                ) / scale_tau_div_mu_arr_pos[:, l]
             else:
                 TMS_correction_pos = np.exp(neg_scaled_tau_div_mu0)[None, :] - np.exp(
                     (scaled_tau - scaled_tau_arr_l)[None, :] / mu_arr_pos[:, None]
@@ -477,123 +468,140 @@ def pydisort(
                     (scaled_tau_arr_lm1 - scaled_tau)[None, :] / mu_arr_pos[:, None]
                     - scaled_tau_arr_lm1[None, :] / mu0
                 )
-
             
             ## Contribution from other layers
             # --------------------------------------------------------------------------------------------------------------------------
-            solution = (
-                mathscr_B[:, l_inv, :]
-                * np.concatenate((TMS_correction_pos, TMS_correction_neg), axis=0)[:, :, None]
-            )
+            Bpos = mathscr_B[:N, :, :]                                      # (N, NLayers, Nphi)
+            Bneg = mathscr_B[N:, :, :]                                      # (N, NLayers, Nphi)
             
-            if is_atmos_multilayered:
-                any_pos = (l.min() < NLayers - 1)
-                any_neg = (l.max() > 0)
-
-                mu0_inv = 1 / mu0
-
-                scaled_tauwith0_front = scaled_tau_arr_with_0[:-1]          # (NLayers,)
-                scaled_tauwith0_back  = scaled_tau_arr_with_0[1:]           # (NLayers,); i.e. `scaled_tau_arr`
+            sol_pos = Bpos[:, l, :] * TMS_correction_pos[:, :, None]        # (N, Ntau, Nphi)
+            sol_neg = Bneg[:, l, :] * TMS_correction_neg[:, :, None]        # (N, Ntau, Nphi)
+            
+            if not is_atmos_multilayered:
+                return np.concatenate((sol_pos, sol_neg), axis=0)
 
-                exp_taufront_mu0_inv = np.concatenate(
-                    ([1], np.exp(-scaled_tauwith0_front[1:] * mu0_inv))
-                )  # (Nlayers, )
+            # ---------------- other-layer contributions ----------------
+            any_pos = (l.min() < NLayers - 1)  # exists tau with layers below
+            any_neg = (l.max() > 0)            # exists tau with layers above
 
+            # Layer boundaries/thickness in scaled-tau coordinates
+            exp_scaled_tau_arr_with_0 = prepend(np.exp(-scaled_tau_arr_with_0[1:] / mu0), NLayers, 1)
 
-                Bpos = mathscr_B[:N, :, :]                       # (N, NLayers_uniq, Nphi)
-                Bneg = mathscr_B[N:, :, :]                       # (N, NLayers_uniq, Nphi)
+            # decay[n,r] = exp(-scaled_thickness_arr[r]/mu[n])
+            delta = scaled_thickness_arr[None, :] * M_inv[:, None]
+            decay = np.exp(-delta)  # (N, NLayers)
 
-                # decay[n,r] = exp(-scaled_thickness_arr[r]/mu[n]); always <= 1
-                log_decay = -scaled_thickness_arr[None, :] * M_inv[:, None]                    # (N, NLayers)
+            # tau_front = scaled_tau_arr_with_0[:-1]
+            exp_taufront_mu0_inv = exp_scaled_tau_arr_with_0[:-1]            # (NLayers,)
+            # tau_back = scaled_tau_arr_with_0[1:]
+            exp_tauback_mu0_inv  = exp_scaled_tau_arr_with_0[1:]             # (NLayers,)
 
-                # decay_prod[n,k] = prod_{m=0}^{k-1} decay[n,m], with decay_prod[:,0]=1
-                log_decay_prod = np.cumsum(log_decay, axis=1) # The cumsum requires ALL layers; `l_uniq` is inapplicable
-                decay_prod = np.concatenate((np.ones((N, 1)), np.exp(log_decay_prod)), axis=1) # (N, NLayers + 1)
+            if is_antiderivative_wrt_tau:
+                # (N, NLayers): multiply by this instead of dividing by (scale_tau/mu)
+                mu_div_scale_layers = mu_arr_pos[:, None] / scale_tau[None, :]
+
+            # ===== POS: sum over r > l =====
+            if any_pos:
+                l_min = int(l.min())  # only Acc_pos at layers >= l_min is ever indexed
+
+                # fac_pos = 1 - exp(-scaled_thickness_arr*(1/mu + 1/mu0)) = -expm1(-scaled_thickness_arr*(...))
+                delta_pos = delta + (scaled_thickness_arr / mu0)[None, :]    # (N, NLayers) >= 0
+                fac_pos = -np.expm1(-delta_pos)                              # (N, NLayers) in [0,1)
+
+                # NOTE: antiderivative target-layer factor `mu/scale_tau[l]` is applied
+                # OUTSIDE the accumulator at line 544, not per-source-layer inside it.
+
+                # coef_pos[n,r] multiplies exp((tau - tau_front[r])/mu - tau_front[r]/mu0) after re-referencing
+                # includes exp(-tau_front[r]/mu0)
+                coef_pos = fac_pos * exp_taufront_mu0_inv[None, :]           # (N, NLayers)
+
+                # Reverse scan to build Acc_l (referenced at tau_back[l]) INCLUDING Bpos:
+                # Acc_{r-1} = coef_pos[:,r]*Bpos[:,r,:] + decay[:,r]*Acc_r
+                # Stored array covers layers [l_min..NLayers-1]; layer r maps to index r-l_min.
+                Npos = NLayers - l_min
+
+                if autograd_compatible:
+                    acc0 = np.zeros((N, Nphi))
+                    acc = acc0
+                    cols_rev = [acc0]  # Acc_{NLayers-1} = 0
+                    for r in range(NLayers - 1, l_min, -1):
+                        term = coef_pos[:, r][:, None] * Bpos[:, r, :]        # (N, Nphi)
+                        acc = term + decay[:, r][:, None] * acc
+                        cols_rev.append(acc)
+                    Acc_pos = np.stack(cols_rev[::-1], axis=1)                # (N, Npos, Nphi)
+                else:
+                    Acc_pos = np.empty((N, Npos, Nphi))
+                    Acc_pos[:, Npos - 1, :] = 0.0                             # Acc_{NLayers-1}
+                    for r in range(NLayers - 1, l_min, -1):
+                        i_dest = r - 1 - l_min
+                        i_source = r - l_min
+                        np.multiply(coef_pos[:, r:r+1], Bpos[:, r, :],
+                                    out=Acc_pos[:, i_dest, :])
+                        Acc_pos[:, i_dest, :] += decay[:, r:r+1] * Acc_pos[:, i_source, :]
+
+                # Apply to tau: multiply by exp((scaled_tau - tau_back[l]) / mu) (argument <= 0 in exact arithmetic)
+                expfac_pos = np.exp(M_inv[:, None] * (scaled_tau - scaled_tau_arr_with_0[1:][l])[None, :])  # (N, Ntau)
+                contrib_pos = Acc_pos[:, l - l_min, :] * expfac_pos[:, :, None]   # (N, Ntau, Nphi)
 
                 if is_antiderivative_wrt_tau:
-                    # (N, NLayers)
-                    integration_factor = mu_arr_pos[:, None] / scale_tau
-
-                # ====================================================================
-                # POS: pre-accumulate "contribution from below" sums Rpos[N, l, Nphi]
-                # ====================================================================
-                if any_pos:
-                    thickness_pos = scaled_thickness_arr[None, :] * (M_inv + mu0_inv)[:, None]  # (N, NLayers); always > 0
-                    # em1_pos = 1 - exp(-scaled_thickness_arr*(1/mu + 1/mu0)) = -expm1(-thickness_pos)
-                    em1_pos = -np.expm1(-thickness_pos)                                         # (N, NLayers)
-                    if is_antiderivative_wrt_tau:
-                        em1_pos = integration_factor * em1_pos
-
-                    layer_term_pos = em1_pos * exp_taufront_mu0_inv[None, :]                    # (N, NLayers)
-
-                    # Reverse-scan via cumsum:
-                    # Rpos[l] = (1/decay_prod[l+1]) * sum_{r=l+1..} (layer_term_pos[r] * decay_prod[r])
-                    weighted_layer_term_pos = layer_term_pos * decay_prod[:, :NLayers]          # (N, NLayers)
-                    cumsum_pos = np.cumsum(weighted_layer_term_pos[:, ::-1], axis=1)[:, ::-1]   # (N, NLayers)
-
-                    Rpos_rest = cumsum_pos[:, 1:] / decay_prod[:, 1 : NLayers]                  # (N, NLayers - 1)
-                    Rpos = np.concatenate([Rpos_rest, np.zeros((N, 1))], axis=1)                # (N, NLayers)
-
-                    # Apply to all `tau`
-                    # expfac_pos = exp(M_inv*tau) * exp(-M_inv*scaled_tauwith0_back[l])
-                    expfac_pos = np.exp(M_inv[:, None] * (scaled_tau - scaled_tauwith0_back[l])[None, :])  # (N, Ntau)
-                    tmp = np.take(Rpos, l, axis=1)                                                         # (N, Ntau)
-                    tmp = (tmp * expfac_pos)[:, :, None] * Bpos[:, l_inv, :]
-                    
-                    if autograd_compatible:
-                        solution = np.concatenate((solution[:N, :, :] + tmp, solution[N:, :, :]), axis=0)
-                    else:
-                        solution[:N, :, :] += tmp
-
-                # ====================================================================
-                # NEG: pre-accumulate "contribution from above" sums Rneg[N, l, Nphi]
-                # ====================================================================
-                if any_neg:
-                    decay = np.exp(log_decay)                                                    # (N, NLayers)
-                    
-                    # Ensure no positive exponents
-                    thickness_neg = scaled_thickness_arr[None, :] * (M_inv - mu0_inv)[:, None]   # (N, NLayers)
-                    mask = (thickness_neg >= 0)
-
-                    exp_x1 = np.exp(-scaled_tauwith0_back * mu0_inv)[None, :]                    # (1, NLayers)
-                    exp_x0 = decay * exp_taufront_mu0_inv[None, :]                               # (N, NLayers)
+                    contrib_pos = contrib_pos * mu_div_scale_layers[:, l, None]
+
+                sol_pos = sol_pos + contrib_pos
+
+            # ===== NEG: sum over r < l =====
+            if any_neg:
+                l_max = int(l.max())  # only Acc_neg at layers <= l_max is ever indexed
+
+                # We need term_base[n,r] = exp(-tau_back[r]/mu0) - exp(-scaled_thickness_arr[r]/mu)*exp(-tau_front[r]/mu0)
+                # computed stably via expm1 using only non-positive inputs.
+                delta_neg = delta - (scaled_thickness_arr / mu0)[None, :]     # (N, NLayers), can be +/-.
+                em1 = np.expm1(-np.abs(delta_neg))                            # (N, NLayers) in [-1,0], argument <= 0
+
+                exp_x1 = exp_tauback_mu0_inv[None, :]                         # (1, NLayers)
+                exp_x0 = decay * exp_taufront_mu0_inv[None, :]                # (N, NLayers)
+
+                # if delta>=0: term = exp_x1*(1-exp(-delta)) = exp_x1*(-expm1(-delta)) = exp_x1*(-em1)
+                # if delta<0 : term = exp_x0*(exp(delta)-1)  = exp_x0*expm1(delta)     = exp_x0*( em1)  (since delta<0 -> em1=expm1(delta))
+                term_base = np.where(delta_neg >= 0.0, -em1 * exp_x1, em1 * exp_x0)  # (N, NLayers)
+
+                coef_neg = term_base                                          # (N, NLayers)
+                # NOTE: antiderivative target-layer factor `-mu/scale_tau[l]` is applied
+                # OUTSIDE the accumulator at line 600, not per-source-layer inside it.
+
+                # Forward scan to build Acc_l (referenced at tau_front[l]) INCLUDING Bneg:
+                # Acc_{r+1} = coef_neg[:,r]*Bneg[:,r,:] + decay[:,r]*Acc_r
+                # Stored array covers layers [0..l_max]; layer index unchanged.
+                Nneg = l_max + 1
+
+                if autograd_compatible:
+                    acc0 = np.zeros((N, Nphi))
+                    acc = acc0
+                    cols = [acc0]  # Acc_0 = 0
+                    for r in range(0, l_max):
+                        term = coef_neg[:, r][:, None] * Bneg[:, r, :]          # (N, Nphi)
+                        acc = term + decay[:, r][:, None] * acc
+                        cols.append(acc)
+                    Acc_neg = np.stack(cols, axis=1)                            # (N, Nneg, Nphi)
+                else:
+                    Acc_neg = np.empty((N, Nneg, Nphi))
+                    Acc_neg[:, 0, :] = 0.0                                      # Acc_0
+                    for r in range(0, l_max):
+                        np.multiply(coef_neg[:, r:r+1], Bneg[:, r, :],
+                                    out=Acc_neg[:, r + 1, :])
+                        Acc_neg[:, r + 1, :] += decay[:, r:r+1] * Acc_neg[:, r, :]
+
+                # Apply to tau: multiply by exp((tau_front[l] - scaled_tau) / mu) (argument <= 0 in exact arithmetic)
+                expfac_neg = np.exp(M_inv[:, None] * (scaled_tau_arr_with_0[l] - scaled_tau)[None, :])  # (N, Ntau)
+                contrib_neg = Acc_neg[:, l, :] * expfac_neg[:, :, None]         # (N, Ntau, Nphi)
 
-                    # Let a = |thickness_neg| >= 0. Then expm1(-a) uses only non-positive inputs.
-                    em1_neg = np.expm1(-np.abs(thickness_neg))                                       # (N, NLayers)
-
-                    # thickness_neg>=0: exp_x1 * (1 - exp(-a)) = exp_x1 * (-expm1(-a)) = exp_x1 * (-em1_neg)
-                    # thickness_neg<0 : exp_x0 * (exp(-a) - 1) = exp_x0 * expm1(-a)    = exp_x0 * ( em1_neg)
-                    layer_term_negthick = em1_neg * exp_x0                                       # (N, NLayers)
-                    layer_term_posthick = -em1_neg * exp_x1                                      # (N, NLayers)
-                    layer_term_neg = layer_term_posthick * mask + layer_term_negthick * (~mask)  # (N, NLayers)
-
-                    if is_antiderivative_wrt_tau:
-                        layer_term_neg = -integration_factor * layer_term_neg
+                if is_antiderivative_wrt_tau:
+                    contrib_neg = -contrib_neg * mu_div_scale_layers[:, l, None]
 
-                    # Forward-scan via cumsum:
-                    # Rneg[l] = decay_prod[l] * sum_{r<l} layer_term_neg[r] / decay_prod[r+1]
-                    weighted_layer_term_neg = layer_term_neg / decay_prod[:, 1:]                 # (N, NLayers)
-                    cumsum_neg = np.cumsum(weighted_layer_term_neg, axis=1)                      # (N, NLayers)
-                    
-                    Rneg_rest = decay_prod[:, 1 : NLayers] * cumsum_neg[:, :-1]                  # (N, NLayers - 1)
-                    Rneg = np.concatenate([np.zeros((N, 1)), Rneg_rest], axis=1)                 # (N, NLayers)
+                sol_neg = sol_neg + contrib_neg
 
-                    # Apply to all `tau`
-                    expfac_neg = np.exp(M_inv[:, None] * (scaled_tauwith0_front[l] - scaled_tau)[None, :])  # (N, Ntau)
-                    tmp = np.take(Rneg, l, axis=1)                                                          # (N, Ntau)
-                    tmp = (tmp * expfac_neg)[:, :, None] * Bneg[:, l_inv, :]
+            return np.concatenate((sol_pos, sol_neg), axis=0)
                     
-                    if autograd_compatible:
-                        solution = np.concatenate((solution[:N, :, :], solution[N:, :, :] + tmp), axis=0)
-                    else:
-                        solution[N:, :, :] += tmp
-
-                return solution
-
             # --------------------------------------------------------------------------------------------------------------------------
-            
-            else:
-                return solution
         # --------------------------------------------------------------------------------------------------------------------------
 
         # IMS correction for the intensity (see section 3.7.2)
@@ -628,7 +636,8 @@ def pydisort(
                 ) / (mu_arr_pos * scaled_mu0 * x)[:, None]
                 
             return (
-                I0_div_4pi
+                I0
+                / (4 * pi)
                 * (omega_avg * f_avg) ** 2
                 / (1 - omega_avg * f_avg)
                 * Legendre(
@@ -640,59 +649,25 @@ def pydisort(
 
         # The corrected intensity
         # --------------------------------------------------------------------------------------------------------------------------
-        def u_corrected(
-            tau,
-            phi,
-            is_antiderivative_wrt_tau=False,
-            return_Fourier_error=False,
-            return_tau_arr=False,
-        ):
+        def u_corrected(tau, phi, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False):
             tau = np.atleast_1d(tau)
             phi = np.atleast_1d(phi)
-
-            u_star_outputs = u_star(
-                tau,
-                phi,
-                is_antiderivative_wrt_tau,
-                return_Fourier_error,
-                return_tau_arr,
-                _return_l=True,
-            )
-            if there_is_iso_source:
-                NT_corrections = TMS_correction(
-                    tau,
-                    phi,
-                    is_antiderivative_wrt_tau,
-                    u_star_outputs[-3],
-                    u_star_outputs[-2],
-                    u_star_outputs[-1],
-                )
-                u_star_outputs = u_star_outputs[:-3]
-            else:
-                NT_corrections = TMS_correction(
-                    tau, phi, is_antiderivative_wrt_tau, u_star_outputs[-1], None, None
-                )
-                u_star_outputs = u_star_outputs[:-1]
-
+            NT_corrections = TMS_correction(tau, phi, is_antiderivative_wrt_tau)
+            
             if autograd_compatible:
-                NT_corrections = np.concatenate(
-                    (
-                        NT_corrections[:N, :, :],
-                        NT_corrections[N:, :, :]
-                        + IMS_correction(tau, phi, is_antiderivative_wrt_tau),
-                    ),
-                    axis=0,
+                NT_corrections = NT_corrections + np.concatenate(
+                    [np.zeros((N, len(tau), len(phi))), IMS_correction(tau, phi, is_antiderivative_wrt_tau)], axis=0
                 )
             else:
                 NT_corrections[N:, :, :] += IMS_correction(tau, phi, is_antiderivative_wrt_tau)
-
+            
             if return_Fourier_error or return_tau_arr:
+                u_star_outputs = u_star(tau, phi, is_antiderivative_wrt_tau, return_Fourier_error, return_tau_arr)
                 return (
                     u_star_outputs[0] + rescale_factor * np.squeeze(NT_corrections),
                 ) + u_star_outputs[1:]
             else:
-                return u_star_outputs[0] + rescale_factor * np.squeeze(NT_corrections)
-
+                return u_star(tau, phi, is_antiderivative_wrt_tau) + rescale_factor * np.squeeze(NT_corrections)
         # --------------------------------------------------------------------------------------------------------------------------
 
         return mu_arr, flux_up, flux_down, u0, u_corrected

{pythonicdisort-1.5 → pythonicdisort-1.7}/src/PythonicDISORT/subroutines.py

@@ -641,60 +641,76 @@ def interpolate(u):
     u_pos_interpol = sc.interpolate.BarycentricInterpolator(mu_arr_pos)
     u_neg_interpol = sc.interpolate.BarycentricInterpolator(-mu_arr_pos)
     
+    # Value-keyed cache for the most recent (tau, phi, is_antideriv) -> u_cache.
+    # Hits when the user loops over `mu` with the same (tau, phi); otherwise re-evaluates `u`.
+    # Cache stores copies of `tau`/`phi` so in-place mutation of the user's arrays
+    # cannot return stale values (compared via `np.array_equal`).
+    # Only used on the simple-return path (skipped when Fourier_error / tau_arr are requested).
+    _cache = [None, None, None, None]  # [tau_copy, phi_copy, is_antideriv, u_cache]
+
     argcount = u.__code__.co_argcount
     if argcount == 5: # Function is `u` instead of `u0`
         def u_interpol(mu, tau, phi, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False):
             if not np.all(np.abs(mu) <= 1):
                 raise ValueError("mu values must be between -1 and 1.")
-            
+
             mu = np.atleast_1d(mu)
 
             if return_Fourier_error or return_tau_arr:
                 u_outputs = u(tau, phi, is_antiderivative_wrt_tau, return_Fourier_error, return_tau_arr)
                 u_cache = u_outputs[0]
+            elif (_cache[3] is not None and is_antiderivative_wrt_tau == _cache[2]
+                  and np.array_equal(tau, _cache[0]) and np.array_equal(phi, _cache[1])):
+                u_cache = _cache[3]
             else:
                 u_cache = u(tau, phi, is_antiderivative_wrt_tau)
-            
+                _cache[:] = [np.array(tau, copy=True), np.array(phi, copy=True),
+                             is_antiderivative_wrt_tau, u_cache]
+
             results = np.empty((len(mu),) + np.shape(u_cache)[1:])
             mask_pos = mu > 0
             mask_else = ~mask_pos
-            
+
             if np.any(mask_pos):
                 u_pos_interpol.set_yi(u_cache[:N])
                 results[mask_pos] = u_pos_interpol(mu[mask_pos])
             if np.any(mask_else):
                 u_neg_interpol.set_yi(u_cache[N:])
                 results[mask_else] = u_neg_interpol(mu[mask_else])
-            
+
             if return_Fourier_error or return_tau_arr:
                 return (np.squeeze(results)[()],) + u_outputs[1:]
             else:
                 return np.squeeze(results)[()]
-    
+
     elif argcount == 4: # Function is u0 instead of u
         def u_interpol(mu, tau, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False):
             if not np.all(np.abs(mu) <= 1):
                 raise ValueError("mu values must be between -1 and 1.")
-            
+
             mu = np.atleast_1d(mu)
 
             if return_Fourier_error or return_tau_arr:
                 u_outputs = u(tau, is_antiderivative_wrt_tau, return_Fourier_error, return_tau_arr)
                 u_cache = u_outputs[0]
+            elif (_cache[3] is not None and is_antiderivative_wrt_tau == _cache[2]
+                  and np.array_equal(tau, _cache[0])):
+                u_cache = _cache[3]
             else:
                 u_cache = u(tau, is_antiderivative_wrt_tau)
-            
+                _cache[:] = [np.array(tau, copy=True), None, is_antiderivative_wrt_tau, u_cache]
+
             results = np.empty((len(mu),) + np.shape(u_cache)[1:])
             mask_pos = mu > 0
             mask_else = ~mask_pos
-            
+
             if np.any(mask_pos):
                 u_pos_interpol.set_yi(u_cache[:N])
                 results[mask_pos] = u_pos_interpol(mu[mask_pos])
             if np.any(mask_else):
                 u_neg_interpol.set_yi(u_cache[N:])
                 results[mask_else] = u_neg_interpol(mu[mask_else])
-            
+
             if return_Fourier_error or return_tau_arr:
                 return (np.squeeze(results)[()],) + u_outputs[1:]
             else:
@@ -764,7 +780,7 @@ def _mathscr_v(tau,                             # Input optical depths
     |          Variable             |                   Shape                   |
     | ----------------------------- | ----------------------------------------- |
     | ``tau``                       | ``Ntau``                                  |
-    | ``scale_tau``                 | `NLayers`                                 |
+    | ``scale_tau``                 | ``NL``                                |
     | ``l``                         | ``Ntau``                                  |
     | ``Nscoeffs``                  | scalar                                    |
     | ``s_poly_coeffs``             | ``NL x Nscoeffs``                         |
@@ -790,11 +806,10 @@ def _mathscr_v(tau,                             # Input optical depths
         d0 = np.einsum("lij, lj -> il", G, c0w, optimize=True)
 
         if is_antiderivative_wrt_tau:
-            p0 = tau / scale_tau                 
+            p0 = tau
+            return (d0 / scale_tau[None, :])[:, l] * p0[None, :]
         else:
-            p0 = np.ones_like(tau)       
-
-        return d0[:, l] * p0[None, :]
+            return d0[:, l]
 
     # --------------------- Nscoeffs == 2 ---------------------
     elif Nscoeffs == 2:
@@ -812,11 +827,11 @@ def _mathscr_v(tau,                             # Input optical depths
         if is_antiderivative_wrt_tau:
             p0 = tau * tau / 2
             p1 = tau
-            return (d0[:, l] * p0[None, :] + d1[:, l] * p1[None, :]) / scale_tau
+            return (d0 / scale_tau[None, :])[:, l] * p0[None, :] + (d1 / scale_tau[None, :])[:, l] * p1[None, :]
         else:
             p0 = tau
-            p1 = np.ones_like(tau) 
-            return d0[:, l] * p0[None, :] + d1[:, l] * p1[None, :]
+            #p1 = np.ones_like(tau) 
+            return d0[:, l] * p0[None, :] + d1[:, l]
         
     # --------------------- Nscoeffs > 2 ----------------------
     else:
@@ -855,11 +870,11 @@ def _mathscr_v(tau,                             # Input optical depths
         powers = np.arange(Nscoeffs - 1, -1, -1)[None, :]
         if is_antiderivative_wrt_tau:
             p = powers + 1
-            tau_poly = (tau[:, None] ** p) / (p * scale_tau)
+            tau_poly = (tau[:, None] ** p) / p
+            return np.einsum("tqp, tp -> qt", (mathscr_b / scale_tau[:, None, None])[l, :, :], tau_poly, optimize=True)
         else:
             tau_poly = tau[:, None] ** powers
-
-        return np.einsum("tqp, tp -> qt", mathscr_b[l, :, :], tau_poly, optimize=True)
+            return np.einsum("tqp, tp -> qt", mathscr_b[l, :, :], tau_poly, optimize=True)
 
 
 

{pythonicdisort-1.5 → pythonicdisort-1.7/src/PythonicDISORT.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PythonicDISORT
-Version: 1.5
+Version: 1.7
 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
@@ -41,7 +41,7 @@ You may contact me, Dion, through dh3065@columbia.edu.
 
 The **GitHub repository** is https://github.com/LDEO-CREW/Pythonic-DISORT.
 
-Accompanying Journal of Open Source Software paper: https://joss.theoj.org/papers/10.21105/joss.06442.
+Accompanying Journal of Open Source Software paper: https://joss.theoj.org/papers/10.21105/joss.06442. **Please cite this paper if you use PythonicDISORT**.
 
 # Documentation
 https://pythonic-disort.readthedocs.io/en/latest/
@@ -76,7 +76,7 @@ These notebooks double up as examples of how to use PythonicDISORT. The tests wh
 # Installation
 
 * From PyPI: `pip install PythonicDISORT`
-* From Conda-forge: (TODO: need to first publish on Conda-forge)
+* From Conda-Forge: `conda install -c conda-forge pythonicdisort` or `mamba install -c conda-forge pythonicdisort`
 * By cloning repository: `pip install .` in the `Pythonic-DISORT` directory; `pip install -r all_optional_dependencies.txt` to install all optional dependencies (see *Requirements to run PythonicDISORT*)
 
 ## Requirements to run PythonicDISORT