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

PythonicDISORT/_assemble_intensity_and_fluxes.py

@@ -26,7 +26,7 @@ def _assemble_intensity_and_fluxes(
     scale_tau,                          # Delta-scale factor for tau
     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?
+    autograd_compatible,                # Should the output functions be compatible with autograd?
 ):  
     """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. 
@@ -73,7 +73,7 @@ def _assemble_intensity_and_fluxes(
     | `scale_tau`                    | `NLayers`                          |
     | `only_flux`                    | boolean                            |
     | `use_banded_solver_NLayers`    | scalar                             |
-    | `_autograd_compatible`         | boolean                            |
+    | `autograd_compatible`          | boolean                            |
     
     Notable internal variables of _assemble_intensity_and_fluxes
     |     Variable      |             Type / Shape               |
@@ -92,7 +92,7 @@ def _assemble_intensity_and_fluxes(
     | `GC_pos`          | `NLayers x NQuad/2 x NQuad`            |
     | `GC_neg`          | `NLayers x NQuad/2 x NQuad`            |
     """
-    if _autograd_compatible:
+    if autograd_compatible:
         import autograd.numpy as np
     else:
         import numpy as np
@@ -114,20 +114,21 @@ def _assemble_intensity_and_fluxes(
         Nscoeffs,
         there_is_iso_source,
     )
-    if there_is_beam_source:
-        if there_is_iso_source:
-            G_collect, K_collect, B_collect, G_inv_collect_0 = outputs
-        else:
-            G_collect, K_collect, B_collect = outputs
-            G_inv_collect_0 = None
+    if there_is_beam_source and there_is_iso_source:
+        G_collect, K_collect, B_collect, G_inv_collect_0 = outputs
+        B_collect_0 = B_collect[0, :, :]
+    elif there_is_beam_source and not there_is_iso_source:
+        G_collect, K_collect, B_collect = outputs
         B_collect_0 = B_collect[0, :, :]
+        G_inv_collect_0 = None
+    elif not there_is_beam_source and there_is_iso_source:
+        G_collect, K_collect, G_inv_collect_0 = outputs
+        B_collect = None
     else:
-        if there_is_iso_source:
-            G_collect, K_collect, G_inv_collect_0 = outputs
-        else:
-            G_collect, K_collect = outputs
-            G_inv_collect_0 = None
+        G_collect, K_collect = outputs
         B_collect = None
+        G_inv_collect_0 = None
+        
             
     GC_collect = _solve_for_coeffs(
         NFourier,
@@ -161,7 +162,7 @@ def _assemble_intensity_and_fluxes(
  
         # Construct the intensity function (refer to Section 3.7 of the Comprehensive Documentation)
         # --------------------------------------------------------------------------------------------------------------------------
-        def u(tau, phi, return_Fourier_error=False, is_antiderivative_wrt_tau=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.
@@ -172,8 +173,8 @@ def _assemble_intensity_and_fluxes(
             """
             tau = np.atleast_1d(tau)
             phi = np.atleast_1d(phi)
-            assert np.all(tau >= 0)
-            assert np.all(tau <= tau_arr[-1])
+            if np.all(tau < 0) or np.all(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)
@@ -196,34 +197,32 @@ def _assemble_intensity_and_fluxes(
                 axis=-1,
             )
             
-            if is_antiderivative_wrt_tau:
-                if there_is_beam_source:
-                    um = (
-                        np.einsum(
-                            "mtij, mtj -> mti",
-                            GC_collect[:, l, :, :],
-                            np.exp(exponent) / (scale_tau[None, :, None] * K_collect)[:, l, :],
-                            optimize=True,
-                        )
-                        + (B_collect / (-scale_tau / mu0)[None, :, None])[:, l, :]
-                        * np.exp(-scaled_tau / mu0)[None, :, None]
-                    )
-                else:
-                    um = np.einsum(
+            if is_antiderivative_wrt_tau and there_is_beam_source:
+                um = (
+                    np.einsum(
                         "mtij, mtj -> mti",
                         GC_collect[:, l, :, :],
-                        np.exp(exponent) / (scale_tau[None, l, None] * K_collect)[:, l, :],
+                        np.exp(exponent) / (scale_tau[None, :, None] * K_collect)[:, l, :],
                         optimize=True,
                     )
+                    + (B_collect / (-scale_tau / mu0)[None, :, None])[:, l, :]
+                    * np.exp(-scaled_tau / mu0)[None, :, None]
+                )
+            elif is_antiderivative_wrt_tau and not there_is_beam_source:
+                um = np.einsum(
+                    "mtij, mtj -> mti",
+                    GC_collect[:, l, :, :],
+                    np.exp(exponent) / (scale_tau[None, l, None] * K_collect)[:, l, :],
+                    optimize=True,
+                )
+            elif not is_antiderivative_wrt_tau and there_is_beam_source:
+                um = np.einsum(
+                    "mtij, mtj -> mti", GC_collect[:, l, :, :], np.exp(exponent), optimize=True
+                ) + B_collect[:, l, :] * np.exp(-scaled_tau / mu0)[None, :, None]
             else:
-                if there_is_beam_source:
-                    um = np.einsum(
-                        "mtij, mtj -> mti", GC_collect[:, l, :, :], np.exp(exponent), optimize=True
-                    ) + B_collect[:, l, :] * np.exp(-scaled_tau / mu0)[None, :, None]
-                else:
-                    um = np.einsum(
-                        "mtij, mtj -> mti", GC_collect[:, l, :, :], np.exp(exponent), optimize=True
-                    )
+                um = np.einsum(
+                    "mtij, mtj -> mti", GC_collect[:, l, :, :], np.exp(exponent), optimize=True
+                )
             
             # Contribution from particular solution for isotropic internal sources
             if there_is_iso_source:
@@ -236,18 +235,17 @@ def _assemble_intensity_and_fluxes(
                     G_inv_collect_0[l_uniq],
                     mu_arr,
                     is_antiderivative_wrt_tau,
-                    _autograd_compatible
+                    autograd_compatible
                 )
                 
-                if _autograd_compatible:
-                    if NFourier == 1:
-                        um = um + _mathscr_v_contribution.T
-                    else:
-                        um = um + np.concatenate((_mathscr_v_contribution.T, np.zeros((NFourier - 1, len(tau), NQuad))))
+                if NFourier == 1:
+                    um = um + _mathscr_v_contribution.T
+                elif autograd_compatible:
+                    um = um + np.concatenate((_mathscr_v_contribution.T, np.zeros((NFourier - 1, len(tau), NQuad))))
                 else:
                     um[0, :, :] += _mathscr_v_contribution.T
                 
-            intensities = np.einsum(
+            u = np.einsum(
                 "mti, mp -> itp",
                 um,
                 np.cos(np.arange(NFourier)[:, None] * (phi0 - phi)[None, :]),
@@ -262,54 +260,59 @@ def _assemble_intensity_and_fluxes(
                     ],
                     axis=-1,
                 )
-                if is_antiderivative_wrt_tau:
-                    if there_is_beam_source:
-                        ulast = (
-                            np.einsum(
-                                "tij, tj -> it",
-                                GC_collect[-1, l, :, :],
-                                np.exp(exponent) / (scale_tau[:, None] * K_collect[-1, :, :])[l, :],
-                                optimize=True,
-                            )
-                            + (B_collect.T[-1, :, :] / (-scale_tau / mu0)[None, :])[:, l]
-                            * np.exp(-scaled_tau / mu0)[None, :]
-                        )
-                    else:
-                        ulast = np.einsum(
+                if is_antiderivative_wrt_tau and there_is_beam_source:
+                    ulast = (
+                        np.einsum(
                             "tij, tj -> it",
                             GC_collect[-1, l, :, :],
                             np.exp(exponent) / (scale_tau[:, None] * K_collect[-1, :, :])[l, :],
                             optimize=True,
                         )
-                else:
-                    if there_is_beam_source:
-                        ulast = np.einsum(
-                            "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
-                        ) + B_collect.T[-1, :, l] * np.exp(-scaled_tau / mu0)[None, :]
-                    else:
-                        ulast = np.einsum(
-                            "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
+                        + (B_collect.T[-1, :, :] / (-scale_tau / mu0)[None, :])[:, l]
+                        * np.exp(-scaled_tau / mu0)[None, :]
+                    )
+                elif is_antiderivative_wrt_tau and not there_is_beam_source:
+                    ulast = np.einsum(
+                        "tij, tj -> it",
+                        GC_collect[-1, l, :, :],
+                        np.exp(exponent) / (scale_tau[:, None] * K_collect[-1, :, :])[l, :],
+                        optimize=True,
                         )
+                elif not is_antiderivative_wrt_tau and not there_is_beam_source:
+                    ulast = np.einsum(
+                        "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
+                    ) + B_collect.T[-1, :, l] * np.exp(-scaled_tau / mu0)[None, :]
+                else:
+                    ulast = np.einsum(
+                        "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
+                    )
                         
-                intensities_abs = np.abs(intensities)
+                u_abs = np.abs(u)
                 Fourier_error = np.max(
                     np.divide(
                         np.abs(
                             ulast[:, :, None] * np.cos((NFourier - 1) * (phi0 - phi))[None, None, :]
                         ),
-                        intensities_abs,
-                        out=np.zeros_like(intensities_abs),
-                        where=intensities_abs > 1e-8,
+                        u_abs,
+                        out=np.zeros_like(u_abs),
+                        where=u_abs > 1e-15,
                     )
                 )
-                return I0_orig * np.squeeze(intensities), Fourier_error
+                
+                if return_tau_arr:
+                    return I0_orig * np.squeeze(u), Fourier_error, tau_arr
+                else:
+                    return I0_orig * np.squeeze(u), Fourier_error
             else:
-                return I0_orig * np.squeeze(intensities)
+                if return_tau_arr:
+                    return I0_orig * np.squeeze(u), tau_arr
+                else:
+                    return I0_orig * np.squeeze(u)
         # --------------------------------------------------------------------------------------------------------------------------
     
     # Construct u0
     # --------------------------------------------------------------------------------------------------------------------------
-    def u0(tau, is_antiderivative_wrt_tau=False, _return_act_dscale_for_reclass=False):
+    def u0(tau, is_antiderivative_wrt_tau=False, return_tau_arr=False, _return_act_dscale_for_reclass=False):
         """        
         Zeroth Fourier mode of the intensity with argument `tau` (type: array or float).
         Returns an ndarray with axes corresponding to variation with `mu` and `tau` respectively.
@@ -320,8 +323,8 @@ def _assemble_intensity_and_fluxes(
         to switch to an antiderivative of the function with respect to `tau`.
         """
         tau = np.atleast_1d(tau)
-        assert np.all(tau >= 0)
-        assert np.all(tau <= tau_arr[-1])
+        if np.all(tau < 0) or np.all(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)
@@ -356,34 +359,32 @@ def _assemble_intensity_and_fluxes(
             axis=-1,
         )
         
-        if is_antiderivative_wrt_tau:
-            if there_is_beam_source:
-                u0 = (
-                    np.einsum(
-                        "tij, tj -> it",
-                        GC_collect_0[l, :, :],
-                        np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :],
-                        optimize=True,
-                    )
-                    + (B_collect_0.T / (-scale_tau / mu0)[None, :])[:, l]
-                    * np.exp(-scaled_tau / mu0)[None, :]
-                )
-            else:
-                u0 = np.einsum(
+        if is_antiderivative_wrt_tau and there_is_beam_source:
+            u0 = (
+                np.einsum(
                     "tij, tj -> it",
                     GC_collect_0[l, :, :],
                     np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :],
                     optimize=True,
                 )
+                + (B_collect_0.T / (-scale_tau / mu0)[None, :])[:, l]
+                * np.exp(-scaled_tau / mu0)[None, :]
+            )
+        if is_antiderivative_wrt_tau and not there_is_beam_source:
+            u0 = np.einsum(
+                "tij, tj -> it",
+                GC_collect_0[l, :, :],
+                np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :],
+                optimize=True,
+            )
+        if not is_antiderivative_wrt_tau and there_is_beam_source:
+            u0 = np.einsum(
+                "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
+            ) + B_collect_0.T[:, l] * np.exp(-scaled_tau / mu0)[None, :]
         else:
-            if there_is_beam_source:
-                u0 = np.einsum(
-                    "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
-                ) + B_collect_0.T[:, l] * np.exp(-scaled_tau / mu0)[None, :]
-            else:
-                u0 = np.einsum(
-                    "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
-                )
+            u0 = np.einsum(
+                "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
+            )
         
         # Contribution from particular solution for isotropic internal sources
         if there_is_iso_source:
@@ -396,11 +397,15 @@ def _assemble_intensity_and_fluxes(
                 G_inv_collect_0[l_uniq],
                 mu_arr,
                 is_antiderivative_wrt_tau,
-                _autograd_compatible
+                autograd_compatible
             )
             u0 = u0 + _mathscr_v_contribution
-        
-        if _return_act_dscale_for_reclass:
+            
+        if return_tau_arr and _return_act_dscale_for_reclass:
+            return I0_orig * np.squeeze(u0), tau_arr, act_dscale_reclassification
+        elif return_tau_arr and not _return_act_dscale_for_reclass:
+            return I0_orig * np.squeeze(u0), tau_arr
+        elif not return_tau_arr and _return_act_dscale_for_reclass:
             return I0_orig * np.squeeze(u0), act_dscale_reclassification
         else:
             return I0_orig * np.squeeze(u0)
@@ -415,7 +420,7 @@ def _assemble_intensity_and_fluxes(
         B_neg = B_collect_0[:, N:].T
 
 
-    def flux_up(tau, is_antiderivative_wrt_tau=False):
+    def flux_up(tau, is_antiderivative_wrt_tau=False, return_tau_arr=False):
         """    
         (Energetic) Flux function with argument `tau` (type: array or float) for positive (upward) `mu` values.
         Returns the diffuse flux magnitudes (same type and size as `tau`).
@@ -423,8 +428,8 @@ def _assemble_intensity_and_fluxes(
         to switch to an antiderivative of the function with respect to `tau`.
         """
         tau = np.atleast_1d(tau)
-        assert np.all(tau >= 0)
-        assert np.all(tau <= tau_arr[-1])
+        if np.all(tau < 0) or np.all(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)
@@ -449,18 +454,17 @@ def _assemble_intensity_and_fluxes(
                 G_inv_collect_0[l_uniq],
                 mu_arr,
                 is_antiderivative_wrt_tau,
-                _autograd_compatible
+                autograd_compatible
             )
         else:
             _mathscr_v_contribution = 0
         
-        if there_is_beam_source:
-            if is_antiderivative_wrt_tau:
-                direct_beam_contribution = (
-                    (B_pos / (-scale_tau / mu0)[None, :])[:, l] * np.exp(-scaled_tau / mu0)[None, :] 
-                )
-            else:
-                direct_beam_contribution = B_pos[:, l] * np.exp(-scaled_tau / mu0)[None, :]
+        if there_is_beam_source and is_antiderivative_wrt_tau:
+            direct_beam_contribution = (
+                (B_pos / (-scale_tau / mu0)[None, :])[:, l] * np.exp(-scaled_tau / mu0)[None, :] 
+            )
+        elif there_is_beam_source and not is_antiderivative_wrt_tau:
+            direct_beam_contribution = B_pos[:, l] * np.exp(-scaled_tau / mu0)[None, :]
         else:
             direct_beam_contribution = 0
 
@@ -488,12 +492,15 @@ def _assemble_intensity_and_fluxes(
                 + direct_beam_contribution
                 + _mathscr_v_contribution
             )
-        
         flux = 2 * pi * (mu_arr_pos * W) @ u0_pos
-        return I0_orig * np.squeeze(flux)[()]
+        
+        if return_tau_arr:
+            return I0_orig * np.squeeze(flux)[()], tau_arr
+        else:
+            return I0_orig * np.squeeze(flux)[()]
 
 
-    def flux_down(tau, is_antiderivative_wrt_tau=False):
+    def flux_down(tau, is_antiderivative_wrt_tau=False, return_tau_arr=False):
         """
         (Energetic) Flux function with argument `tau` (type: array or float) for negative (downward) `mu` values.
         Returns a tuple of the diffuse and direct flux magnitudes respectively where each entry is of the
@@ -502,8 +509,8 @@ def _assemble_intensity_and_fluxes(
         to switch to an antiderivative of the function with respect to `tau`.
         """
         tau = np.atleast_1d(tau)
-        assert np.all(tau >= 0)
-        assert np.all(tau <= tau_arr[-1])
+        if np.all(tau < 0) or np.all(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)
@@ -528,20 +535,19 @@ def _assemble_intensity_and_fluxes(
                 G_inv_collect_0[l_uniq],
                 mu_arr,
                 is_antiderivative_wrt_tau,
-                _autograd_compatible
+                autograd_compatible
             )
         else:
             _mathscr_v_contribution = 0
         
-        if there_is_beam_source:
-            if is_antiderivative_wrt_tau:
-                direct_beam_contribution = (B_neg / (-scale_tau / mu0)[None, l])[:, l] * np.exp(-scaled_tau / mu0)[None, :]
-                direct_beam = I0 * mu0 * np.exp(-tau / mu0) * -mu0
-                direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0) / (-scale_tau / mu0)[None, l]
-            else:
-                direct_beam_contribution = B_neg[:, l] * np.exp(-scaled_tau / mu0)[None, :]
-                direct_beam = I0 * mu0 * np.exp(-tau / mu0)
-                direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0)
+        if there_is_beam_source and is_antiderivative_wrt_tau:
+            direct_beam_contribution = (B_neg / (-scale_tau / mu0)[None, l])[:, l] * np.exp(-scaled_tau / mu0)[None, :]
+            direct_beam = I0 * mu0 * np.exp(-tau / mu0) * -mu0
+            direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0) / (-scale_tau / mu0)[None, l]
+        elif there_is_beam_source and not is_antiderivative_wrt_tau:
+            direct_beam_contribution = B_neg[:, l] * np.exp(-scaled_tau / mu0)[None, :]
+            direct_beam = I0 * mu0 * np.exp(-tau / mu0)
+            direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0)
         else:
             direct_beam_contribution = 0
             direct_beam = 0
@@ -567,12 +573,19 @@ def _assemble_intensity_and_fluxes(
                 + direct_beam_contribution
                 + _mathscr_v_contribution
             )
-        
         diffuse_flux = 2 * pi * (mu_arr_pos * W) @ u0_neg + direct_beam_scaled - direct_beam
-        return (
-            I0_orig * np.squeeze(diffuse_flux)[()],
-            I0_orig * I0 * np.squeeze(direct_beam)[()],
-        )
+        
+        if return_tau_arr:
+            return (
+                I0_orig * np.squeeze(diffuse_flux)[()],
+                I0_orig * I0 * np.squeeze(direct_beam)[()],
+                tau_arr
+            )
+        else:
+            return (
+                I0_orig * np.squeeze(diffuse_flux)[()],
+                I0_orig * I0 * np.squeeze(direct_beam)[()],
+            )
         # --------------------------------------------------------------------------------------------------------------------------
 
     if only_flux:

PythonicDISORT/_solve_for_coeffs.py

@@ -109,18 +109,17 @@ def _solve_for_coeffs(
         # Assemble RHS
         # --------------------------------------------------------------------------------------------------------------------------
         # Ensure the BCs are of the correct shape
-        if b_pos_is_scalar:
-            if m_equals_0:
-                b_pos_m = np.full(N, b_pos)
-            else:
-                b_pos_m = np.zeros(N)
+        if b_pos_is_scalar and m_equals_0:
+            b_pos_m = np.full(N, b_pos)
+        elif b_pos_is_scalar and not m_equals_0:
+            b_pos_m = np.zeros(N)
         else:
             b_pos_m = b_pos[:, m]
-        if b_neg_is_scalar:
-            if m_equals_0:
-                b_neg_m = np.full(N, b_neg)
-            else:
-                b_neg_m = np.zeros(N)
+            
+        if b_neg_is_scalar and m_equals_0:
+            b_neg_m = np.full(N, b_neg)
+        elif b_neg_is_scalar and not m_equals_0:
+            b_neg_m = np.zeros(N)
         else:
             b_neg_m = b_neg[:, m]
         
@@ -205,7 +204,7 @@ def _solve_for_coeffs(
                 l_range = np.arange(1, NLayers)
                 RHS_middle = (
                     (B_collect_m[l_range, :] - B_collect_m[l_range - 1, :])
-                    * np.exp(-mu0 * scaled_tau_arr_with_0)[l_range, None]
+                    * np.exp(-scaled_tau_arr_with_0 / mu0)[l_range, None]
                 ).ravel()
                 
             if BDRF_bool:

PythonicDISORT/_solve_for_gen_and_part_sols.py

@@ -103,9 +103,7 @@ def _solve_for_gen_and_part_sols(
         for l in range(NLayers):
             weighted_asso_Leg_coeffs_l = weighted_scaled_Leg_coeffs[l, :][ells] * fac
             
-            if np.any(weighted_asso_Leg_coeffs_l > 0) and np.all( # Due to `fac` or the coeffs being zero
-                np.isfinite(asso_leg_term_pos) # We take precautions against overflow and underflow (this shouldn't happen though)
-            ):  
+            if np.any(weighted_asso_Leg_coeffs_l > 0): # There are shortcuts if the coefficients are all zero
                 scaled_omega_l = scaled_omega_arr[l]
                 
                 # Generate D
@@ -213,29 +211,27 @@ def _solve_for_gen_and_part_sols(
         # --------------------------------------------------------------------------------------------------------------------------
 
         
-    if there_is_beam_source:
-        if there_is_iso_source:
-            return (
-                G_collect.reshape((NFourier, NLayers, NQuad, NQuad)),
-                K_collect.reshape((NFourier, NLayers, NQuad)),
-                B_collect.reshape((NFourier, NLayers, NQuad)),
-                G_inv_collect_0,
-            )
-        else:
-            return (
-                G_collect.reshape((NFourier, NLayers, NQuad, NQuad)),
-                K_collect.reshape((NFourier, NLayers, NQuad)),
-                B_collect.reshape((NFourier, NLayers, NQuad)),
-            )
+    if there_is_beam_source and there_is_iso_source:
+        return (
+            G_collect.reshape((NFourier, NLayers, NQuad, NQuad)),
+            K_collect.reshape((NFourier, NLayers, NQuad)),
+            B_collect.reshape((NFourier, NLayers, NQuad)),
+            G_inv_collect_0,
+        )
+    elif there_is_beam_source and not there_is_iso_source:
+        return (
+            G_collect.reshape((NFourier, NLayers, NQuad, NQuad)),
+            K_collect.reshape((NFourier, NLayers, NQuad)),
+            B_collect.reshape((NFourier, NLayers, NQuad)),
+        )
+    elif not there_is_beam_source and there_is_iso_source:
+        return (
+            G_collect.reshape((NFourier, NLayers, NQuad, NQuad)),
+            K_collect.reshape((NFourier, NLayers, NQuad)),
+            G_inv_collect_0,
+        )
     else:
-        if there_is_iso_source:
-            return (
-                G_collect.reshape((NFourier, NLayers, NQuad, NQuad)),
-                K_collect.reshape((NFourier, NLayers, NQuad)),
-                G_inv_collect_0,
-            )
-        else:
-            return (
-                G_collect.reshape((NFourier, NLayers, NQuad, NQuad)), 
-                K_collect.reshape((NFourier, NLayers, NQuad)),
-            )
+        return (
+            G_collect.reshape((NFourier, NLayers, NQuad, NQuad)), 
+            K_collect.reshape((NFourier, NLayers, NQuad)),
+        )

PythonicDISORT/pydisort.py

@@ -1,5 +1,6 @@
 from PythonicDISORT import subroutines
 from PythonicDISORT._assemble_intensity_and_fluxes import _assemble_intensity_and_fluxes
+import warnings
 
 import numpy as np
 import scipy as sc
@@ -23,12 +24,12 @@ def pydisort(
     BDRF_Fourier_modes=[],
     s_poly_coeffs=np.array([[]]),
     use_banded_solver_NLayers=13,
-    _autograd_compatible = False,
+    autograd_compatible = False,
 ):
     """Solves the 1D RTE for the fluxes, and optionally intensity,
     of a multi-layer atmosphere with the specified optical properties, boundary conditions
     and sources. Optionally performs delta-M scaling and NT corrections. 
-    Refer to the `*_test.ipynb` Jupyter Notebooks in the `pydisotest` directory for examples of use.
+    Refer to the ``*_test.ipynb`` Jupyter Notebooks in the ``pydisotest`` directory for examples of use.
     
         See https://pythonic-disort.readthedocs.io/en/latest/Pythonic-DISORT.html#1.-USER-INPUT-REQUIRED:-Choose-parameters
         for a more detailed explanation of each parameter.
@@ -43,7 +44,7 @@ def pydisort(
     omega_arr : array or float
         Single-scattering albedo of each atmospheric layer.
     NQuad : int
-        Number of `mu` quadrature nodes.
+        Number of ``mu`` quadrature nodes.
     Leg_coeffs_all : 2darray
         All available unweighted phase function Legendre coefficients.
         Each row of coefficients pertain to an atmospheric layer.
@@ -70,47 +71,55 @@ def pydisort(
     NT_cor : optional, bool
         Perform Nakajima-Tanaka intensity corrections?
     BDRF_Fourier_modes : optional, list
-        BDRF Fourier modes, each a function with arguments `mu, -mu_p` of type array
+        BDRF Fourier modes, each a function with arguments ``mu, -mu_p`` of type array
         and which output has the same dimensions as the outer product of the two arrays.
     s_poly_coeffs : optional, array
         Polynomial coefficients of isotropic internal sources.
         Arrange coefficients from lowest order term to highest.
     use_banded_solver_NLayers : optional, int
-        At or above how many atmospheric layers should `scipy.linalg.solve_banded` be used?
+        At or above how many atmospheric layers should ``scipy.linalg.solve_banded`` be used?
+    autograd_compatible : optional, bool
+        If ``True``, the autograd package: https://github.com/HIPS/autograd can be used to compute
+        the tau-derivatives of the function outputs but ``pydisort`` will be less efficient.  
 
     Returns
     -------
     mu_arr : array or float
-        All `mu` (cosine of polar angle) quadrature nodes.
+        All ``mu`` (cosine of polar angle) quadrature nodes.
     Fp(tau) : function
-        (Energetic) Flux function with argument `tau` (type: array or float) for positive (upward) `mu` values.
-        Returns the diffuse flux magnitudes (same type and size as `tau`).
-        Pass `is_antiderivative_wrt_tau = True` (defaults to `False`)
-        to switch to an antiderivative of the function with respect to `tau`.
-        For example, `Fp(tau0, is_antiderivative_wrt_tau=True) - Fp(0, is_antiderivative_wrt_tau=True)` 
-        will give the tau-integral of `Fp` from `tau0` to `0`.
+        (Energetic) Flux function with argument ``tau`` (type: array or float) for positive (upward) ``mu`` values.
+        Returns the diffuse flux magnitudes (same type and size as ``tau``).
+        Pass ``is_antiderivative_wrt_tau = True`` (defaults to ``False``)
+        to switch to an antiderivative of the function with respect to ``tau``.
+        For example, ``Fp(1, is_antiderivative_wrt_tau=True) - Fp(0, is_antiderivative_wrt_tau=True)`` 
+        will produce the tau-integral over ``[1, 0]`` assuming both values belong to the same layer, and
+        ``Fp(tau_arr, is_antiderivative_wrt_tau = True) - Fp(np.insert(tau_arr[:-1] + 1e-15, 0, 0), is_antiderivative_wrt_tau = True)``
+        will produce an array of the tau-integral over each layer.
+        Pass ``return_tau_arr`` to return ``tau_arr`` (defaults to ``False``).
     Fm(tau) : function
-        (Energetic) Flux function with argument `tau` (type: array or float) for negative (downward) `mu` values.
+        (Energetic) Flux function with argument ``tau`` (type: array or float) for negative (downward) ``mu`` values.
         Returns a tuple of the diffuse and direct flux magnitudes respectively where each entry is of the
-        same type and size as `tau`.
-        Pass `is_antiderivative_wrt_tau = True` (defaults to `False`)
-        to switch to an antiderivative of the function with respect to `tau`.
+        same type and size as ``tau``.
+        Pass ``is_antiderivative_wrt_tau = True`` (defaults to ``False``)
+        to switch to an antiderivative of the function with respect to ``tau``.
+        Pass ``return_tau_arr`` to return ``tau_arr`` (defaults to ``False``).
     u0(tau) : function
-        Zeroth Fourier mode of the intensity with argument `tau` (type: array or float).
-        Returns an ndarray with axes corresponding to variation with `mu` and `tau` respectively.
+        Zeroth Fourier mode of the intensity with argument ``tau`` (type: array or float).
+        Returns an ndarray with axes corresponding to variation with ``mu`` and ``tau`` respectively.
         This function is useful for calculating actinic fluxes and other quantities of interest,
         but reclassification of delta-scaled flux and other corrections must be done manually
-        (for actinic flux `generate_diff_act_flux_funcs` will automatically perform the reclassification).
-        Pass `is_antiderivative_wrt_tau = True` (defaults to `False`)
-        to switch to an antiderivative of the function with respect to `tau`.
+        (for actinic flux ``subroutines.generate_diff_act_flux_funcs`` will automatically perform the reclassification).
+        Pass ``is_antiderivative_wrt_tau = True`` (defaults to ``False``)
+        to switch to an antiderivative of the function with respect to ``tau``.
+        Pass ``return_tau_arr`` to return ``tau_arr`` (defaults to ``False``).
     u(tau, phi) : function, optional
-        Intensity function with arguments `(tau, phi)` each of type array or float.
-        Returns an ndarray with axes corresponding to variation with `mu, tau, phi` respectively.
-        Pass `return_Fourier_error = True` (defaults to `False`) to return the 
+        Intensity function with arguments ``(tau, phi)`` each of type array or float.
+        Returns an ndarray with axes corresponding to variation with ``mu, tau, phi`` respectively.
+        Pass ``is_antiderivative_wrt_tau = True`` (defaults to ``False``)
+        to switch to an antiderivative of the function with respect to ``tau``.
+        Pass ``return_Fourier_error = True`` (defaults to ``False``) to return the 
         Cauchy / Fourier convergence evaluation (type: float) for the last Fourier term.
-        Pass `is_antiderivative_wrt_tau = True` (defaults to `False`)
-        to switch to an antiderivative of the function with respect to `tau`.
-
+        Pass ``return_tau_arr`` to return ``tau_arr`` (defaults to ``False``).
     """
     
     """
@@ -159,7 +168,7 @@ def pydisort(
     | `scaled_mu0`                 | scalar               |
     """
 
-    if _autograd_compatible:
+    if autograd_compatible:
         import autograd.numpy as np
     else:
         import numpy as np
@@ -201,50 +210,71 @@ def pydisort(
     # Input checks (refer to Section 1 of the Comprehensive Documentation)
     # --------------------------------------------------------------------------------------------------------------------------
     # Optical depths and thickness must be positive
-    assert np.all(tau_arr > 0)
-    assert np.all(thickness_arr > 0)
+    if not np.all(tau_arr > 0):
+        raise ValueError("tau values cannot be negative.")
+    if not np.all(thickness_arr > 0):
+        raise ValueError("Layer thickness cannot be negative.")
     # Single-scattering albedo must be between 0 and 1, excluding 1
-    assert np.all(omega_arr >= 0)
-    assert np.all(omega_arr < 1)
+    if not (np.all(omega_arr >= 0) and np.all(omega_arr < 1)):
+        raise ValueError("Single-scattering albedo must be between 0 and 1, excluding 1.") 
     # There must be a positive number of Legendre coefficients each with magnitude <= 1
     # The user must supply at least as many phase function Legendre coefficients as intended for use
-    assert NLeg > 0
-    assert np.all(Leg_coeffs_all[:, 0] == 1)
-    assert np.all(0 <= Leg_coeffs_all) and np.all(Leg_coeffs_all <= 1)
-    assert np.all(np.array([f.__code__.co_argcount for f in BDRF_Fourier_modes]) == 2)
-    assert NLeg <= NLeg_all
+    if not NLeg > 0:
+        raise ValueError("The number of phase function Legendre coefficients must be positive.") 
+    if not np.all(Leg_coeffs_all[:, 0] == 1):
+        raise ValueError("The first phase function Legendre coefficient must equal 1.") 
+    if not (np.all(0 <= Leg_coeffs_all) and np.all(Leg_coeffs_all <= 1)):
+        raise ValueError("The phase function Legendre coefficients must all be between 0 and 1.") 
+    if not np.all(np.array([f.__code__.co_argcount for f in BDRF_Fourier_modes]) == 2):
+        raise ValueError("Some functions in the list `BDRF_Fourier_modes` are of the wrong form.")
+    if not NLeg <= NLeg_all:
+        raise ValueError("`NLeg` cannot be larger than the number of phase function Legendre coefficients provided.")
     # Ensure that the first dimension of the following inputs corresponds to the number of layers
-    assert np.shape(Leg_coeffs_all)[0] == NLayers
-    assert len(omega_arr) == NLayers
-    if len(f_arr) != 1 or f_arr[0] != 0:
-        assert len(f_arr) == NLayers
-    if there_is_iso_source:
-        assert np.shape(s_poly_coeffs)[0] == NLayers
+    if not np.shape(Leg_coeffs_all)[0] == NLayers:
+        raise ValueError("The zeroth dimension of the shape of `Leg_coeffs_all` does not match the number of layers as deduced from the length of `tau_arr`.")
+    if not len(omega_arr) == NLayers:
+        raise ValueError("The zeroth dimension of the shape of `omega_arr` does not match the number of layers as deduced from the length of `tau_arr`.")
+    if np.any(f_arr != 0) and not len(f_arr) == NLayers:
+        raise ValueError("The length of `f_arr` does not match the number of layers as deduced from the length of `tau_arr`.")
+    if there_is_iso_source and not np.shape(s_poly_coeffs)[0] == NLayers:
+        raise ValueError("The zeroth dimension of the shape of `s_poly_coeffs` does not match the number of layers as deduced from the length of `tau_arr`.")
     # Conditions on the number of quadrature angles (NQuad), Legendre coefficients (NLeg) and loops (NFourier)
-    assert NQuad >= 2
-    assert NQuad % 2 == 0
-    assert NFourier > 0
-    assert NFourier <= NLeg
+    if not NQuad >= 2:
+        raise ValueError("There must be at least two streams.")
+    if not NQuad % 2 == 0:
+        raise ValueError("The number of streams must be even.")
+    if not NFourier > 0:
+        raise ValueError("The number of Fourier modes to use in the solution must be positive.")
+    if not NFourier <= NLeg:
+        raise ValueError("The number of Fourier modes to use in the solution must be less than or equal to the number of phase function Legendre coefficients used.")
+    if NFourier > 64 and not only_flux:
+        warnings.warn("`NFourier` is large and may cause errors, consider decreasing `NFourier` to 64 and it probably should be even less.")
     # Not strictly necessary but there will be tremendous inaccuracies if this is violated
-    assert NLeg <= NQuad
+    if not NLeg <= NQuad:
+        raise ValueError("There should be more streams than the number of phase function Legendre coefficients used.")
     # We require principal angles and a downward incident beam
-    assert I0 >= 0
+    if not I0 >= 0:
+        raise ValueError("The incident beam must have positive intensity.")
     if there_is_beam_source:
-        assert 0 < mu0 and mu0 <= 1
-        assert 0 <= phi0 and phi0 < 2 * pi
+        if not 0 < mu0 and mu0 <= 1:
+            raise ValueError("The cosine of the polar angle of the incident beam must be between 0 and 1, excluding 0.")
+        if not 0 <= phi0 and phi0 < 2 * pi:
+            raise ValueError("Provide the principal azimuthal angle for the incident beam (must be between 0 and 2pi, excluding 2pi).")
     # Ensure that the BC inputs are of the correct shape
     if len(np.atleast_1d(b_pos)) == 1:
         b_pos_is_scalar = True
-    else:
-        assert np.shape(b_pos) == (N, NFourier)
+    elif not np.shape(b_pos) == (N, NFourier):
+        raise ValueError("The shape of the bottom of atmosphere boundary conditions must equal (NQuad // 2, NFourier).")
     if len(np.atleast_1d(b_neg)) == 1:
         b_neg_is_scalar = True
-    else:
-        assert np.shape(b_neg) == (N, NFourier)
+    elif not np.shape(b_neg) == (N, NFourier):
+        raise ValueError("The shape of the top of atmosphere boundary conditions must equal (NQuad // 2, NFourier).")
     # The fractional scattering must be between 0 and 1
-    assert np.all(0 <= f_arr) and np.all(f_arr <= 1)
+    if not (np.all(0 <= f_arr) and np.all(f_arr <= 1)):
+        raise ValueError("The fractional scattering must be between 0 and 1.")
     # The minimum threshold is 3 else the matrix will not be banded
-    assert(use_banded_solver_NLayers >= 3)
+    if not use_banded_solver_NLayers >= 3:
+        raise ValueError("The minimum threshold `use_banded_solver_NLayers` is 3, else the matrix will not be banded.")
     # --------------------------------------------------------------------------------------------------------------------------
     
     # Some more setup
@@ -268,7 +298,8 @@ def pydisort(
     M_inv = 1 / mu_arr_pos
     
     # We do not allow mu0 to equal a quadrature / computational angle
-    assert not np.any(np.abs(mu_arr_pos - mu0) < 1e-8)
+    if np.any(np.abs(mu_arr_pos - mu0) < 1e-8):
+        raise ValueError("Some quadrature angles come too close to `mu0`. Perturb `NQuad` or `mu0` to rectify this error.")
     # --------------------------------------------------------------------------------------------------------------------------
 
     # Delta-M scaling; there is no scaling if f = 0
@@ -317,7 +348,7 @@ def pydisort(
             scale_tau,
             only_flux,
             use_banded_solver_NLayers,
-            _autograd_compatible,
+            autograd_compatible,
         )
         
         # TMS correction for the intensity (see Section 3.7.2)
@@ -527,7 +558,7 @@ def pydisort(
                             axis=1,
                         )
 
-                if _autograd_compatible:
+                if autograd_compatible:
                     if not any_pos_contribution:
                         TMS_correction_pos = np.zeros((N, Ntau, Nphi))
                     if not any_neg_contribution:
@@ -579,6 +610,7 @@ def pydisort(
                     (tau[None, :] - 1 / x[:, None]) * np.exp(-tau / scaled_mu0)[None, :]
                     + np.exp(-tau[None, :] / mu_arr_pos[:, None]) / x[:, None]
                 ) / (mu_arr_pos * scaled_mu0 * x)[:, None]
+                
             return (
                 I0
                 / (4 * pi)
@@ -593,26 +625,25 @@ def pydisort(
 
         # The corrected intensity
         # --------------------------------------------------------------------------------------------------------------------------
-        def u_corrected(tau, phi, return_Fourier_error=False, is_antiderivative_wrt_tau=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)
             NT_corrections = TMS_correction(tau, phi, is_antiderivative_wrt_tau)
 
-            if _autograd_compatible:
+            if autograd_compatible:
                 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:
-                u_star_outputs = u_star(tau, phi, True, 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] + I0_orig * np.squeeze(NT_corrections),
-                    u_star_outputs[1],
-                )
+                ) + u_star_outputs[1:]
             else:
-                return u_star(tau, phi, False, is_antiderivative_wrt_tau) + I0_orig * np.squeeze(NT_corrections)
+                return u_star(tau, phi, is_antiderivative_wrt_tau) + I0_orig * np.squeeze(NT_corrections)
         # --------------------------------------------------------------------------------------------------------------------------
 
         return mu_arr, flux_up, flux_down, u0, u_corrected
@@ -641,5 +672,5 @@ def pydisort(
             scale_tau,
             only_flux,
             use_banded_solver_NLayers,
-            _autograd_compatible,
+            autograd_compatible,
         )

PythonicDISORT/subroutines.py

@@ -28,6 +28,7 @@ def transform_interval(arr, c, d, a, b):
     return (((arr - a) * (d - c)) / (b - a)) + c
 
 
+
 def transform_weights(weights, c, d, a, b):
     """Transforms an array of quadrature weights from interval [a, b] to [c, d].
 
@@ -53,9 +54,10 @@ def transform_weights(weights, c, d, a, b):
     return weights * (d - c) / (b - a)
 
 
+
 def calculate_nu(mu, phi, mu_p, phi_p):
-    """Calculates the cosine of the scattering angle `nu` 
-    between incident angle `(mu_p, phi_p)` and scattering angle `(mu, phi)`.
+    """Calculates the cosine of the scattering angle ``nu`` 
+    between incident angle ``(mu_p, phi_p)`` and scattering angle ``(mu, phi)``.
 
     Parameters
     ----------
@@ -71,7 +73,7 @@ def calculate_nu(mu, phi, mu_p, phi_p):
     Returns
     -------
     nu : ndarray
-        Cosine of scattering angles which axes capture variation with `mu, phi, mu_p, phi_p` respectively.
+        Cosine of scattering angles which axes capture variation with ``mu, phi, mu_p, phi_p`` respectively.
 
     """
     mu, phi, mu_p, phi_p = np.atleast_1d(mu, phi, mu_p, phi_p)
@@ -83,12 +85,13 @@ def calculate_nu(mu, phi, mu_p, phi_p):
     return np.squeeze(nu)
 
 
+
 def Gauss_Legendre_quad(N, c=0, d=1):
     """Generates Gauss-Legendre quadrature zero points and weights for integration from c to d.
 
     Parameters
     ----------
-    N : int
+    N : int (will be converted to int)
         Number of quadrature nodes.
     c : float, optional
         Start of integration interval.
@@ -108,6 +111,7 @@ def Gauss_Legendre_quad(N, c=0, d=1):
     return transform_interval(mu_arr_pos, c, d, -1, 1), transform_weights(W, c, d, -1, 1)
 
 
+
 def Clenshaw_Curtis_quad(Nphi, c=0, d=(2 * pi)):
     """Generates Clenshaw-Curtis quadrature zero points and weights for integration from c to d.
 
@@ -129,8 +133,8 @@ def Clenshaw_Curtis_quad(Nphi, c=0, d=(2 * pi)):
 
     """
     # Ensure that the number of nodes is odd and greater than 2
-    assert Nphi > 2
-    assert Nphi % 2 == 1
+    if not (Nphi > 2 and Nphi % 2 == 1):
+        raise ValueError("The number of quadrature nodes must be odd and greater than 2.")
 
     Nphi -= 1  # The extra index corresponds to the point 0 which we will add later
     Nphi_pos = Nphi // 2
@@ -144,9 +148,10 @@ def Clenshaw_Curtis_quad(Nphi, c=0, d=(2 * pi)):
     return transform_interval(phi_arr, c, d, -1, 1), transform_weights(full_weights_phi, c, d, -1, 1)
 
 
+
 def generate_FD_mat(Ntau, a, b):
     """Generates a sparse first derivative central difference (second-order accuracy) 
-    matrix on [a,b] in `csr` format with `Ntau` grid points.
+    matrix on [a,b] in ``csr`` format with ``Ntau`` grid points.
     Second order forward or backward differences are used at the boundaries.
 
     Parameters
@@ -183,6 +188,7 @@ def generate_FD_mat(Ntau, a, b):
     return tau_arr, first_deriv.tocsr()
   
 
+
 def atleast_2d_append(*arys):
     """View inputs as arrays with at least two dimensions. 
     Dimensions are added as necessary to the back of the shape tuple rather than to the front.
@@ -219,7 +225,8 @@ def atleast_2d_append(*arys):
         return res[0]
     else:
         return res
-        
+       
+       
         
 def generate_diff_act_flux_funcs(u0):
     """Generates the up and down diffuse actinic flux functions respectively.
@@ -237,73 +244,95 @@ def generate_diff_act_flux_funcs(u0):
     Returns
     -------
     Fp_act(tau) : function
-        Actinic flux function with argument `tau` (type: array) for positive (upward) `mu` values.
+        Actinic flux function with argument ``tau`` (type: array) for positive (upward) ``mu`` values.
         Returns the diffuse flux magnitudes (type: array).
-        Pass `is_antiderivative_wrt_tau = True` (defaults to `False`)
-        to switch to an antiderivative of the function with respect to `tau`.
+        Pass ``is_antiderivative_wrt_tau = True`` (defaults to ``False``)
+        to switch to an antiderivative of the function with respect to ``tau``.
+        Pass ``return_tau_arr`` to return ``tau_arr`` (defaults to ``False``).
     Fm_act(tau) : function
-        Actinic flux function with argument `tau` (type: array) for negative (downward) `mu` values.
+        Actinic flux function with argument ``tau`` (type: array) for negative (downward) ``mu`` values.
         Returns the diffuse flux magnitudes (type: array).
-        Pass `is_antiderivative_wrt_tau = True` (defaults to `False`)
-        to switch to an antiderivative of the function with respect to `tau`.
+        Pass ``is_antiderivative_wrt_tau = True`` (defaults to ``False``)
+        to switch to an antiderivative of the function with respect to ``tau``.
+        Pass ``return_tau_arr`` to return ``tau_arr`` (defaults to ``False``).
 
     """
-    N = np.shape(u0(0))[0] // 2
+    N = len(u0(0)) // 2
     GL_weights = Gauss_Legendre_quad(N, 0, 1)[1]
 
     # Note that the zeroth axis of the array u0(tau) captures variation with mu
-    def flux_act_up(tau, is_antiderivative_wrt_tau=False):
-        return np.squeeze(2 * pi * GL_weights @ u0(tau, is_antiderivative_wrt_tau)[:N])[()]
+    def flux_act_up(tau, is_antiderivative_wrt_tau=False, return_tau_arr=False):
+        if return_tau_arr:
+            (
+                u0_cache, 
+                tau_arr,
+            ) = u0(tau, is_antiderivative_wrt_tau, True)
+            return np.squeeze(2 * pi * GL_weights @ u0_cache[:N])[()], tau_arr
+        else:
+            return np.squeeze(2 * pi * GL_weights @ u0(tau, is_antiderivative_wrt_tau)[:N])[()]
         
-    def flux_act_down_diffuse(tau, is_antiderivative_wrt_tau=False):
-        u0_cache, act_dscale_reclassification = u0(tau, is_antiderivative_wrt_tau, _return_act_dscale_for_reclass=True)
-        result_without_reclassification = 2 * pi * GL_weights @ u0_cache[N:]
-        return np.squeeze(result_without_reclassification + act_dscale_reclassification)[()]
+    def flux_act_down_diffuse(tau, is_antiderivative_wrt_tau=False, return_tau_arr=False):
+        if return_tau_arr:
+            (
+                u0_cache, 
+                tau_arr,
+                act_dscale_reclassification,
+            ) = u0(tau, is_antiderivative_wrt_tau, True, _return_act_dscale_for_reclass=True)
+            result_without_reclassification = 2 * pi * GL_weights @ u0_cache[N:]
+            return np.squeeze(result_without_reclassification + act_dscale_reclassification)[()], tau_arr
+        else:
+            (
+                u0_cache, 
+                act_dscale_reclassification,
+            ) = u0(tau, is_antiderivative_wrt_tau, False, _return_act_dscale_for_reclass=True)
+            result_without_reclassification = 2 * pi * GL_weights @ u0_cache[N:]
+            return np.squeeze(result_without_reclassification + act_dscale_reclassification)[()]
     
     return flux_act_up, flux_act_down_diffuse
 
 
-def interpolate(mu_arr_pos, u):
+
+def interpolate(u):
     """Polynomial (Barycentric) interpolation with respect to mu. The output 
     is a function that is continuous and variable in all three arguments: mu, tau and phi.
     Discussed in sections 3.7 and 6.3 in the Comprehensive Documentation.
 
     Parameters
     ----------
-    mu_arr_pos : array
-        Positive `mu` (cosine of polar angle) quadrature nodes.
-        Equivalent to `mu_arr[:N]`.
     u : function
-        Non-interpolated intensity function as outputted by `pydisort`.
+        Non-interpolated intensity function as given by ``pydisort``.
 
     Returns
     -------
     u_interpol : function
-        Intensity function with arguments `(mu, tau, phi)` each of type array or float.
-        Returns an ndarray with axes corresponding to variation with `mu, tau, phi` respectively.
-        Pass `return_Fourier_error = True` (defaults to `False`) to return the 
+        Intensity function with arguments ``(mu, tau, phi)`` each of type array or float.
+        Returns an ndarray with axes corresponding to variation with ``mu, tau, phi`` respectively.
+        Pass ``is_antiderivative_wrt_tau = True`` (defaults to ``False``)
+        to switch to an antiderivative of the function with respect to ``tau``.
+        Pass ``return_Fourier_error = True`` (defaults to ``False``) to return the 
         Cauchy / Fourier convergence evaluation (type: float) for the last Fourier term.
-        Pass `is_antiderivative_wrt_tau = True` (defaults to `False`)
-        to switch to an antiderivative of the function with respect to `tau`.
+        Pass ``return_tau_arr`` to return ``tau_arr`` (defaults to ``False``).
 
     """
-    N = len(mu_arr_pos)
+    N = len(u(0, 0)) // 2
+    mu_arr_pos = Gauss_Legendre_quad(N)[0]
+    
     u_pos_interpol = sc.interpolate.BarycentricInterpolator(mu_arr_pos)
     u_neg_interpol = sc.interpolate.BarycentricInterpolator(-mu_arr_pos)
 
-    def u_interpol(mu, tau, phi, return_Fourier_error=False, is_antiderivative_wrt_tau=False):
-        assert np.all(np.abs(mu) <= 1)
+    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:
-            u_cache, Fourier_error = u(tau, phi, True, is_antiderivative_wrt_tau)
+        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]
         else:
-            u_cache = u(tau, phi, False, is_antiderivative_wrt_tau)
+            u_cache = u(tau, phi, is_antiderivative_wrt_tau)
         
-        u_shape =  np.shape(u_cache)
-        assert N == u_shape[0] // 2 # Check that mu_arr_pos (of length N) matches the function u
-        
-        results = np.empty((len(mu),) + u_shape[1:])
+        results = np.empty((len(mu),) + np.shape(u_cache)[1:])
         mask_pos = mu > 0
         mask_else = ~mask_pos
         
@@ -314,17 +343,18 @@ def interpolate(mu_arr_pos, u):
             u_neg_interpol.set_yi(u_cache[N:])
             results[mask_else] = u_neg_interpol(mu[mask_else])
         
-        if return_Fourier_error:
-            return np.squeeze(results)[()], Fourier_error
+        if return_Fourier_error or return_tau_arr:
+            return (np.squeeze(results)[()],) + u_outputs[1:]
         else:
             return np.squeeze(results)[()]
 
     return u_interpol
-          
+         
+         
 
 def to_diag_ordered_form(A, sym_offset):
     """
-    Convert a matrix A to the diagonal ordered form required by `scipy.linalg.solve_banded`.
+    Convert a matrix A to the diagonal ordered form required by ``scipy.linalg.solve_banded``.
     We assume that the matrix has the same number of super- and sub-diagonals.
 
     Parameters
@@ -356,6 +386,7 @@ def to_diag_ordered_form(A, sym_offset):
         axis=0,
     )
 
+
     
 def _mathscr_v(tau,                             # Input optical depths
                 l,                              # Layer index of each input optical depth
@@ -366,40 +397,40 @@ def _mathscr_v(tau,                             # Input optical depths
                 G_inv,                          # Inverse of eigenvector matrix
                 mu_arr,                         # Quadrature nodes for both hemispheres
                 is_antiderivative_wrt_tau=False, # Switch to an antiderivative of the function?
-                _autograd_compatible=False,      # Should the output functions be compatible with autograd?
+                autograd_compatible=False,      # Should the output functions be compatible with autograd?
                 ):
     """Particular solution for isotropic internal sources.
     Refer to Section 3.6.1 of the Comprehensive Documentation.
     It has many seemingly redundant arguments to maximize 
-    precomputation in the `_assemble_intensity_and_fluxes` 
-    and `_solve_for_coeffs` functions which call it.
+    precomputation in the ``_assemble_intensity_and_fluxes`` 
+    and ``_solve_for_coeffs`` functions which call it.
     
     Arguments of _mathscr_v
     |          Variable           |                 Shape                 |
     | --------------------------- | ------------------------------------- |
-    | `tau`                       | `Ntau`                                |
-    | `l`                         | `Ntau`                                |
-    | `Nscoeffs`                  | scalar                                |
-    | `s_poly_coeffs`             | `NLayers x Nscoeffs` or `Nscoeffs`    |
-    | `G`                         | `NLayers<= x NQuad x NQuad`           |
-    | `K`                         | `NLayers<= x NQuad`                   |
-    | `G_inv`                     | `NLayers<= x NQuad x NQuad` or `None` |
-    | `mu_arr`                    | `NQuad`                               |
-    | `is_antiderivative_wrt_tau` | boolean                               |
-    | `_autograd_compatible`      | boolean                               |
+    | ``tau``                       | ``Ntau``                                |
+    | ``l``                         | ``Ntau``                                |
+    | ``Nscoeffs``                  | scalar                                |
+    | ``s_poly_coeffs``             | ``NLayers x Nscoeffs`` or ``Nscoeffs``    |
+    | ``G``                         | ``NLayers<= x NQuad x NQuad``           |
+    | ``K``                         | ``NLayers<= x NQuad``                   |
+    | ``G_inv``                     | ``NLayers<= x NQuad x NQuad`` or ``None`` |
+    | ``mu_arr``                    | ``NQuad``                               |
+    | ``is_antiderivative_wrt_tau`` | boolean                               |
+    | ``autograd_compatible``      | boolean                               |
     
     Notable internal variables of _mathscr_v
     |     Variable     |                Shape                | 
     | ---------------- | ----------------------------------- |
-    | i_arr            | `Nscoeffs`                          |
-    | i_arr_repeat     | `Nscoeffs*(Nscoeffs+1)/2`           | Using triangular number formula
-    | j_arr            | `Nscoeffs*(Nscoeffs+1)/2`           | Using triangular number formula
-    | s_poly_coeffs_nj | `NLayers x Nscoeffs*(Nscoeffs+1)/2` |
-    | OUTPUT           | `NQuad x Ntau`                      |
+    | i_arr            | ``Nscoeffs``                          |
+    | i_arr_repeat     | ``Nscoeffs*(Nscoeffs+1)/2``           | Using triangular number formula
+    | j_arr            | ``Nscoeffs*(Nscoeffs+1)/2``           | Using triangular number formula
+    | s_poly_coeffs_nj | ``NLayers x Nscoeffs*(Nscoeffs+1)/2`` |
+    | OUTPUT           | ``NQuad x Ntau``                      |
     """
     n = Nscoeffs - 1
     
-    if _autograd_compatible:
+    if autograd_compatible:
         import autograd.numpy as np
     
         def mathscr_b(i):
@@ -407,9 +438,9 @@ def _mathscr_v(tau,                             # Input optical depths
             Notable internal variables of mathscr_b
             |     Variable     |                 Shape                 |
             | ---------------- | ------------------------------------- |
-            | j_arr            | `i + 1`                               |
-            | s_poly_coeffs_nj | `i + 1`                               |
-            | OUTPUT           | `Nscoeffs x (i + 1) x NQuad`          |
+            | j_arr            | ``i + 1``                               |
+            | s_poly_coeffs_nj | ``i + 1``                               |
+            | OUTPUT           | ``Nscoeffs x (i + 1) x NQuad``          |
             """
         
             j_arr = np.arange(i + 1)
@@ -460,6 +491,7 @@ def _mathscr_v(tau,                             # Input optical depths
     )
     
 
+
 def _compare(results, mu_to_compare, reorder_mu, flux_up, flux_down, u):
     """Performs pointwise comparisons between results from Stamnes' DISORT,
     which are stored in .npz files, against results from PythonicDISORT. Used in our PyTests.

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

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

PythonicDISORT-0.9.3.dist-info/RECORD

@@ -0,0 +1,11 @@
+PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
+PythonicDISORT/_assemble_intensity_and_fluxes.py,sha256=PsZJN-eHKKMP2VHqk8zFVw1CZ5Em2tc0aJOej3qdg2k,29134
+PythonicDISORT/_solve_for_coeffs.py,sha256=8xnjR6fenaTiPXXzxvtXmETFZRnwOnGkHIuuOFpz7Lc,15402
+PythonicDISORT/_solve_for_gen_and_part_sols.py,sha256=1IvyEZrcqLcafAQgAFpXtJ2H6hCvIM11OgBHF3KqmGc,12960
+PythonicDISORT/pydisort.py,sha256=AZmhV2niPW_XK5cGi_tMSgQQemBtApQB1sKLbEQXqLY,37039
+PythonicDISORT/subroutines.py,sha256=tq9_m4iXE8zLE8hh69aWh4wNfRHyCkgZFxZI48DRNE0,21312
+PythonicDISORT-0.9.3.dist-info/LICENSE.md,sha256=q1i1QVh0OLzK0fj6kaNRNWsM_KZca66JwWgWwbuCPlA,1092
+PythonicDISORT-0.9.3.dist-info/METADATA,sha256=elndi6Hlfg3SmgDWdV-Esuu3JML574zlvDvpwrI4I8k,7102
+PythonicDISORT-0.9.3.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+PythonicDISORT-0.9.3.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
+PythonicDISORT-0.9.3.dist-info/RECORD,,

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

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

PythonicDISORT-0.9.1.dist-info/RECORD

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