PythonicDISORT 0.7.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

PythonicDISORT/_assemble_intensity_and_fluxes.py

@@ -28,9 +28,9 @@ def _assemble_intensity_and_fluxes(
     NLayers_to_use_sparse,              # Number of layers above or equal which to use sparse matrices
     _autograd_compatible,               # Should the output functions be compatible with autograd?
 ):  
-    """In this function the solution functions: intensity `u`, upward flux `flux_up`, downward flux `flux_down` 
-    are assembled from the eigenpairs and the coefficients which have been solved from the boundary conditions. 
-    These solution functions are then returned.
+    """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.
     This function is wrapped by the `pydisort` function.
     It should be called through `pydisort` and never directly.
     It has many seemingly redundant arguments, which are described below,
@@ -161,7 +161,15 @@ 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):
+        def u(tau, phi, return_Fourier_error=False, is_antiderivative_wrt_tau=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.
+            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`.
+            """
             tau = np.atleast_1d(tau)
             phi = np.atleast_1d(phi)
             assert np.all(tau >= 0)
@@ -179,7 +187,7 @@ def _assemble_intensity_and_fluxes(
                 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],
@@ -187,14 +195,33 @@ def _assemble_intensity_and_fluxes(
                 ],
                 axis=-1,
             )
-            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[None, :, None] / mu0)
+            
+            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, l, None] * K_collect[:, l, :]),
+                        optimize=True,
+                    ) + B_collect[:, l, :] * np.exp(-scaled_tau[None, :, None] / mu0) / (
+                        -scale_tau[None, l, None] / mu0
+                    )
+                else:
+                    um = np.einsum(
+                        "mtij, mtj -> mti",
+                        GC_collect[:, l, :, :],
+                        np.exp(exponent) / (scale_tau[None, l, None] * K_collect[:, l, :]),
+                        optimize=True,
+                    )
             else:
-                um = np.einsum(
-                    "mtij, mtj -> mti", GC_collect[:, l, :, :], np.exp(exponent), optimize=True
-                )
+                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[None, :, None] / mu0)
+                else:
+                    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:
@@ -206,11 +233,15 @@ def _assemble_intensity_and_fluxes(
                     K_collect_0,
                     G_inv_collect_0,
                     mu_arr,
-                    _autograd_compatible,
+                    is_antiderivative_wrt_tau,
+                    _autograd_compatible
                 )
                 
                 if _autograd_compatible:
-                    um = um + np.concatenate((_mathscr_v_contribution.T, np.zeros((NFourier - 1, len(tau), NQuad))))
+                    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))))
                 else:
                     um[0, :, :] += _mathscr_v_contribution.T
                 
@@ -229,20 +260,43 @@ def _assemble_intensity_and_fluxes(
                     ],
                     axis=-1,
                 )
-                if there_is_beam_source:
-                    ulast = np.einsum(
-                        "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
-                    ) + B_collect[-1, l, :].T * np.exp(-scaled_tau[None, :] / mu0)
+                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[l, None] * K_collect[-1, l, :]),
+                            optimize=True,
+                        ) + B_collect[-1, l, :].T * np.exp(-scaled_tau[None, :] / mu0) / (
+                            -scale_tau[None, l] / mu0
+                        )
+                    else:
+                        ulast = np.einsum(
+                            "tij, tj -> it",
+                            GC_collect[-1, l, :, :],
+                            np.exp(exponent) / (scale_tau[l, None] * K_collect[-1, l, :]),
+                            optimize=True,
+                        )
                 else:
-                    ulast = np.einsum(
-                        "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
-                    )
+                    if there_is_beam_source:
+                        ulast = np.einsum(
+                            "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
+                        ) + B_collect[-1, l, :].T * np.exp(-scaled_tau[None, :] / mu0)
+                    else:
+                        ulast = np.einsum(
+                            "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
+                        )
+                        
                 intensities_abs = np.abs(intensities)
                 Fourier_error = np.max(
-                    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,
+                    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,
+                    )
                 )
                 return I0_orig * np.squeeze(intensities), Fourier_error
             else:
@@ -251,7 +305,16 @@ def _assemble_intensity_and_fluxes(
     
     # Construct u0
     # --------------------------------------------------------------------------------------------------------------------------
-    def u0(tau, return_act_dscale_reclassification=False):
+    def u0(tau, is_antiderivative_wrt_tau=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.
+        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`.
+        """
         tau = np.atleast_1d(tau)
         assert np.all(tau >= 0)
         assert np.all(tau <= tau_arr[-1])
@@ -268,11 +331,14 @@ def _assemble_intensity_and_fluxes(
             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_reclassification: 
-                act_dscale_reclassification = I0 * np.exp(-scaled_tau / mu0) - I0 * np.exp(-tau / mu0)
+            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_reclassification:
+            if _return_act_dscale_for_reclass:
                 act_dscale_reclassification = 0
 
         exponent = np.concatenate(
@@ -282,14 +348,33 @@ def _assemble_intensity_and_fluxes(
             ],
             axis=-1,
         )
-        if there_is_beam_source:
-            u0 = np.einsum(
-                "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
-            ) + B_collect_0[l, :].T * np.exp(-scaled_tau[None, :] / mu0)
+        
+        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[l, None] * K_collect_0[l, :]),
+                    optimize=True,
+                ) + B_collect_0[l, :].T * np.exp(-scaled_tau[None, :] / mu0) / (
+                    -scale_tau[None, l] / mu0
+                )
+            else:
+                u0 = np.einsum(
+                    "tij, tj -> it",
+                    GC_collect_0[l, :, :],
+                    np.exp(exponent) / (scale_tau[l, None] * K_collect_0[l, :]),
+                    optimize=True,
+                )
         else:
-            u0 = np.einsum(
-                "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
-            )
+            if there_is_beam_source:
+                u0 = np.einsum(
+                    "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
+                ) + B_collect_0[l, :].T * np.exp(-scaled_tau[None, :] / mu0)
+            else:
+                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:
@@ -301,11 +386,12 @@ def _assemble_intensity_and_fluxes(
                 K_collect_0,
                 G_inv_collect_0,
                 mu_arr,
-                _autograd_compatible,
+                is_antiderivative_wrt_tau,
+                _autograd_compatible
             )
-            u0 += _mathscr_v_contribution
+            u0 = u0 + _mathscr_v_contribution
         
-        if return_act_dscale_reclassification:
+        if _return_act_dscale_for_reclass:
             return I0_orig * np.squeeze(u0), act_dscale_reclassification
         else:
             return I0_orig * np.squeeze(u0)
@@ -320,7 +406,13 @@ def _assemble_intensity_and_fluxes(
         B_neg = B_collect_0[:, N:].T
 
 
-    def flux_up(tau):
+    def flux_up(tau, is_antiderivative_wrt_tau=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`).
+        Pass `is_antiderivative_wrt_tau = True` (defaults to `False`)
+        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])
@@ -347,13 +439,19 @@ def _assemble_intensity_and_fluxes(
                 K_collect_0,
                 G_inv_collect_0,
                 mu_arr,
-                _autograd_compatible,
+                is_antiderivative_wrt_tau,
+                _autograd_compatible
             )[:N, :]
         else:
             _mathscr_v_contribution = 0
-
+        
         if there_is_beam_source:
-            direct_beam_contribution = B_pos[:, l] * np.exp(-scaled_tau[None, :] / mu0)
+            if is_antiderivative_wrt_tau:
+                direct_beam_contribution = (
+                    B_pos[:, l] * np.exp(-scaled_tau[None, :] / mu0) / (-scale_tau[None, l] / mu0)
+                )
+            else:
+                direct_beam_contribution = B_pos[:, l] * np.exp(-scaled_tau[None, :] / mu0)
         else:
             direct_beam_contribution = 0
 
@@ -364,17 +462,36 @@ def _assemble_intensity_and_fluxes(
             ],
             axis=-1,
         )
-        u0_pos = (
-            np.einsum("tij, tj -> it", GC_pos[l, :, :], np.exp(exponent))
-            + direct_beam_contribution
-            + _mathscr_v_contribution
-        )
+        
+        if is_antiderivative_wrt_tau:
+            u0_pos = (
+                np.einsum(
+                    "tij, tj -> it",
+                    GC_pos[l, :, :],
+                    np.exp(exponent) / (scale_tau[l, None] * K_collect_0[l, :]),
+                )
+                + direct_beam_contribution
+                + _mathscr_v_contribution
+            )
+        else:
+            u0_pos = (
+                np.einsum("tij, tj -> it", GC_pos[l, :, :], np.exp(exponent))
+                + direct_beam_contribution
+                + _mathscr_v_contribution
+            )
         
         flux = 2 * pi * (mu_arr_pos * W) @ u0_pos
         return I0_orig * np.squeeze(flux)[()]
 
 
-    def flux_down(tau):
+    def flux_down(tau, is_antiderivative_wrt_tau=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
+        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`.
+        """
         tau = np.atleast_1d(tau)
         assert np.all(tau >= 0)
         assert np.all(tau <= tau_arr[-1])
@@ -401,15 +518,21 @@ def _assemble_intensity_and_fluxes(
                 K_collect_0,
                 G_inv_collect_0,
                 mu_arr,
-                _autograd_compatible,
+                is_antiderivative_wrt_tau,
+                _autograd_compatible
             )[N:, :]
         else:
             _mathscr_v_contribution = 0
         
         if there_is_beam_source:
-            direct_beam_contribution = B_neg[:, l] * np.exp(-scaled_tau[None, :] / mu0)
-            direct_beam = I0 * mu0 * np.exp(-tau / mu0)
-            direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0)
+            if is_antiderivative_wrt_tau:
+                direct_beam_contribution = B_neg[:, l] * np.exp(-scaled_tau[None, :] / mu0) / (-scale_tau[None, l] / mu0)
+                direct_beam = I0 * mu0 * np.exp(-tau / mu0) * -mu0
+                direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0) / (-scale_tau[None, l] / mu0)
+            else:
+                direct_beam_contribution = B_neg[:, l] * np.exp(-scaled_tau[None, :] / mu0)
+                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
@@ -422,11 +545,19 @@ def _assemble_intensity_and_fluxes(
             ],
             axis=-1,
         )
-        u0_neg = (
-            np.einsum("tij, tj -> it", GC_neg[l, :, :], np.exp(exponent))
-            + direct_beam_contribution
-            + _mathscr_v_contribution
-        )
+        
+        if is_antiderivative_wrt_tau:
+            u0_neg = (
+                np.einsum("tij, tj -> it", GC_neg[l, :, :], np.exp(exponent) / (scale_tau[l, None] * K_collect_0[l, :]))
+                + direct_beam_contribution
+                + _mathscr_v_contribution
+            )
+        else:
+            u0_neg = (
+                np.einsum("tij, tj -> it", GC_neg[l, :, :], np.exp(exponent))
+                + direct_beam_contribution
+                + _mathscr_v_contribution
+            )
         
         diffuse_flux = 2 * pi * (mu_arr_pos * W) @ u0_neg + direct_beam_scaled - direct_beam
         return (

PythonicDISORT/_solve_for_coeffs.py

@@ -29,9 +29,9 @@ def _solve_for_coeffs(
     NLayers_to_use_sparse,                  # Number of layers above or equal which to use sparse matrices
 ):
     """
-    In this function the unknown coefficients of the general solution to the system
-    of ODEs for each Fourier mode is solved using the boundary conditions 
-    and the product of the coefficients and the eigenvectors is returned.
+    Uses the boundary conditions to solve for the unknown coefficients 
+    of the general solution to the system of ordinary differential equations for each Fourier mode. 
+    Returns the product of the coefficients and the eigenvectors.
     This function is wrapped by the `_assemble_intensity_and_fluxes` function.
     It has many seemingly redundant arguments to maximize precomputation in the `pydisort` function.
     See the Jupyter Notebook, especially section 3, for documentation, explanation and derivation.
@@ -124,7 +124,8 @@ def _solve_for_coeffs(
         # _mathscr_v_contribution
         if there_is_iso_source and m_equals_0:
             _mathscr_v_contribution_top = -_mathscr_v(
-                        np.array([0]), np.array([0]),
+                        np.array([0]), 
+                        np.array([0]),
                         s_poly_coeffs[0:1, :],
                         Nscoeffs,
                         G_collect_m[0:1, N:, :],
@@ -139,38 +140,51 @@ def _solve_for_coeffs(
                 _mathscr_v_contribution_middle = (
                     _mathscr_v(
                         tau_arr[:-1],
-                        indices,
-                        s_poly_coeffs[:-1, :],
+                        indices + 1,
+                        s_poly_coeffs,
                         Nscoeffs,
-                        G_collect_m[:-1, :, :],
-                        K_collect_m[:-1, :],
-                        G_inv_collect_0[:-1, :, :],
+                        G_collect_m,
+                        K_collect_m,
+                        G_inv_collect_0,
                         mu_arr,
                     )
                     - _mathscr_v(
                         tau_arr[:-1],
                         indices,
-                        s_poly_coeffs[1:, :],
+                        s_poly_coeffs,
                         Nscoeffs,
-                        G_collect_m[1:, :, :],
-                        K_collect_m[1:, :],
-                        G_inv_collect_0[1:, :, :],
+                        G_collect_m,
+                        K_collect_m,
+                        G_inv_collect_0,
                         mu_arr,
                     )
-                ).ravel()
+                ).ravel(order='F')
             
-            _mathscr_v_contribution_bottom = _mathscr_v(
-                    tau_arr[-1:], np.array([0]),
+            _mathscr_v_contribution_bottom = -_mathscr_v(
+                    tau_arr[-1:], 
+                    np.array([0]),
                     s_poly_coeffs[-1:, :],
                     Nscoeffs,
-                    G_collect_m[-1:, N:, :],
+                    G_collect_m[-1:, :N, :],
                     K_collect_m[-1:, :],
                     G_inv_collect_0[-1:, :, :],
-                    mu_arr).ravel()
+                    mu_arr
+                ).ravel()
             if NBDRF > 0:
                 _mathscr_v_contribution_bottom = (
-                    R + np.eye(N)
-                ) @ _mathscr_v_contribution_bottom
+                    _mathscr_v_contribution_bottom
+                    + R
+                    @ _mathscr_v(
+                        tau_arr[-1:],
+                        np.array([0]),
+                        s_poly_coeffs[-1:, :],
+                        Nscoeffs,
+                        G_collect_m[-1:, N:, :],
+                        K_collect_m[-1:, :],
+                        G_inv_collect_0[-1:, :, :],
+                        mu_arr
+                    ).ravel()
+                )
                     
             _mathscr_v_contribution = np.concatenate(
                 [

PythonicDISORT/_solve_for_gen_and_part_sols.py

@@ -17,9 +17,9 @@ def _solve_for_gen_and_part_sols(
     there_is_iso_source,        # Is there an isotropic source?
 ):
     """
-    In this function the coefficient matrix of the system of ODEs for each Fourier mode is diagonalized 
-    and the eigenpairs are returned; the general solution to each system of ODEs is determined up to unknown coefficients.
-    The particular solutions to each system of ODEs is also determined and its coefficient vector is returned.
+    Diagonalizes the coefficient matrix of the system of ordinary differential equations (ODEs) 
+    for each Fourier mode and returns the eigenpairs which give the general solution up to unknown coefficients.
+    Also solves for the particular solution to each system of ODEs and returns its coefficient vector.
     This function is wrapped by the `_assemble_intensity_and_fluxes` function.
     It has many seemingly redundant arguments to maximize precomputation in the `pydisort` function.
     See the Jupyter Notebook, especially section 3, for documentation, explanation and derivation.
@@ -46,16 +46,18 @@ def _solve_for_gen_and_part_sols(
     | `there_is_iso_source`          | boolean                            |
     
     Notable internal variables of _solve_for_gen_and_part_sols
-    |      Variable     |             Type / Shape               |
-    | ----------------- | -------------------------------------- |
-    | `ells_all`        | `NLeg`                                 |
-    | `G_collect`       | `NFourier*NLayers x NQuad x NQuad`     | Reshaped to NFourier x NLayers x NQuad x NQuad
-    | `K_collect`       | `NFourier*NLayers x NQuad`             | Reshaped to NFourier x NLayers x NQuad
-    | `alpha_arr`       | `NFourier*NLayers x NQuad/2 x NQuad/2` | Reshaped to NFourier x NLayers x NQuad/2 x NQuad/2
-    | `beta_arr`        | `NFourier*NLayers x NQuad/2 x NQuad/2` | Reshaped to NFourier x NLayers x NQuad/2 x NQuad/2
-    | `X_tilde_arr`     | `NFourier*NLayers x NQuad`             | Reshaped to NFourier x NLayers x NQuad
-    | `B_collect`       | `NFourier*NLayers x NQuad` or `None`   | Reshaped to NFourier x NLayers x NQuad
-    | `G_inv_collect_0` | `NLayers x NQuad x NQuad` or `None`    |
+    |       Variable      |             Type / Shape               |
+    | ------------------- | -------------------------------------- |
+    | `ells_all`          | `NLeg`                                 |
+    | `G_collect`         | `NFourier*NLayers x NQuad x NQuad`     | Reshaped to NFourier x NLayers x NQuad x NQuad
+    | `K_collect`         | `NFourier*NLayers x NQuad`             | Reshaped to NFourier x NLayers x NQuad
+    | `alpha_arr`         | `NFourier*NLayers x NQuad/2 x NQuad/2` | Reshaped to NFourier x NLayers x NQuad/2 x NQuad/2
+    | `beta_arr`          | `NFourier*NLayers x NQuad/2 x NQuad/2` | Reshaped to NFourier x NLayers x NQuad/2 x NQuad/2
+    | `X_tilde_arr`       | `NFourier*NLayers x NQuad`             | Reshaped to NFourier x NLayers x NQuad
+    | `B_collect`         | `NFourier*NLayers x NQuad` or `None`   | Reshaped to NFourier x NLayers x NQuad
+    | `eigenvecs_GpG_arr` | `NFourier*NLayers x NQuad/2 x NQuad    |
+    | `eigenvecs_GmG_arr` | `NFourier*NLayers x NQuad/2 x NQuad    |
+    | `G_inv_collect_0`   | `NLayers x NQuad x NQuad` or `None`    |
 
     """
     ############################### Assemble system and diagonalize coefficient matrix #########################################
@@ -109,7 +111,7 @@ def _solve_for_gen_and_part_sols(
             if np.any(weighted_asso_Leg_coeffs_l > 0) and np.all(
                 np.isfinite(asso_leg_term_pos)
             ):  
-                # Generate mathscr_D and mathscr_X (BDRF terms)
+                # Generate D and X (phase function terms)
                 # --------------------------------------------------------------------------------------------------------------------------
                 D_temp = weighted_asso_Leg_coeffs_l[None, :] * asso_leg_term_pos.T
                 D_pos = (scaled_omega_l / 2) * D_temp @ asso_leg_term_pos
@@ -190,7 +192,7 @@ def _solve_for_gen_and_part_sols(
             G_inv_collect_0[no_shortcut_indices_0, :, :] = G_inv_arr[: len(no_shortcut_indices_0), :, :]
         # --------------------------------------------------------------------------------------------------------------------------
 
-        # Particular solution for the sunbeam source (refer to Section 3.6.1 of the Comprehensive Documentation)
+        # Particular solution for the direct beam source (refer to Section 3.6.1 of the Comprehensive Documentation)
         # --------------------------------------------------------------------------------------------------------------------------
         if there_is_beam_source:
             X_tilde_arr = X_tilde_arr[: len(no_shortcut_indices), :]

PythonicDISORT/pydisort.py

@@ -28,6 +28,7 @@ def pydisort(
     """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.
     
         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.
@@ -64,7 +65,7 @@ def pydisort(
         Dirichlet boundary condition for the downward direction.
     only_flux : optional, bool
         Do NOT compute the intensity function?
-    f_arr : optional, array
+    f_arr : optional, array or float
         Fractional scattering into peak for each atmospheric layer.
     NT_cor : optional, bool
         Perform Nakajima-Tanaka intensity corrections?
@@ -79,25 +80,36 @@ def pydisort(
 
     Returns
     -------
-    mu_arr : array
+    mu_arr : array or float
         All `mu` (cosine of polar angle) quadrature nodes.
-    Fp(tau) : function, function
-        (Energetic) Flux function with argument `tau` (type: array) for positive (upward) `mu` values.
-        Returns the diffuse flux magnitudes (type: array).
+    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`.
     Fm(tau) : function
-        (Energetic) Flux function with argument `tau` (type: array) for negative (downward) `mu` values.
-        Returns a tuple of the diffuse and direct flux magnitudes respectively (type: (array, array)).
+        (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`.
     u0(tau) : function
-        Zeroth Fourier mode of the intensity with argument `tau` (type: array).
+        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`.
     u(tau, phi) : function, optional
-        Intensity function with arguments `(tau, phi, return_Fourier_error=False)` of types `(array, array, bool)`.
+        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.
-        The optional flag `return_Fourier_error` determines whether the function will also return
-        the Cauchy / Fourier convergence evaluation (type: float) for the last Fourier term.
+        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`.
 
     """
     
@@ -310,9 +322,9 @@ def pydisort(
             _autograd_compatible,
         )
         
-        # TMS correction
+        # TMS correction for the intensity (see Section 3.7.2)
         # --------------------------------------------------------------------------------------------------------------------------
-        def TMS_correction(tau, phi):
+        def TMS_correction(tau, phi, is_antiderivative_wrt_tau):
             Ntau = len(tau)
             Nphi = len(phi)
             # Atmospheric layer indices
@@ -357,20 +369,36 @@ def pydisort(
             )
             # --------------------------------------------------------------------------------------------------------------------------
             
-            TMS_correction_pos = (
-                np.exp(-scaled_tau / mu0)[None, :]
-                - np.exp(
-                    (scaled_tau - scaled_tau_arr_l)[None, :] / mu_arr_pos[:, None]
-                    - scaled_tau_arr_l[None, :] / mu0
+            if is_antiderivative_wrt_tau:
+                TMS_correction_pos = (
+                    np.exp(-scaled_tau / mu0)[None, :] / (-scale_tau[None, l] / mu0)
+                    - np.exp(
+                        (scaled_tau - scaled_tau_arr_l)[None, :] / mu_arr_pos[:, None]
+                        - scaled_tau_arr_l[None, :] / mu0
+                    ) / (scale_tau[None, l] / mu_arr_pos[:, None])
                 )
-            )
-            TMS_correction_neg = (
-                np.exp(-scaled_tau / mu0)[None, :]
-                - np.exp(
-                    (scaled_tau_arr_lm1 - scaled_tau)[None, :] / mu_arr_pos[:, None]
-                    - scaled_tau_arr_lm1[None, :] / mu0
+                TMS_correction_neg = (
+                    np.exp(-scaled_tau / mu0)[None, :] / (-scale_tau[None, l] / mu0)
+                    - np.exp(
+                        (scaled_tau_arr_lm1 - scaled_tau)[None, :] / mu_arr_pos[:, None]
+                        - scaled_tau_arr_lm1[None, :] / mu0
+                    ) / (-scale_tau[None, l] / mu_arr_pos[:, None])
+                )
+            else:
+                TMS_correction_pos = (
+                    np.exp(-scaled_tau / mu0)[None, :]
+                    - np.exp(
+                        (scaled_tau - scaled_tau_arr_l)[None, :] / mu_arr_pos[:, None]
+                        - scaled_tau_arr_l[None, :] / mu0
+                    )
+                )
+                TMS_correction_neg = (
+                    np.exp(-scaled_tau / mu0)[None, :]
+                    - np.exp(
+                        (scaled_tau_arr_lm1 - scaled_tau)[None, :] / mu_arr_pos[:, None]
+                        - scaled_tau_arr_lm1[None, :] / mu0
+                    )
                 )
-            )
             
             ## Contribution from other layers
             # TODO: Despite being vectorized and despite many attempts at further optimization,
@@ -392,72 +420,114 @@ def pydisort(
 
                 layers_arr_repeat = np.repeat(layers_arr, Ntau)
                 scaled_tau_tile = np.tile(scaled_tau, NLayers)
-                scaled_tau_arr_with_0_repeat = np.repeat(scaled_tau_arr_with_0, Ntau)
 
                 if any_pos_contribution:
                     contribution_from_each_other_layer_pos = np.zeros((N, NLayers * Ntau, Nphi))
-                    scaled_tau_l_tile_pos = scaled_tau_tile[pos_contribution_mask]
-                    mathscr_B_l_repeat_posmasked = mathscr_B[
-                        :N, layers_arr_repeat[pos_contribution_mask], :
-                    ]
-                    scaled_tau_arr_with_0_l_tile_posmasked = scaled_tau_arr_with_0_repeat[:-Ntau][
-                        pos_contribution_mask
-                    ]
-                    scaled_tau_arr_with_0_lp1_tile_posmasked = scaled_tau_arr_with_0_repeat[Ntau:][
-                        pos_contribution_mask
-                    ]
-
-                    contribution_from_each_other_layer_pos[:, pos_contribution_mask, :] = (
-                        mathscr_B_l_repeat_posmasked
-                        * (
-                            np.exp(
-                                (scaled_tau_l_tile_pos - scaled_tau_arr_with_0_l_tile_posmasked)
-                                / mu_arr_pos[:, None]
-                                - scaled_tau_arr_with_0_l_tile_posmasked / mu0
-                            )
-                            - np.exp(
-                                (scaled_tau_l_tile_pos - scaled_tau_arr_with_0_lp1_tile_posmasked)
-                                / mu_arr_pos[:, None]
-                                - scaled_tau_arr_with_0_lp1_tile_posmasked / mu0
-                            )
-                        )[:, :, None]
-                    )
-                    contribution_from_other_layers_pos = np.sum(
-                        contribution_from_each_other_layer_pos.reshape((N, NLayers, Ntau, Nphi)),
-                        axis=1,
-                    )
+                    scaled_tau_tile_pos = scaled_tau_tile[pos_contribution_mask]
+                    
+                    layers_arr_repeat_pos = layers_arr_repeat[pos_contribution_mask]
+                    mathscr_B_repeat_pos = mathscr_B[:N, layers_arr_repeat_pos, :]
+                    scaled_tau_arr_with_0_l_repeat_pos = scaled_tau_arr_with_0[:-1][layers_arr_repeat_pos]
+                    scaled_tau_arr_with_0_lp1_repeat_pos = scaled_tau_arr_with_0[1:][layers_arr_repeat_pos]
+                    
+                    if is_antiderivative_wrt_tau:
+                        scale_tau_repeat_pos = scale_tau[layers_arr_repeat_pos]
+                        contribution_from_each_other_layer_pos[:, pos_contribution_mask, :] = (
+                            mathscr_B_repeat_pos
+                            * (
+                                (
+                                    np.exp(
+                                        (scaled_tau_tile_pos - scaled_tau_arr_with_0_l_repeat_pos)
+                                        / mu_arr_pos[:, None]
+                                        - scaled_tau_arr_with_0_l_repeat_pos / mu0
+                                    )
+                                    - np.exp(
+                                        (scaled_tau_tile_pos - scaled_tau_arr_with_0_lp1_repeat_pos)
+                                        / mu_arr_pos[:, None]
+                                        - scaled_tau_arr_with_0_lp1_repeat_pos / mu0
+                                    )
+                                )
+                                / (scale_tau_repeat_pos / mu_arr_pos[:, None])
+                            )[:, :, None]
+                        )
+                        contribution_from_other_layers_pos = np.sum(
+                            contribution_from_each_other_layer_pos.reshape((N, NLayers, Ntau, Nphi)),
+                            axis=1,
+                        )
+                    else:
+                        contribution_from_each_other_layer_pos[:, pos_contribution_mask, :] = (
+                            mathscr_B_repeat_pos
+                            * (
+                                np.exp(
+                                    (scaled_tau_tile_pos - scaled_tau_arr_with_0_l_repeat_pos)
+                                    / mu_arr_pos[:, None]
+                                    - scaled_tau_arr_with_0_l_repeat_pos / mu0
+                                )
+                                - np.exp(
+                                    (scaled_tau_tile_pos - scaled_tau_arr_with_0_lp1_repeat_pos)
+                                    / mu_arr_pos[:, None]
+                                    - scaled_tau_arr_with_0_lp1_repeat_pos / mu0
+                                )
+                            )[:, :, None]
+                        )
+                        contribution_from_other_layers_pos = np.sum(
+                            contribution_from_each_other_layer_pos.reshape((N, NLayers, Ntau, Nphi)),
+                            axis=1,
+                        )
+                        
                 if any_neg_contribution:
                     contribution_from_each_other_layer_neg = np.zeros((N, NLayers * Ntau, Nphi))
-                    scaled_tau_l_tile_neg = scaled_tau_tile[neg_contribution_mask]
-                    mathscr_B_l_repeat_negmasked = mathscr_B[
-                        N:, layers_arr_repeat[neg_contribution_mask], :
-                    ]
-                    scaled_tau_arr_with_0_l_tile_negmasked = scaled_tau_arr_with_0_repeat[:-Ntau][
-                        neg_contribution_mask
-                    ]
-                    scaled_tau_arr_with_0_lp1_tile_negmasked = scaled_tau_arr_with_0_repeat[Ntau:][
-                        neg_contribution_mask
-                    ]
+                    scaled_tau_tile_neg = scaled_tau_tile[neg_contribution_mask]
+                    
+                    layers_arr_repeat_neg = layers_arr_repeat[neg_contribution_mask]
+                    mathscr_B_repeat_neg = mathscr_B[N:, layers_arr_repeat_neg, :]
+                    scaled_tau_arr_with_0_l_repeat_neg = scaled_tau_arr_with_0[:-1][layers_arr_repeat_neg]
+                    scaled_tau_arr_with_0_lp1_repeat_neg = scaled_tau_arr_with_0[1:][layers_arr_repeat_neg]
 
-                    contribution_from_each_other_layer_neg[:, neg_contribution_mask, :] = (
-                        mathscr_B_l_repeat_negmasked
-                        * (
-                            np.exp(
-                                (scaled_tau_arr_with_0_lp1_tile_negmasked - scaled_tau_l_tile_neg)
-                                / mu_arr_pos[:, None]
-                                - scaled_tau_arr_with_0_lp1_tile_negmasked / mu0
-                            )
-                            - np.exp(
-                                (scaled_tau_arr_with_0_l_tile_negmasked - scaled_tau_l_tile_neg)
-                                / mu_arr_pos[:, None]
-                                - scaled_tau_arr_with_0_l_tile_negmasked / mu0
-                            )
-                        )[:, :, None]
-                    )
-                    contribution_from_other_layers_neg = np.sum(
-                        contribution_from_each_other_layer_neg.reshape((N, NLayers, Ntau, Nphi)),
-                        axis=1,
-                    )
+                    if is_antiderivative_wrt_tau:
+                        scale_tau_repeat_neg = scale_tau[layers_arr_repeat_neg]
+                        contribution_from_each_other_layer_neg[:, neg_contribution_mask, :] = (
+                            mathscr_B_repeat_neg
+                            * (
+                                (
+                                    np.exp(
+                                        (scaled_tau_arr_with_0_lp1_repeat_neg - scaled_tau_tile_neg)
+                                        / mu_arr_pos[:, None]
+                                        - scaled_tau_arr_with_0_lp1_repeat_neg / mu0
+                                    )
+                                    - np.exp(
+                                        (scaled_tau_arr_with_0_l_repeat_neg - scaled_tau_tile_neg)
+                                        / mu_arr_pos[:, None]
+                                        - scaled_tau_arr_with_0_l_repeat_neg / mu0
+                                    )
+                                )
+                                / (-scale_tau_repeat_pos / mu_arr_pos[:, None])
+                            )[:, :, None]
+                        )
+                        contribution_from_other_layers_neg = np.sum(
+                            contribution_from_each_other_layer_neg.reshape((N, NLayers, Ntau, Nphi)),
+                            axis=1,
+                        )
+                    else:
+                        contribution_from_each_other_layer_neg[:, neg_contribution_mask, :] = (
+                            mathscr_B_repeat_neg
+                            * (
+                                np.exp(
+                                    (scaled_tau_arr_with_0_lp1_repeat_neg - scaled_tau_tile_neg)
+                                    / mu_arr_pos[:, None]
+                                    - scaled_tau_arr_with_0_lp1_repeat_neg / mu0
+                                )
+                                - np.exp(
+                                    (scaled_tau_arr_with_0_l_repeat_neg - scaled_tau_tile_neg)
+                                    / mu_arr_pos[:, None]
+                                    - scaled_tau_arr_with_0_l_repeat_neg / mu0
+                                )
+                            )[:, :, None]
+                        )
+                        contribution_from_other_layers_neg = np.sum(
+                            contribution_from_each_other_layer_neg.reshape((N, NLayers, Ntau, Nphi)),
+                            axis=1,
+                        )
 
                 if _autograd_compatible:
                     if not any_pos_contribution:
@@ -481,7 +551,7 @@ def pydisort(
                 return solution
         # --------------------------------------------------------------------------------------------------------------------------
 
-        # IMS correction
+        # IMS correction for the intensity (see Section 3.7.2)
         # --------------------------------------------------------------------------------------------------------------------------
         sum1 = np.sum(omega_arr * tau_arr)
         omega_avg = sum1 / np.sum(tau_arr)
@@ -495,15 +565,22 @@ def pydisort(
         )
         scaled_mu0 = mu0 / (1 - omega_avg * f_avg)
 
-        def IMS_correction(tau, phi):
+        def IMS_correction(tau, phi, is_antiderivative_wrt_tau):
             nu = subroutines.atleast_2d_append(
                 subroutines.calculate_nu(-mu_arr_pos, phi, -mu0, phi0)
             )
             x = 1 / mu_arr_pos - 1 / scaled_mu0
-            chi = (1 / (mu_arr_pos * scaled_mu0 * x))[:, None] * (
-                (tau[None, :] - 1 / x[:, None]) * np.exp(-tau / scaled_mu0)[None, :]
-                + np.exp(-tau[None, :] / mu_arr_pos[:, None]) / x[:, None]
-            )
+            if is_antiderivative_wrt_tau:
+                chi = (
+                    (scaled_mu0 - x[:, None] * scaled_mu0 * (scaled_mu0 + tau)[None, :])
+                    * np.exp(-tau / scaled_mu0)[None, :]
+                    - mu_arr_pos[:, None] * np.exp(-tau[None, :] / mu_arr_pos[:, None])
+                ) / (mu_arr_pos * scaled_mu0 * x**2)[:, None]
+            else:
+                chi = (
+                    (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)
@@ -518,26 +595,26 @@ def pydisort(
 
         # The corrected intensity
         # --------------------------------------------------------------------------------------------------------------------------
-        def u_corrected(tau, phi, return_Fourier_error=False):
+        def u_corrected(tau, phi, return_Fourier_error=False, is_antiderivative_wrt_tau=False):
             tau = np.atleast_1d(tau)
             phi = np.atleast_1d(phi)
-            NT_corrections = TMS_correction(tau, phi)
+            NT_corrections = TMS_correction(tau, phi, is_antiderivative_wrt_tau)
 
             if _autograd_compatible:
                 NT_corrections = NT_corrections + np.concatenate(
-                    [np.zeros((N, len(tau), len(phi))), IMS_correction(tau, phi)], axis=0
+                    [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)
+                NT_corrections[N:, :, :] += IMS_correction(tau, phi, is_antiderivative_wrt_tau)
                
             if return_Fourier_error:
-                u_star_outputs = u_star(tau, phi, True)
+                u_star_outputs = u_star(tau, phi, True, is_antiderivative_wrt_tau)
                 return (
                     u_star_outputs[0] + I0_orig * np.squeeze(NT_corrections),
                     u_star_outputs[1],
                 )
             else:
-                return u_star(tau, phi, False) + I0_orig * np.squeeze(NT_corrections)
+                return u_star(tau, phi, False, is_antiderivative_wrt_tau) + I0_orig * np.squeeze(NT_corrections)
         # --------------------------------------------------------------------------------------------------------------------------
 
         return mu_arr, flux_up, flux_down, u0, u_corrected

PythonicDISORT/subroutines.py

@@ -1,7 +1,6 @@
 import numpy as np
 import scipy as sc
 from math import pi
-from numpy.polynomial.legendre import leggauss
 
 
 def transform_interval(arr, c, d, a, b):
@@ -71,7 +70,7 @@ def calculate_nu(mu, phi, mu_p, phi_p):
 
     Returns
     -------
-    ndarray
+    nu : ndarray
         Cosine of scattering angles which axes capture variation with `mu, phi, mu_p, phi_p` respectively.
 
     """
@@ -104,7 +103,7 @@ def Gauss_Legendre_quad(N, c=0, d=1):
         Quadrature weights.
 
     """
-    mu_arr_pos, W = leggauss(int(N))
+    mu_arr_pos, W = np.polynomial.legendre.leggauss(int(N))
 
     return transform_interval(mu_arr_pos, c, d, -1, 1), transform_weights(W, c, d, -1, 1)
 
@@ -224,8 +223,10 @@ def atleast_2d_append(*arys):
         
 def generate_diff_act_flux_funcs(u0):
     """Generates the up and down diffuse actinic flux functions respectively.
-    This is a use case of the ``u0`` function that is an output of ``pydisort``.
+    This is a use case of the ``u0`` function which is an output of ``pydisort``.
     The reclassification of delta-scaled actinic flux is automatically performed.
+    The computation of actinic fluxes is similar to that of energetic fluxes
+    and the latter is discussed in section 3.8 of the Comprehensive Documentation.
 
     Parameters
     ----------
@@ -235,36 +236,102 @@ def generate_diff_act_flux_funcs(u0):
 
     Returns
     -------
-    function
+    Fp_act(tau) : function
         Actinic flux function with argument `tau` (type: array) for positive (upward) `mu` values.
         Returns the diffuse flux magnitudes (type: array).
-    function
+        Pass `is_antiderivative_wrt_tau = True` (defaults to `False`)
+        to switch to an antiderivative of the function with respect to `tau`.
+    Fm_act(tau) : function
         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`.
 
     """
     N = np.shape(u0(0))[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
-    flux_act_up = lambda tau: 2 * pi * GL_weights @ u0(tau)[:N]
-    def flux_act_down_diffuse(tau):
-        u0_cache, act_dscale_reclassification = u0(tau, True)
+    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_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 result_without_reclassification + act_dscale_reclassification
+        return np.squeeze(result_without_reclassification + act_dscale_reclassification)[()]
     
     return flux_act_up, flux_act_down_diffuse
 
 
-def _mathscr_v(tau,                         # Input optical depths
-                l,                          # Layer index of each input optical depth
-                s_poly_coeffs,              # Polynomial coefficients of isotropic source
-                Nscoeffs,                   # Number of isotropic source polynomial coefficients
-                G,                          # Eigenvector matrices
-                K,                          # Eigenvalues
-                G_inv,                      # Inverse of eigenvector matrix
-                mu_arr,                     # Quadrature nodes for both hemispheres
-                _autograd_compatible=False  # Should the output functions be compatible with autograd?
+def interpolate(mu_arr_pos, 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`.
+
+    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 
+        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`.
+
+    """
+    N = len(mu_arr_pos)
+    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)
+        mu = np.atleast_1d(mu)
+
+        if return_Fourier_error:
+            u_cache, Fourier_error = u(tau, phi, True, is_antiderivative_wrt_tau)
+        else:
+            u_cache = u(tau, phi, False, 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:])
+        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:
+            return np.squeeze(results)[()], Fourier_error
+        else:
+            return np.squeeze(results)[()]
+
+    return u_interpol
+                
+
+def _mathscr_v(tau,                             # Input optical depths
+                l,                              # Layer index of each input optical depth
+                s_poly_coeffs,                  # Polynomial coefficients of isotropic source
+                Nscoeffs,                       # Number of isotropic source polynomial coefficients
+                G,                              # Eigenvector matrices
+                K,                              # Eigenvalues
+                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?
                 ):
     """Particular solution for isotropic internal sources.
     Refer to Section 3.6.1 of the Comprehensive Documentation.
@@ -273,16 +340,17 @@ def _mathscr_v(tau,                         # Input optical depths
     and `_solve_for_coeffs` functions which call it.
     
     Arguments of _mathscr_v
-    |        Variable        |                 Shape                 |
-    | ---------------------- | ------------------------------------- |
-    | `tau`                  | `Ntau`                                |
-    | `l`                    | `Ntau`                                |
-    | `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`                               |
-    | `_autograd_compatible` | boolean                               |
+    |          Variable           |                 Shape                 |
+    | --------------------------- | ------------------------------------- |
+    | `tau`                       | `Ntau`                                |
+    | `l`                         | `Ntau`                                |
+    | `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                | 
@@ -340,11 +408,17 @@ def _mathscr_v(tau,                         # Input optical depths
                 0,
             ),
         )
-
+    
+    powers = np.flip(np.arange(Nscoeffs))[None, :]
+    if is_antiderivative_wrt_tau:
+        tau_poly = tau[:, None] ** (powers + 1) / (powers + 1)
+    else:
+        tau_poly = tau[:, None] ** powers
+        
     return np.einsum(
         "tik, tc, ctk, tkj, j -> it",
         G[l, :, :],
-        tau[:, None] ** np.flip(np.arange(Nscoeffs))[None, :],
+        tau_poly,
         mathscr_v_coeffs[:, l, :],
         G_inv[l, :, :],
         1 / mu_arr,
@@ -365,6 +439,7 @@ def _compare(results, mu_to_compare, reorder_mu, flux_up, flux_down, u):
     which are stored in .npz files, against results from PythonicDISORT. Used in our PyTests.
 
     """
+    # Load saved results from Stamnes' DISORT
     uu = results["uu"]
     flup = results["flup"]
     rfldn = results["rfldn"]
@@ -438,6 +513,7 @@ def _compare(results, mu_to_compare, reorder_mu, flux_up, flux_down, u):
     diff_tau_pt = tau_test_arr[max_diff_tau_index]
     ratio_tau_pt = tau_test_arr[max_ratio_tau_index]
     print("Intensities")
+    print()
     print("At tau = " + str(diff_tau_pt))
     print("Max pointwise difference =", np.max(diff[:, max_diff_tau_index, :]))
     print("At tau = " + str(ratio_tau_pt))

{PythonicDISORT-0.7.0.dist-info → PythonicDISORT-0.9.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PythonicDISORT
-Version: 0.7.0
+Version: 0.9.0
 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
@@ -20,42 +20,45 @@ Provides-Extra: pytest
 Requires-Dist: pytest >=6.2.5 ; extra == 'pytest'
 
 # Introduction
-The PythonicDISORT package is a Discrete Ordinates Solver for the (1D) Radiative Transfer Equation in a single or multi-layer plane-parallel atmosphere.
+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: 
-delta-M scaling, Nakajima-Tanaka (NT) corrections, only flux option, isotropic internal sources (thermal sources), 
-Dirichlet boundary conditions (diffuse flux boundary sources), Bi-Directional Reflectance Function (BDRF) for surface reflection, and more.
-In addition, upon the request of a user, we added a subroutine to calculate actinic fluxes and we are open to further feature requests as well as feedback. 
+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, 
+interpolation with respect to polar angle and more.
+In addition, we added a subroutine to calculate actinic fluxes to satisfy a user request, and integration with respect to tau was also added.
+Further feature requests as well as feedback are welcome.
 
 You may contact me, Dion, through dh3065@columbia.edu.
 
-Our **GitHub repository:** https://github.com/LDEO-CREW/Pythonic-DISORT also includes our F2PY-wrapped Stamnes DISORT (version 4.0.99) in the `disort4.0.99_f2py` directory.
-The original was downloaded from http://www.rtatmocn.com/disort/. The wrapper is inspired by https://github.com/kconnour/pyRT_DISORT.
+The **GitHub repository** is https://github.com/LDEO-CREW/Pythonic-DISORT.
 
 # Documentation
 https://pythonic-disort.readthedocs.io/en/latest/
 
 Also see the accompanying Jupyter Notebook `Pythonic-DISORT.ipynb` in the `docs` directory
 of our GitHub repository.
-The Jupyter Notebook provides comprehensive documentation, suggested inputs, explanations, 
+This Jupyter Notebook provides comprehensive documentation, suggested inputs, explanations, 
 mathematical derivations and verification tests.
-We highly recommend reading the non-optional parts of sections 1 and 2 before use.
+It is highly recommended that new users read the non-optional parts of sections 1 and 2.
 
 ## PyTest and examples of how to use PythonicDISORT
 
-Not only do we have verification tests in the Jupyter Notebook, we have also recreated most of the test problems in Stamnes' `disotest.f90`; 
-`disotest.f90` is included in the `disort4.0.99_f2py` directory of our GitHub repository. In these tests, the solutions from PythonicDISORT are compared 
-against solutions from our F2PY-wrapped Stamnes' DISORT (version 4.0.99). With PyTest installed, execute the console command `pytest` 
-in the `pydisotest` directory to run these tests. The `pydisotest` directory also contains Jupyter Notebooks, one for each test, 
-to show how the tests were implemented. These notebooks double up as examples of how to use PythonicDISORT.
+Not only are there verification tests in `Pythonic-DISORT.ipynb`, 
+most of the test problems in Stamnes' `disotest.f90` (download DISORT 4.0.99 from http://www.rtatmocn.com/disort/) have also been recreated.
+In these tests, the solutions from PythonicDISORT are compared against solutions 
+from a F2PY-wrapped Stamnes' DISORT (version 4.0.99; wrapper inspired by https://github.com/kconnour/pyRT_DISORT). With PyTest installed, execute the console command `pytest` 
+in the `pydisotest` directory to run these tests. The `pydisotest` directory also contains Jupyter Notebooks to show the implementation of each test.
+These notebooks double up as examples of how to use PythonicDISORT.
 
 # Installation
 
 * From PyPI: `pip install PythonicDISORT`
-* From Conda-forge: (TODO: we need to first publish on Conda-forge)
+* From Conda-forge: (TODO: need to first publish on Conda-forge)
 * By cloning repository: `pip install .` in the `Pythonic-DISORT` directory; `pip install -r all_optional_dependencies.txt` to install all optional dependencies (see *Requirements to run PythonicDISORT*)
 
 ## Requirements to run PythonicDISORT
@@ -64,21 +67,18 @@ to show how the tests were implemented. These notebooks double up as examples of
 * `scipy >= 1.8.0`
 * (OPTIONAL) `pytest >= 6.2.5` (Required to use the command `pytest`, see *PyTest and examples of how to use PythonicDISORT*)
 
-## Additional requirements to run the Jupyter Notebook
+## (OPTIONAL) Additional requirements to run the Jupyter Notebook
 * `autograd >= 1.5`
 * `jupyter > 1.0.0`
 * `notebook > 6.5.2`
+* `matplotlib >= 3.6.0`
 
-In addition, our F2PY-wrapped Stamnes' DISORT (in the `disort4.0.99_f2py` directory), or equivalent,
-must be set up to run the last section (section 6).
+In addition, a F2PY-wrapped Stamnes' DISORT, or equivalent, is required to properly run the last section (section 6).
 
 ## Compatibility
 
-The PythonicDISORT package *should* be system agnostic given its minimal dependencies and pure Python code.
-We do not guarantee that our Jupyter Notebook and F2PY-wrapped Stamnes' DISORT will work on other systems though.
-The latter will almost certainly need user edits to work and in any case it requires FORTRAN compilers which
-are not included in our GitHub repository. Everything in the repository was built and tested on Windows 11 but 
-not yet tested on other systems.
+The PythonicDISORT package should be system agnostic given its minimal dependencies and pure Python code.
+Everything in the repository was built and tested on Windows 11.
 
 # Acknowledgements
 
@@ -106,6 +106,8 @@ Science and Technology Center (STC) (Award #2019625) under which this package wa
 
 10) Trefethen, L. N.. 1996. *Finite difference and spectral methods for ordinary and partial differential equations.* https://people.maths.ox.ac.uk/trefethen/pdetext.html.
 
-11) T. Nakajima and M. Tanaka. 1988. *Algorithms for radiative intensity calculations in moderately thick atmospheres using a truncation approximation.* https://www.sciencedirect.com/science/article/pii/0022407388900313.
+11) Knut Stamnes. 1982. *On the computation of angular distributions of radiation in planetary atmospheres.* https://www.sciencedirect.com/science/article/pii/0022407382900966.
 
-12) K. Connour and A. Stcherbinine. 2022. *pyRT_DISORT.* https://github.com/kconnour/pyRT_DISORT.
+12) T. Nakajima and M. Tanaka. 1988. *Algorithms for radiative intensity calculations in moderately thick atmospheres using a truncation approximation.* https://www.sciencedirect.com/science/article/pii/0022407388900313.
+
+13) Connour, Kyle and Wolff, Michael. 2020. *pyRT_DISORT: A pre-processing front-end to help make DISORT simulations easier in Python.* https://github.com/kconnour/pyRT_DISORT.

PythonicDISORT-0.9.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
+PythonicDISORT/_assemble_intensity_and_fluxes.py,sha256=SFX5SMUhKNt_MU-0t3kv5lLXxc1EzVK2EIDLnggoIEw,27206
+PythonicDISORT/_solve_for_coeffs.py,sha256=ZGThwif73zAlhzXYdVzV2nJxdxRioq4MIFsni3_5-Hc,19594
+PythonicDISORT/_solve_for_gen_and_part_sols.py,sha256=9ScaxY7SH3lYaE2UBnEIsmsDeXdjcf74_QWaeqyOIRs,12973
+PythonicDISORT/pydisort.py,sha256=HmAdvbDHQoHQy1xXToJmT3X5MGRudMKoI0L7zAaOd-8,33188
+PythonicDISORT/subroutines.py,sha256=2cJLQc46ctAiqY_qARlsxQZPYVcBOwkXf0-diCtnGXo,19446
+PythonicDISORT-0.9.0.dist-info/LICENSE.md,sha256=q1i1QVh0OLzK0fj6kaNRNWsM_KZca66JwWgWwbuCPlA,1092
+PythonicDISORT-0.9.0.dist-info/METADATA,sha256=dLzNm79C3rWzVMzC2GIboGpMloxs4Fiy3RpQgOMIMP4,7102
+PythonicDISORT-0.9.0.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+PythonicDISORT-0.9.0.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
+PythonicDISORT-0.9.0.dist-info/RECORD,,

{PythonicDISORT-0.7.0.dist-info → PythonicDISORT-0.9.0.dist-info}/WHEEL

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

PythonicDISORT-0.7.0.dist-info/RECORD

@@ -1,11 +0,0 @@
-PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
-PythonicDISORT/_assemble_intensity_and_fluxes.py,sha256=IH1OYHfMLh1LrNiwSCnaJduELvHpStJgXqTycQ869X8,20410
-PythonicDISORT/_solve_for_coeffs.py,sha256=8x7bDARYjS_fnoFGt5LRnm6yUtjyPa_YQ2Mj0LA2AFY,19162
-PythonicDISORT/_solve_for_gen_and_part_sols.py,sha256=okJFq60712DN48-jdodipvP5S11yYK_yobDVpg0kTi8,12841
-PythonicDISORT/pydisort.py,sha256=xjhiRnwmVIPGKJqEuapwLwGLicL3mjizO-w3r5kaP1M,27880
-PythonicDISORT/subroutines.py,sha256=oqPIluJaFF88d7N7iS72wFOKfiYAmT3uA33tpPjuLP0,15857
-PythonicDISORT-0.7.0.dist-info/LICENSE.md,sha256=q1i1QVh0OLzK0fj6kaNRNWsM_KZca66JwWgWwbuCPlA,1092
-PythonicDISORT-0.7.0.dist-info/METADATA,sha256=1V3vyKC8Qa9uEH_WvCfdUnpU28HqOEnb3PdPCIR7Ij4,7230
-PythonicDISORT-0.7.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-PythonicDISORT-0.7.0.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
-PythonicDISORT-0.7.0.dist-info/RECORD,,