PythonicDISORT 1.6__tar.gz → 1.7__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PythonicDISORT
-Version: 1.6
+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

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

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

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

@@ -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],
@@ -339,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
@@ -404,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],
@@ -448,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
@@ -463,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, :],
@@ -530,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
@@ -545,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:, :],
@@ -570,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.6 → 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.6 → 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.6 → 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
     
 
@@ -409,13 +409,13 @@ def pydisort(
         def TMS_correction(tau, phi, is_antiderivative_wrt_tau):
             Ntau = len(tau)
             Nphi = len(phi)
-            # Atmospheric layer indices
-            l = np.argmax(tau[:, None] <= tau_arr[None, :], axis=1)
+            # 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
-            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
@@ -427,25 +427,15 @@ def pydisort(
             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)
-                ],
-                axis=1,
-            )
-            p_trun = np.concatenate(
-                [
-                    f(nu)[:, None, :]
-                    # Iterates over the 0th axis
-                    for f in map(Legendre, weighted_scaled_Leg_coeffs)
-                ],
-                axis=1,
-            )
+            # 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 * I0)[None, :, None]
-                / (4 * np.pi)
+                (scaled_omega_arr * I0_div_4pi)[None, :, None]
                 * (mu0 / (mu0 + mu_arr))[:, None, None]
                 * (p_true / (1 - f_arr[None, :, None]) - p_trun)
             )
@@ -495,14 +485,16 @@ def pydisort(
             any_neg = (l.max() > 0)            # exists tau with layers above
 
             # Layer boundaries/thickness in scaled-tau coordinates
-            tau_front = scaled_tau_arr_with_0[:-1]                           # (NLayers,) = tau_r
-            tau_back  = scaled_tau_arr_with_0[1:]                            # (NLayers,) = tau_{r+1}
+            exp_scaled_tau_arr_with_0 = prepend(np.exp(-scaled_tau_arr_with_0[1:] / mu0), NLayers, 1)
 
-            # decay[n,r] = exp(-scaled_thickness_arr[r]/mu[n])  (argument <= 0; underflow is safe)
-            decay = np.exp(-scaled_thickness_arr[None, :] * M_inv[:, None])                  # (N, NLayers)
+            # decay[n,r] = exp(-scaled_thickness_arr[r]/mu[n])
+            delta = scaled_thickness_arr[None, :] * M_inv[:, None]
+            decay = np.exp(-delta)  # (N, NLayers)
 
-            exp_taufront_mu0_inv = np.exp(-tau_front / mu0)              # (NLayers,)
-            exp_tauback_mu0_inv  = np.exp(-tau_back / mu0)              # (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,)
 
             if is_antiderivative_wrt_tau:
                 # (N, NLayers): multiply by this instead of dividing by (scale_tau/mu)
@@ -510,73 +502,100 @@ def pydisort(
 
             # ===== 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 = scaled_thickness_arr[None, :] * (M_inv + 1 / mu0)[:, None]        # (N, NLayers) >= 0
-                fac_pos = -np.expm1(-delta_pos)                               # (N, NLayers) in [0,1)
+                delta_pos = delta + (scaled_thickness_arr / mu0)[None, :]    # (N, NLayers) >= 0
+                fac_pos = -np.expm1(-delta_pos)                              # (N, NLayers) in [0,1)
 
-                if is_antiderivative_wrt_tau:
-                    fac_pos = fac_pos * mu_div_scale_layers                  # (N, NLayers)
+                # 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)
+                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
-                acc0 = np.zeros((N, Nphi))
-                acc = acc0
-                cols_rev = [acc0]  # Acc_{NLayers-1} = 0
-
-                for r in range(NLayers - 1, 0, -1):  # r = NLayers-1 ... 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, NLayers, Nphi)
+                # 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 - tau_back[l])[None, :])  # (N, Ntau)
-                contrib_pos = Acc_pos[:, l, :] * expfac_pos[:, :, None]         # (N, Ntau, Nphi)
+                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:
+                    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 = scaled_thickness_arr[None, :] * (M_inv - 1 / mu0)[:, None]        # (N, NLayers), can be +/-.
-                abs_delta = np.abs(delta_neg)
-                em1 = np.expm1(-abs_delta)                                    # (N, NLayers) in [-1,0], argument <= 0
+                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)
+                term_base = np.where(delta_neg >= 0.0, -em1 * exp_x1, em1 * exp_x0)  # (N, NLayers)
 
-                if is_antiderivative_wrt_tau:
-                    # original divides by (-scale_tau/mu) == multiply by (-(mu/scale_tau))
-                    coef_neg = term_base * -mu_div_scale_layers               # (N, NLayers)
-                else:
-                    coef_neg = term_base                                      # (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
-                acc0 = np.zeros((N, Nphi))
-                acc = acc0
-                cols = [acc0]  # Acc_0 = 0
-
-                for r in range(0, NLayers - 1):  # r = 0 ... NLayers-2
-                    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, NLayers, Nphi)
+                # 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] * (tau_front[l] - scaled_tau)[None, :])  # (N, Ntau)
-                contrib_neg = Acc_neg[:, l, :] * expfac_neg[:, :, None]          # (N, Ntau, Nphi)
+                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)
+
+                if is_antiderivative_wrt_tau:
+                    contrib_neg = -contrib_neg * mu_div_scale_layers[:, l, None]
 
                 sol_neg = sol_neg + contrib_neg
 

{pythonicdisort-1.6 → 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.6 → pythonicdisort-1.7/src/PythonicDISORT.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PythonicDISORT
-Version: 1.6
+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