PythonicDISORT 1.6__tar.gz → 1.8__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PythonicDISORT
-Version: 1.6
+Version: 1.8
 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
@@ -26,15 +26,15 @@ Requires-Dist: matplotlib>=3.6.0; extra == "notebook-dependencies"
 # Introduction
 The PythonicDISORT package is a Discrete Ordinates Solver for the (1D) Radiative Transfer Equation 
 in a plane-parallel, horizontally homogeneous atmosphere.
-It is coded entirely in Python 3 and is a reimplementation instead of a wrapper.
-While PythonicDISORT has been optimized for speed, it will naturally be slower than similar FORTRAN algorithms.
-On the other hand, PythonicDISORT should be easier to install, use, and modify than FORTRAN-based Discrete Ordinates Solvers.
-
-PythonicDISORT is based on Stamnes' FORTRAN DISORT (see References, in particular [2, 3, 8]) and has its main features: multi-layer solver, 
-delta-_M_ scaling, Nakajima-Tanaka (NT) corrections, only flux option, direct beam source, isotropic internal source (blackbody emission), 
-Dirichlet boundary conditions (diffuse flux boundary sources), Bi-Directional Reflectance Function (BDRF) for surface reflection,
-and interpolation with respect to polar angle.
-In addition, we added a subroutine to compute actinic fluxes to satisfy a user request, and integration with respect to optical depth was also added.
+It is coded entirely in Python 3 and is not a wrapper.
+PythonicDISORT has been optimized for speed and **as of v1.8 is one-third the speed of FORTRAN DISORT while being NumPy and Scipy only**.
+PythonicDISORT is intended to be easier to install, use, and modify than FORTRAN-based Discrete Ordinates Solvers.
+
+PythonicDISORT is based on Stamnes' FORTRAN DISORT (see References, in particular [2, 3, 8]),
+as of v1.8 the only **missing features** are $\delta$-M+ scaling (PythonicDISORT instead implements NT corrections), 
+the Chapman function for near-horizon solar angles, and pseudo-spherical corrections. 
+On the other hand, we added a subroutine to compute actinic fluxes to satisfy a user request, 
+added analytic integration with respect to optical depth, and improved the handling of singularities.
 Further feature requests as well as feedback are welcome.
 
 You may contact me, Dion, through dh3065@columbia.edu.
@@ -70,8 +70,11 @@ These notebooks double up as examples of how to use PythonicDISORT. The tests wh
 * Test Problem 7: Absorption + Scattering + All Possible Sources, Lambertian and Hapke Surface Reflectivities (one layer)
 * Test Problem 8: Absorbing / Isotropic-Scattering Medium (multiple layers)
 * Test Problem 9: General Emitting / Absorbing / Scattering Medium (multiple layers)
-* Test Problem 11: Single-Layer vs. Multiple Layers (no corresponding Jupyter Notebook)
-* Test Problem I: Antiderivative / integration functionality (no corresponding Jupyter Notebook)
+
+(no corresponding Jupyter Notebook)
+* Test Problem 11: Single-Layer vs. Multiple Layers
+* Test Problem I: Antiderivative / integration functionality
+* Test Problem Interp: Interpolation + NT Correction functionality
 
 # Installation
 

{pythonicdisort-1.6 → pythonicdisort-1.8}/README.md

@@ -4,15 +4,15 @@
 # Introduction
 The PythonicDISORT package is a Discrete Ordinates Solver for the (1D) Radiative Transfer Equation 
 in a plane-parallel, horizontally homogeneous atmosphere.
-It is coded entirely in Python 3 and is a reimplementation instead of a wrapper.
-While PythonicDISORT has been optimized for speed, it will naturally be slower than similar FORTRAN algorithms.
-On the other hand, PythonicDISORT should be easier to install, use, and modify than FORTRAN-based Discrete Ordinates Solvers.
-
-PythonicDISORT is based on Stamnes' FORTRAN DISORT (see References, in particular [2, 3, 8]) and has its main features: multi-layer solver, 
-delta-_M_ scaling, Nakajima-Tanaka (NT) corrections, only flux option, direct beam source, isotropic internal source (blackbody emission), 
-Dirichlet boundary conditions (diffuse flux boundary sources), Bi-Directional Reflectance Function (BDRF) for surface reflection,
-and interpolation with respect to polar angle.
-In addition, we added a subroutine to compute actinic fluxes to satisfy a user request, and integration with respect to optical depth was also added.
+It is coded entirely in Python 3 and is not a wrapper.
+PythonicDISORT has been optimized for speed and **as of v1.8 is one-third the speed of FORTRAN DISORT while being NumPy and Scipy only**.
+PythonicDISORT is intended to be easier to install, use, and modify than FORTRAN-based Discrete Ordinates Solvers.
+
+PythonicDISORT is based on Stamnes' FORTRAN DISORT (see References, in particular [2, 3, 8]),
+as of v1.8 the only **missing features** are $\delta$-M+ scaling (PythonicDISORT instead implements NT corrections), 
+the Chapman function for near-horizon solar angles, and pseudo-spherical corrections. 
+On the other hand, we added a subroutine to compute actinic fluxes to satisfy a user request, 
+added analytic integration with respect to optical depth, and improved the handling of singularities.
 Further feature requests as well as feedback are welcome.
 
 You may contact me, Dion, through dh3065@columbia.edu.
@@ -48,8 +48,11 @@ These notebooks double up as examples of how to use PythonicDISORT. The tests wh
 * Test Problem 7: Absorption + Scattering + All Possible Sources, Lambertian and Hapke Surface Reflectivities (one layer)
 * Test Problem 8: Absorbing / Isotropic-Scattering Medium (multiple layers)
 * Test Problem 9: General Emitting / Absorbing / Scattering Medium (multiple layers)
-* Test Problem 11: Single-Layer vs. Multiple Layers (no corresponding Jupyter Notebook)
-* Test Problem I: Antiderivative / integration functionality (no corresponding Jupyter Notebook)
+
+(no corresponding Jupyter Notebook)
+* Test Problem 11: Single-Layer vs. Multiple Layers
+* Test Problem I: Antiderivative / integration functionality
+* Test Problem Interp: Interpolation + NT Correction functionality
 
 # Installation
 

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

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

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

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