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

PythonicDISORT/{_assemble_solution_functions.py → _assemble_intensity_and_fluxes.py}

@@ -1,51 +1,107 @@
 from PythonicDISORT.subroutines import _mathscr_v
-from PythonicDISORT._diagonalize import _diagonalize
-from PythonicDISORT._solve_for_coefs import _solve_for_coefs
+from PythonicDISORT._solve_for_gen_and_part_sols import _solve_for_gen_and_part_sols
+from PythonicDISORT._solve_for_coeffs import _solve_for_coeffs
 
 import numpy as np
 from math import pi
 
-def _assemble_solution_functions(
-    scaled_omega_arr,
-    tau_arr,
-    scaled_tau_arr_with_0,
-    mu_arr_pos, mu_arr,
-    M_inv, W,
-    N, NQuad, NLeg, NFourier,
-    NLayers, NBDRF,
-    atmos_is_multilayered,
-    weighted_scaled_Leg_coeffs,
-    BDRF_Fourier_modes,
-    mu0, I0, I0_orig, phi0,
-    there_is_beam_source,
-    b_pos, b_neg,
-    b_pos_is_scalar, b_neg_is_scalar,
-    s_poly_coeffs,
-    Nscoeffs,
-    there_is_iso_source,
-    scale_tau,
-    only_flux,
-    use_sparse_NLayers,
-    _is_compatible_with_autograd,
+def _assemble_intensity_and_fluxes(
+    scaled_omega_arr,                   # Delta-scaled single-scattering albedos
+    tau_arr,                            # Lower boundary of layers
+    scaled_tau_arr_with_0,              # Delta-scaled lower boundary of layers with 0 inserted in front
+    mu_arr_pos, mu_arr,                 # Quadrature nodes for 1) upper 2) both hemispheres
+    M_inv, W,                           # 1) 1 / mu; 2) quadrature weights for each hemisphere
+    N, NQuad, NLeg,                     # Number of 1) upper 2) both hemispheres quadrature nodes; 3) phase function Legendre coefficients 
+    NFourier, NLayers, NBDRF,           # Number of 1) intensity Fourier modes; 2) layers; 3) BDRF Fourier modes
+    is_atmos_multilayered,              # Is the atmosphere multilayered?
+    weighted_scaled_Leg_coeffs,         # Weighted and delta-scaled Legendre coefficients
+    BDRF_Fourier_modes,                 # BDRF Fourier modes
+    mu0, I0, I0_orig, phi0,             # Properties of direct beam. If I0 was rescaled to 1, I0_orig is the original I0
+    there_is_beam_source,               # Is there a beam source?
+    b_pos, b_neg,                       # Dirichlet BCs
+    b_pos_is_scalar, b_neg_is_scalar,   # Is each Dirichlet BCs scalar (isotropic)?
+    s_poly_coeffs,                      # Polynomial coefficients of isotropic source
+    Nscoeffs,                           # Number of isotropic source polynomial coefficients
+    there_is_iso_source,                # Is there an isotropic source?
+    scale_tau,                          # Delta-scale factor for tau
+    only_flux,                          # Only compute 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?
 ):  
-    """This function is wrapped by the `pydisort` function.
+    """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.
+    This function is wrapped by the `pydisort` function.
     It should be called through `pydisort` and never directly.
-    It has many seemingly redundant arguments to maximize precomputation in `pydisort`.
-    These arguments are passed on to the `_diagonalize` and `_solve_for_coefs` functions 
+    It has many seemingly redundant arguments, which are described below,
+    to maximize precomputation in `pydisort`
+    These arguments are passed on to the `_solve_for_gen_and_part_sols` and `_solve_for_coeffs` functions 
     which this function wraps. See the Jupyter Notebook, especially section 3, for 
     documentation, explanation and derivation.
     
+    Arguments of _assemble_intensity_and_fluxes
+    |            Variable            |            Type / Shape            |
+    | ------------------------------ | ---------------------------------- |
+    | `scaled_omega_arr`             | `NLayers`                          |
+    | `tau_arr`                      | `NLayers`                          |
+    | `scaled_tau_arr_with_0`        | `NLayers + 1`                      |
+    | `mu_arr_pos`                   | `NQuad/2`                          |
+    | `mu_arr`                       | `NQuad`                            |
+    | `M_inv`                        | `NQuad/2`                          |
+    | `W`                            | `NQuad/2`                          |
+    | `N`                            | scalar                             |
+    | `NQuad`                        | scalar                             |
+    | `NLeg`                         | scalar                             |
+    | `NFourier`                     | scalar                             |
+    | `NLayers`                      | scalar                             |
+    | `NBDRF`                        | scalar                             |
+    | `is_atmos_multilayered`        | boolean                            |
+    | `weighted_scaled_Leg_coeffs`   | `NLayers x NLeg`                   |
+    | `BDRF_Fourier_modes`           | `NBDRF`                            |
+    | `mu0`                          | scalar                             |
+    | `I0`                           | scalar                             |
+    | `I0_orig`                      | scalar                             |
+    | `phi0`                         | scalar                             |
+    | `there_is_beam_source`         | scalar                             |
+    | `b_pos`                        | `NQuad/2 x NFourier` or scalar     |
+    | `b_neg`                        | `NQuad/2 x NFourier` or scalar     |
+    | `b_pos_is_scalar`              | boolean                            |
+    | `b_neg_is_scalar`              | boolean                            |              
+    | `s_poly_coeffs`                | `NLayers x Nscoeffs` or `Nscoeffs` |
+    | `Nscoeffs`                     | scalar                             |
+    | `there_is_iso_source`          | boolean                            |
+    | `scale_tau`                    | `NLayers`                          |
+    | `only_flux`                    | boolean                            |
+    | `NLayers_to_use_sparse`        | scalar                             |
+    | `_autograd_compatible`         | boolean                            |
+    
+    Notable internal variables of _assemble_intensity_and_fluxes
+    |     Variable      |             Type / Shape               |
+    | ----------------- | -------------------------------------- |
+    | `G_collect`       | `NFourier x NLayers x NQuad x NQuad`   |
+    | `G_collect_0`     | `NLayers x NQuad x NQuad`              |
+    | `G_inv_collect_0` | `NLayers x NQuad x NQuad` or `None`    |
+    | `K_collect`       | `NFourier x NLayers x NQuad`           |
+    | `K_collect_0`     | `NLayers x NQuad`                      |
+    | `B_collect`       | `NFourier x NLayers x NQuad` or `None` |
+    | `B_collect_0`     | `NLayers x NQuad` or `None`            |
+    | `B_pos`           | `NLayers x NQuad/2` or `None`          |
+    | `B_neg`           | `NLayers x NQuad/2` or `None`          |
+    | `GC_collect`      | `NFourier x NLayers x NQuad x NQuad`   |
+    | `GC_collect_0`    | `NLayers x NQuad x NQuad`              |
+    | `GC_pos`          | `NLayers x NQuad/2 x NQuad`            |
+    | `GC_neg`          | `NLayers x NQuad/2 x NQuad`            |
     """
-    if _is_compatible_with_autograd:
+    if _autograd_compatible:
         import autograd.numpy as np
     else:
         import numpy as np
     
-    ###################################### Assemble uncorrected solution functions #############################################
+    ################################## Assemble uncorrected intensity and flux functions #######################################
     
     # Compute all the necessary quantities
     # --------------------------------------------------------------------------------------------------------------------------
-    outputs = _diagonalize(
+    outputs = _solve_for_gen_and_part_sols(
         NFourier,
         scaled_omega_arr,
         mu_arr_pos, mu_arr,
@@ -73,7 +129,7 @@ def _assemble_solution_functions(
             G_inv_collect_0 = None
         B_collect = None
             
-    GC_collect = _solve_for_coefs(
+    GC_collect = _solve_for_coeffs(
         NFourier,
         G_collect,
         K_collect,
@@ -81,10 +137,10 @@ def _assemble_solution_functions(
         G_inv_collect_0,
         tau_arr,
         scaled_tau_arr_with_0,
-        mu_arr_pos, mu_arr_pos * W, mu_arr,
+        mu_arr, mu_arr_pos, mu_arr_pos * W,
         N, NQuad,
         NLayers, NBDRF,
-        atmos_is_multilayered,
+        is_atmos_multilayered,
         BDRF_Fourier_modes,
         mu0, I0,
         there_is_beam_source,
@@ -93,7 +149,7 @@ def _assemble_solution_functions(
         s_poly_coeffs,
         Nscoeffs,
         there_is_iso_source,
-        use_sparse_NLayers,
+        NLayers_to_use_sparse,
     )
     
     G_collect_0 = G_collect[0, :, :, :]
@@ -103,7 +159,7 @@ def _assemble_solution_functions(
         
     if not only_flux:
  
-        # Construct the intensity function
+        # Construct the intensity function (refer to Section 3.7 of the Comprehensive Documentation)
         # --------------------------------------------------------------------------------------------------------------------------
         def u(tau, phi, return_Fourier_error=False):
             tau = np.atleast_1d(tau)
@@ -150,10 +206,10 @@ def _assemble_solution_functions(
                     K_collect_0,
                     G_inv_collect_0,
                     mu_arr,
-                    _is_compatible_with_autograd,
+                    _autograd_compatible,
                 )
                 
-                if _is_compatible_with_autograd:
+                if _autograd_compatible:
                     um = um + np.concatenate((_mathscr_v_contribution.T, np.zeros((NFourier - 1, len(tau), NQuad))))
                 else:
                     um[0, :, :] += _mathscr_v_contribution.T
@@ -245,7 +301,7 @@ def _assemble_solution_functions(
                 K_collect_0,
                 G_inv_collect_0,
                 mu_arr,
-                _is_compatible_with_autograd,
+                _autograd_compatible,
             )
             u0 += _mathscr_v_contribution
         
@@ -255,7 +311,7 @@ def _assemble_solution_functions(
             return I0_orig * np.squeeze(u0)
     # --------------------------------------------------------------------------------------------------------------------------
     
-    # Construct the flux functions
+    # Construct the flux functions (refer to Section 3.8 of the Comprehensive Documentation)
     # --------------------------------------------------------------------------------------------------------------------------
     GC_pos = GC_collect_0[:, :N, :]
     GC_neg = GC_collect_0[:, N:, :]
@@ -291,7 +347,7 @@ def _assemble_solution_functions(
                 K_collect_0,
                 G_inv_collect_0,
                 mu_arr,
-                _is_compatible_with_autograd,
+                _autograd_compatible,
             )[:N, :]
         else:
             _mathscr_v_contribution = 0
@@ -345,7 +401,7 @@ def _assemble_solution_functions(
                 K_collect_0,
                 G_inv_collect_0,
                 mu_arr,
-                _is_compatible_with_autograd,
+                _autograd_compatible,
             )[N:, :]
         else:
             _mathscr_v_contribution = 0

PythonicDISORT/{_solve_for_coefs.py → _solve_for_coeffs.py}

@@ -6,38 +6,79 @@ import scipy as sc
 from math import pi
 
 
-def _solve_for_coefs(
-    NFourier,
-    G_collect,
-    K_collect,
-    B_collect,
-    G_inv_collect_0,
-    tau_arr,
-    scaled_tau_arr_with_0,
-    mu_arr_pos, mu_arr_pos_times_W, mu_arr,
-    N, NQuad,
-    NLayers, NBDRF,
-    atmos_is_multilayered,
-    BDRF_Fourier_modes,
-    mu0, I0,
-    there_is_beam_source,
-    b_pos, b_neg,
-    b_pos_is_scalar, b_neg_is_scalar,
-    s_poly_coeffs,
-    Nscoeffs,
-    there_is_iso_source,
-    use_sparse_NLayers,
+def _solve_for_coeffs(
+    NFourier,                               # Number of intensity Fourier modes
+    G_collect,                              # Eigenvector matrices
+    K_collect,                              # Eigenvalues
+    B_collect,                              # Coefficients vectors for particular solutions
+    G_inv_collect_0,                        # Inverse of eigenvector matrix for the 0th Fourier mode
+    tau_arr,                                # Lower boundary of layers
+    scaled_tau_arr_with_0,                  # Delta-scaled lower boundary of layers with 0 inserted in front
+    mu_arr, mu_arr_pos, mu_arr_pos_times_W, # Quadrature nodes for 1) both 2) upper hemispheres; 3) upper hemisphere quadrature nodes times weights
+    N, NQuad,                               # Number of 1) upper 2) both hemispheres quadrature nodes
+    NLayers, NBDRF,                         # Number of 1) intensity Fourier modes; 2) layers
+    is_atmos_multilayered,                  # Is the atmosphere multilayered?
+    BDRF_Fourier_modes,                     # BDRF Fourier modes
+    mu0, I0,                                # Properties of direct beam
+    there_is_beam_source,                   # Is there a beam source?
+    b_pos, b_neg,                           # Dirichlet BCs
+    b_pos_is_scalar, b_neg_is_scalar,       # Is each Dirichlet BCs scalar (isotropic)?
+    s_poly_coeffs,                          # Polynomial coefficients of isotropic source           
+    Nscoeffs,                               # Number of isotropic source polynomial coefficients
+    there_is_iso_source,                    # Is there an isotropic source?
+    NLayers_to_use_sparse,                  # Number of layers above or equal which to use sparse matrices
 ):
-    """This function is wrapped by the `_assemble_solution_functions` function.
+    """
+    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.
+    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.
     The labels in this file reference labels in the Jupyter Notebook, especially sections 3 and 4.
 
+    Arguments of _solve_for_coeffs
+    |            Variable            |              Type / Shape              |
+    | ------------------------------ | -------------------------------------- |
+    | `NFourier`                     | scalar                                 |
+    | `G_collect`                    | `NFourier x NLayers x NQuad x NQuad`   |
+    | `K_collect`                    | `NFourier x NLayers x NQuad`           |
+    | `B_collect`                    | `NFourier x NLayers x NQuad` or `None` |
+    | `G_inv_collect_0`              | `NLayers x NQuad x NQuad` or `None`    |
+    | `tau_arr`                      | `NLayers`                              |
+    | `scaled_tau_arr_with_0`        | `NLayers + 1`                          |
+    | `mu_arr`                       | `NQuad`                                |
+    | `mu_arr_pos`                   | `NQuad/2`                              |
+    | `mu_arr_pos_times_W`           | `NQuad/2`                              |
+    | `N`                            | scalar                                 |
+    | `NQuad`                        | scalar                                 |
+    | `NLayers`                      | scalar                                 |
+    | `NBDRF`                        | scalar                                 |
+    | `is_atmos_multilayered`        | boolean                                |
+    | `BDRF_Fourier_modes`           | `NBDRF`                                |
+    | `mu0`                          | scalar                                 |
+    | `I0`                           | scalar                                 |
+    | `there_is_beam_source`         | boolean                                |
+    | `b_pos`                        | `NQuad/2 x NFourier` or scalar         |
+    | `b_neg`                        | `NQuad/2 x NFourier` or scalar         |
+    | `b_pos_is_scalar`              | boolean                                |
+    | `b_neg_is_scalar`              | boolean                                |              
+    | `s_poly_coeffs`                | `NLayers x Nscoeffs` or `Nscoeffs`     |
+    | `Nscoeffs`                     | scalar                                 |
+    | `there_is_iso_source`          | boolean                                |
+    | `NLayers_to_use_sparse`        | scalar                                 |
+    
+    Notable internal variables of _solve_for_coeffs
+    |   Variable   |                 Shape                |
+    | ------------ | ------------------------------------ |
+    | `GC_collect` | `NFourier x NLayers x NQuad x NQuad` |
+
     """
     ################################## Solve for coefficients of homogeneous solution ##########################################
+    ############# Refer to Section 3.6.2 (single-layer) and 4 (multi-layer) of the Comprehensive Documentation #################
     
     GC_collect = np.empty((NFourier, NLayers, NQuad, NQuad))
-    use_sparse_framework = NLayers >= use_sparse_NLayers
+    use_sparse_framework = NLayers >= NLayers_to_use_sparse
     
     # The following loops can easily be parallelized, but the speed-up is unlikely to be worth the overhead
     for m in range(NFourier):
@@ -50,6 +91,7 @@ def _solve_for_coefs(
             B_collect_m = B_collect[m, :, :]
             
         # Generate mathscr_D and mathscr_X (BDRF terms)
+        # Just for this part, refer to Section 3.4.2 of the Comprehensive Documentation 
         # --------------------------------------------------------------------------------------------------------------------------
         if BDRF_bool:
             mathscr_D_neg = (1 + m_equals_0 * 1) * BDRF_Fourier_modes[m](mu_arr_pos, -mu_arr_pos)
@@ -92,7 +134,7 @@ def _solve_for_coefs(
                     ).ravel()
         
             _mathscr_v_contribution_middle = np.array([])
-            if atmos_is_multilayered:
+            if is_atmos_multilayered:
                 indices = np.arange(NLayers - 1)
                 _mathscr_v_contribution_middle = (
                     _mathscr_v(
@@ -147,7 +189,7 @@ def _solve_for_coefs(
                 BDRF_RHS_contribution = 0
             
             RHS_middle = np.array([])
-            if atmos_is_multilayered:
+            if is_atmos_multilayered:
                 RHS_middle = np.array(
                     [
                         (B_collect_m[l + 1, :] - B_collect_m[l, :])
@@ -173,7 +215,7 @@ def _solve_for_coefs(
             RHS = np.concatenate([b_neg_m, RHS_middle, b_pos_m]) + _mathscr_v_contribution
         # --------------------------------------------------------------------------------------------------------------------------
         
-        # Assemble LHS
+        # Assemble LHS (much of this code is replicated in Section 4 of the Comprehensive Documentation)
         dim = NLayers * NQuad
 
         G_0_nn = G_collect_m[0, N:, :N]
@@ -193,8 +235,8 @@ def _solve_for_coefs(
             BDRF_LHS_contribution_pos = 0
         
         if use_sparse_framework:
-            N_squared = N**2
-            density = 4 * N_squared * (2 * NLayers - 1)
+            N_sq = N**2
+            density = 4 * N_sq * (2 * NLayers - 1)
             row_indices = np.empty(density)
             col_indices = np.empty(density)
             block_rows = np.empty(density)
@@ -203,35 +245,35 @@ def _solve_for_coefs(
             block_row_ind = np.arange(block_row_len).reshape(NQuad, 2 * NQuad)
             block_row_sort = np.empty(block_row_len, dtype='int')
             
-            block_row_sort[:N_squared] = block_row_ind[:N, :N].ravel()
-            block_row_sort[N_squared : 2 * N_squared] = block_row_ind[N : 2 * N, :N].ravel()
-            block_row_sort[2 * N_squared : 4 * N_squared] = block_row_ind[:2 * N, N : 2 * N].ravel()
-            block_row_sort[4 * N_squared : 6 * N_squared] = block_row_ind[:2 * N, 2 * N : 3 * N].ravel()
-            block_row_sort[6 * N_squared : 7 * N_squared] = block_row_ind[:N, 3 * N : 4 * N].ravel()
-            block_row_sort[7 * N_squared : 8 * N_squared] = block_row_ind[N : 2 * N, 3 * N : 4 * N].ravel()
+            block_row_sort[:N_sq] = block_row_ind[:N, :N].ravel()
+            block_row_sort[N_sq : 2 * N_sq] = block_row_ind[N : 2 * N, :N].ravel()
+            block_row_sort[2 * N_sq : 4 * N_sq] = block_row_ind[:2 * N, N : 2 * N].ravel()
+            block_row_sort[4 * N_sq : 6 * N_sq] = block_row_ind[:2 * N, 2 * N : 3 * N].ravel()
+            block_row_sort[6 * N_sq : 7 * N_sq] = block_row_ind[:N, 3 * N : 4 * N].ravel()
+            block_row_sort[7 * N_sq : 8 * N_sq] = block_row_ind[N : 2 * N, 3 * N : 4 * N].ravel()
 
             # BCs for the entire atmosphere
-            row_indices[:N_squared], col_indices[:N_squared] = _nd_slice_to_indexes(np.s_[:N, :N])
-            block_rows[:N_squared] = G_0_nn.ravel()
+            row_indices[:N_sq], col_indices[:N_sq] = _nd_slice_to_indexes(np.s_[:N, :N])
+            block_rows[:N_sq] = G_0_nn.ravel()
             (
-                row_indices[N_squared : 2 * N_squared],
-                col_indices[N_squared : 2 * N_squared],
+                row_indices[N_sq : 2 * N_sq],
+                col_indices[N_sq : 2 * N_sq],
             ) = _nd_slice_to_indexes(np.s_[:N, N:NQuad])
-            block_rows[N_squared : 2 * N_squared] = (
+            block_rows[N_sq : 2 * N_sq] = (
                 G_0_np * np.exp(K_collect_m[0, :N] * scaled_tau_arr_with_0[1])[None, :]
             ).ravel()
             (
-                row_indices[2 * N_squared : 3 * N_squared],
-                col_indices[2 * N_squared : 3 * N_squared],
+                row_indices[2 * N_sq : 3 * N_sq],
+                col_indices[2 * N_sq : 3 * N_sq],
             ) = _nd_slice_to_indexes(np.s_[dim - N : dim, dim - NQuad : dim - N])
-            block_rows[2 * N_squared : 3 * N_squared] = (
+            block_rows[2 * N_sq : 3 * N_sq] = (
                 (G_L_pn - BDRF_LHS_contribution_neg) * E_Lm1L[None, :]
             ).ravel()
             (
-                row_indices[3 * N_squared : 4 * N_squared],
-                col_indices[3 * N_squared : 4 * N_squared],
+                row_indices[3 * N_sq : 4 * N_sq],
+                col_indices[3 * N_sq : 4 * N_sq],
             ) = _nd_slice_to_indexes(np.s_[dim - N : dim, dim - N : dim])
-            block_rows[3 * N_squared : 4 * N_squared] = (
+            block_rows[3 * N_sq : 4 * N_sq] = (
                 G_L_pp - BDRF_LHS_contribution_pos
             ).ravel()
 
@@ -256,19 +298,19 @@ def _solve_for_coefs(
                     np.s_[N + l * NQuad : N + (l + 1) * NQuad, l * NQuad : (l + 2) * NQuad]
                 )
                 row_indices[
-                    4 * N_squared + l * block_row_len : 4 * N_squared + (l + 1) * block_row_len
+                    4 * N_sq + l * block_row_len : 4 * N_sq + (l + 1) * block_row_len
                 ] = row_indices_l[block_row_sort]
                 col_indices[
-                    4 * N_squared + l * block_row_len : 4 * N_squared + (l + 1) * block_row_len
+                    4 * N_sq + l * block_row_len : 4 * N_sq + (l + 1) * block_row_len
                 ] = col_indices_l[block_row_sort]
                 
-                starting_ind = 4 * N_squared + l * block_row_len
-                block_rows[starting_ind : N_squared + starting_ind] = (G_l_pn * E_lm1l[None, :]).ravel()
-                block_rows[starting_ind + N_squared : 2 * N_squared + starting_ind] = (G_l_nn * E_lm1l[None, :]).ravel()
-                block_rows[starting_ind + 2 * N_squared : 4 * N_squared + starting_ind] = G_l_ap.ravel()
-                block_rows[starting_ind + 4 * N_squared : 6 * N_squared + starting_ind] = -G_lp1_an.ravel()
-                block_rows[starting_ind + 6 * N_squared : 7 * N_squared + starting_ind] = -(G_lp1_pp * E_llp1[None, :]).ravel()
-                block_rows[starting_ind + 7 * N_squared : 8 * N_squared + starting_ind] = -(G_lp1_np * E_llp1[None, :]).ravel()
+                start_ind = 4 * N_sq + l * block_row_len
+                block_rows[start_ind : N_sq + start_ind] = (G_l_pn * E_lm1l[None, :]).ravel()
+                block_rows[start_ind + N_sq : 2 * N_sq + start_ind] = (G_l_nn * E_lm1l[None, :]).ravel()
+                block_rows[start_ind + 2 * N_sq : 4 * N_sq + start_ind] = G_l_ap.ravel()
+                block_rows[start_ind + 4 * N_sq : 6 * N_sq + start_ind] = -G_lp1_an.ravel()
+                block_rows[start_ind + 6 * N_sq : 7 * N_sq + start_ind] = -(G_lp1_pp * E_llp1[None, :]).ravel()
+                block_rows[start_ind + 7 * N_sq : 8 * N_sq + start_ind] = -(G_lp1_np * E_llp1[None, :]).ravel()
     
             LHS = sc.sparse.coo_matrix(
                 (
@@ -284,7 +326,7 @@ def _solve_for_coefs(
             # Solve the system
             C_m = sc.sparse.linalg.spsolve(LHS.tocsr(), RHS).reshape(NLayers, NQuad)
         
-        else:
+        else:            
             LHS = np.zeros((dim, dim))
 
             # BCs for the entire atmosphere
@@ -313,14 +355,14 @@ def _solve_for_coefs(
                 E_lm1l = np.exp(K_l_pos * (scaled_tau_arr_lm1 - scaled_tau_arr_l))
                 E_llp1 = np.exp(K_lp1_pos * (scaled_tau_arr_l - scaled_tau_arr_lp1))
                 
-                starting_row = N + l * NQuad
-                starting_col = l * NQuad
-                LHS[starting_row : N + starting_row, starting_col : N + starting_col] = G_l_pn * E_lm1l[None, :]
-                LHS[N + starting_row : 2 * N + starting_row, starting_col : N + starting_col] = G_l_nn * E_lm1l[None, :]
-                LHS[starting_row : 2 * N + starting_row, N + starting_col : 2 * N + starting_col] = G_l_ap
-                LHS[starting_row : 2 * N + starting_row, 2 * N + starting_col : 3 * N + starting_col] = -G_lp1_an
-                LHS[starting_row : N + starting_row, 3 * N + starting_col : 4 * N + starting_col] = -G_lp1_pp * E_llp1[None, :]
-                LHS[N + starting_row : 2 * N + starting_row, 3 * N + starting_col : 4 * N + starting_col] = -G_lp1_np * E_llp1[None, :]
+                start_row = N + l * NQuad
+                start_col = l * NQuad
+                LHS[start_row : N + start_row, start_col : N + start_col] = G_l_pn * E_lm1l[None, :]
+                LHS[N + start_row : 2 * N + start_row, start_col : N + start_col] = G_l_nn * E_lm1l[None, :]
+                LHS[start_row : 2 * N + start_row, N + start_col : 2 * N + start_col] = G_l_ap
+                LHS[start_row : 2 * N + start_row, 2 * N + start_col : 3 * N + start_col] = -G_lp1_an
+                LHS[start_row : N + start_row, 3 * N + start_col : 4 * N + start_col] = -G_lp1_pp * E_llp1[None, :]
+                LHS[N + start_row : 2 * N + start_row, 3 * N + start_col : 4 * N + start_col] = -G_lp1_np * E_llp1[None, :]
             
             # Solve the system
             C_m = np.linalg.solve(LHS, RHS).reshape(NLayers, NQuad)

PythonicDISORT/{_diagonalize.py → _solve_for_gen_and_part_sols.py}

@@ -3,26 +3,63 @@ import scipy as sc
 from math import pi
 
 
-def _diagonalize(
-    NFourier,
-    scaled_omega_arr,
-    mu_arr_pos, mu_arr,
-    M_inv, W,
-    N, NQuad, NLeg,
-    NLayers,
-    weighted_scaled_Leg_coeffs,
-    mu0, I0,
-    there_is_beam_source,
-    Nscoeffs,
-    there_is_iso_source,
+def _solve_for_gen_and_part_sols(
+    NFourier,                   # Number of intensity Fourier modes
+    scaled_omega_arr,           # Delta-scaled single-scattering albedos
+    mu_arr_pos, mu_arr,         # Quadrature nodes for 1) upper 2) both hemispheres
+    M_inv, W,                   # 1) 1 / mu; 2) quadrature weights for each hemisphere
+    N, NQuad, NLeg,             # Number of 1) upper 2) both hemispheres quadrature nodes; 3) phase function Legendre coefficients 
+    NLayers,                    # Number of layers
+    weighted_scaled_Leg_coeffs, # Weighted and delta-scaled Legendre coefficients
+    mu0, I0,                    # Properties of direct beam
+    there_is_beam_source,       # Is there a beam source?
+    Nscoeffs,                   # Number of isotropic source polynomial coefficients
+    there_is_iso_source,        # Is there an isotropic source?
 ):
-    """This function is wrapped by the `_assemble_solution_functions` function.
+    """
+    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.
+    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.
     The labels in this file reference labels in the Jupyter Notebook, especially sections 3 and 4.
 
+    Arguments of _solve_for_gen_and_part_sols
+    |            Variable            |            Type / Shape            |
+    | ------------------------------ | ---------------------------------- |
+    | `NFourier`                     | scalar                             |
+    | `scaled_omega_arr`             | `NLayers`                          |
+    | `mu_arr_pos`                   | `NQuad/2`                          |
+    | `mu_arr`                       | `NQuad`                            |
+    | `M_inv`                        | `NQuad/2`                          |
+    | `W`                            | `NQuad/2`                          |
+    | `N`                            | scalar                             |
+    | `NQuad`                        | scalar                             |
+    | `NLeg`                         | scalar                             |
+    | `NLayers`                      | scalar                             |
+    | `weighted_scaled_Leg_coeffs`   | `NLayers x NLeg`                   |
+    | `mu0`                          | scalar                             |
+    | `I0`                           | scalar                             |
+    | `there_is_beam_source`         | boolean                            |
+    | `Nscoeffs`                     | scalar                             |
+    | `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`    |
+
     """
     ############################### Assemble system and diagonalize coefficient matrix #########################################
+    ########################### Refer to Section 3.4.2 of the Comprehensive Documentation  #####################################
     
     # Initialization
     # --------------------------------------------------------------------------------------------------------------------------
@@ -116,7 +153,7 @@ def _diagonalize(
                 
     if len(no_shortcut_indices) > 0:
     
-        # Diagonalization of coefficient matrix
+        # Diagonalization of coefficient matrix (refer to Section 3.4.2 of the Comprehensive Documentation)
         # --------------------------------------------------------------------------------------------------------------------------
         alpha_arr = alpha_arr[: len(no_shortcut_indices), :, :]
         beta_arr = beta_arr[: len(no_shortcut_indices), :, :]
@@ -153,7 +190,7 @@ def _diagonalize(
             G_inv_collect_0[no_shortcut_indices_0, :, :] = G_inv_arr[: len(no_shortcut_indices_0), :, :]
         # --------------------------------------------------------------------------------------------------------------------------
 
-        # Particular solution for the sunbeam source
+        # Particular solution for the sunbeam 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

@@ -1,5 +1,5 @@
 from PythonicDISORT import subroutines
-from PythonicDISORT._assemble_solution_functions import _assemble_solution_functions
+from PythonicDISORT._assemble_intensity_and_fluxes import _assemble_intensity_and_fluxes
 
 import numpy as np
 import scipy as sc
@@ -22,8 +22,8 @@ def pydisort(
     NT_cor=False,
     BDRF_Fourier_modes=[],
     s_poly_coeffs=np.array([[]]),
-    use_sparse_NLayers=13,
-    _is_compatible_with_autograd = False,
+    NLayers_to_use_sparse=13,
+    _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
@@ -37,14 +37,15 @@ def pydisort(
 
     Parameters
     ----------
-    tau_arr : array
+    tau_arr : array or float
         Optical depth of the lower boundary of each atmospheric layer.
-    omega_arr : array
+    omega_arr : array or float
         Single-scattering albedo of each atmospheric layer.
     NQuad : int
         Number of `mu` quadrature nodes.
-    Leg_coeffs_all : ndarray
-        All available unweighted phase function Legendre coefficients. 
+    Leg_coeffs_all : 2darray
+        All available unweighted phase function Legendre coefficients.
+        Each row of coefficients pertain to an atmospheric layer.
         Each coefficient should be between 0 and 1 inclusive.
     mu0 : float
         Cosine of polar angle of the incident beam.
@@ -53,9 +54,10 @@ def pydisort(
     phi0 : float
         Azimuthal angle of the incident beam.
     NLeg : optional, int
-        Number of phase function Legendre coefficients.
+        Number of phase function Legendre coefficients
+        to use in the pre-correction solver.
     NFourier : optional, int
-        Number of Fourier modes to use in the intensity function.
+        Number of Fourier modes to use to construct the intensity function.
     b_pos : optional, 2darray or float
         Dirichlet boundary condition for the upward direction.
     b_neg : optional, 2darray or float
@@ -72,33 +74,80 @@ def pydisort(
     s_poly_coeffs : optional, array
         Polynomial coefficients of isotropic internal sources.
         Arrange coefficients from lowest order term to highest.
-    use_sparse_NLayers : optional, int
+    NLayers_to_use_sparse : optional, int
         At or above how many atmospheric layers should SciPy's sparse matrix framework be used?
 
     Returns
     -------
-    array
+    mu_arr : array
         All `mu` (cosine of polar angle) quadrature nodes.
-    function
+    Fp(tau) : function, function
         (Energetic) Flux function with argument `tau` (type: array) for positive (upward) `mu` values.
         Returns the diffuse flux magnitudes (type: array).
-    function
+    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)).
-    function
+    u0(tau) : function
         Zeroth Fourier mode of the intensity with argument `tau` (type: array).
         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).
-    function, optional
+    u(tau, phi) : function, optional
         Intensity function with arguments `(tau, phi, return_Fourier_error=False)` of types `(array, array, bool)`.
         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.
 
     """
-    if _is_compatible_with_autograd:
+    
+    """
+    Arguments of pydisort
+    |         Variable        |            Type / Shape            |
+    | ----------------------- | ---------------------------------- |
+    | `tau_arr`               | `NLayers`                          |
+    | `omega_arr`             | `NLayers`                          |
+    | `NQuad`                 | scalar                             |
+    | `Leg_coeffs_all`        | `NLeg_all`                         |
+    | `mu0`                   | scalar                             |
+    | `I0`                    | scalar                             |
+    | `phi0`                  | scalar                             |
+    | `NLeg`                  | scalar                             |
+    | `NFourier`              | scalar                             |
+    | `b_pos`                 | `NQuad/2 x NFourier` or scalar     |
+    | `b_neg`                 | `NQuad/2 x NFourier` or scalar     |              
+    | `only_flux`             | boolean                            |
+    | `f_arr`                 | `NLayers` or scalar                |
+    | `NT_cor`                | boolean                            |
+    | `BDRF_Fourier_modes`    | `NBDRF`                            |
+    | `s_poly_coeffs`         | `NLayers x Nscoeffs` or `Nscoeffs` |
+    | `NLayers_to_use_sparse` | scalar                             |
+    
+    Notable internal variables of pydisort
+    |          Variable            |     Type / Shape     |
+    | ---------------------------- | -------------------- |
+    | `I0_orig`                    | scalar               |
+    | `thickness_arr`              | `NLayers`            |
+    | `weighted_Leg_coeffs_all`    | `NLayers x NLeg_all` |
+    | `Leg_coeffs`                 | `NLayers x NLeg`     |
+    | `weighted_scaled_Leg_coeffs` | `NLayers x NLeg`     |
+    | `mu_arr_pos`                 | `NQuad/2`            |
+    | `W`                          | `NQuad/2`            |
+    | `mu_arr`                     | `NQuad`              |
+    | `M_inv`                      | `NQuad/2`            |
+    | `scale_tau`                  | `NLayers`            |
+    | `scaled_tau_arr_with_0`      | `NLayers + 1`        |
+    | `scaled_omega_arr`           | `NLayers`            |
+    | `sum1`                       | scalar               |
+    | `omega_avg`                  | scalar               |
+    | `sum2`                       | scalar               |
+    | `f_avg`                      | scalar               |
+    | `Leg_coeffs_residue`         | `NLayers x NLeg_all` |
+    | `Leg_coeffs_residue_avg`     | `NLeg_all`           |
+    | `scaled_mu0`                 | scalar               |
+    """
+
+    if _autograd_compatible:
         import autograd.numpy as np
     else:
         import numpy as np
@@ -134,10 +183,10 @@ def pydisort(
     N = NQuad // 2
     there_is_beam_source = I0 > 0
     there_is_iso_source = Nscoeffs > 0
-    atmos_is_multilayered = NLayers > 1
+    is_atmos_multilayered = NLayers > 1
     # --------------------------------------------------------------------------------------------------------------------------
 
-    # Input checks
+    # Input checks (refer to Section 1 of the Comprehensive Documentation)
     # --------------------------------------------------------------------------------------------------------------------------
     # Optical depths and thickness must be positive
     assert np.all(tau_arr > 0)
@@ -183,7 +232,7 @@ def pydisort(
     # The fractional scattering must be between 0 and 1
     assert np.all(0 <= f_arr) and np.all(f_arr <= 1)
     # The minimum threshold is the minimum numbers of layers: 1
-    assert(use_sparse_NLayers >= 1)
+    assert(NLayers_to_use_sparse >= 1)
     # --------------------------------------------------------------------------------------------------------------------------
     
     # Some more setup
@@ -191,6 +240,7 @@ def pydisort(
     NBDRF = len(BDRF_Fourier_modes)
     weighted_Leg_coeffs_all = (2 * np.arange(NLeg_all) + 1) * Leg_coeffs_all
     Leg_coeffs = Leg_coeffs_all[:, :NLeg]
+    # The following is explained in Section 1.4 of the Comprehensive Documentation
     if (b_pos_is_scalar and b_pos == 0) and (b_neg_is_scalar and b_neg == 0) and not there_is_iso_source and there_is_beam_source:
         I0_orig = I0
         I0 = 1
@@ -198,7 +248,7 @@ def pydisort(
         I0_orig = 1
     # --------------------------------------------------------------------------------------------------------------------------
     
-    # Generation of Double Gauss-Legendre quadrature weights and points
+    # Generation of "double-Gauss" quadrature weights and points (refer to Section 3.4 of the Comprehensive Documentation)
     # --------------------------------------------------------------------------------------------------------------------------
     # For positive mu values (the weights are identical for both domains)
     mu_arr_pos, W = subroutines.Gauss_Legendre_quad(N)  # mu_arr_neg = -mu_arr_pos
@@ -210,6 +260,7 @@ def pydisort(
     # --------------------------------------------------------------------------------------------------------------------------
 
     # Delta-M scaling; there is no scaling if f = 0
+    # Refer to Sections 1.3.1 and 3.3 of the Comprehensive Documentation
     # --------------------------------------------------------------------------------------------------------------------------
     if np.any(f_arr > 0):
         scale_tau = 1 - omega_arr * f_arr
@@ -231,18 +282,19 @@ def pydisort(
     
     if NT_cor and not only_flux and there_is_beam_source and np.any(f_arr > 0) and NLeg < NLeg_all:
         
-        ############################### Perform NT corrections on the intensity but not the flux ###################################
+        ############################### Perform NT corrections on the intensity but not the flux ###############################
+        ############################### Refer to Section 3.7.2 of the Comprehensive Documentation ##############################
         
         # Delta-M scaled solution; no further corrections to the flux
-        flux_up, flux_down, u0, u_star = _assemble_solution_functions(
+        flux_up, flux_down, u0, u_star = _assemble_intensity_and_fluxes(
             scaled_omega_arr,
             tau_arr,
             scaled_tau_arr_with_0,
             mu_arr_pos, mu_arr,
             M_inv, W,
-            N, NQuad, NLeg, NFourier,
-            NLayers, NBDRF,
-            atmos_is_multilayered,
+            N, NQuad, NLeg,
+            NFourier, NLayers, NBDRF,
+            is_atmos_multilayered,
             weighted_scaled_Leg_coeffs,
             BDRF_Fourier_modes,
             mu0, I0, I0_orig, phi0,
@@ -254,8 +306,8 @@ def pydisort(
             there_is_iso_source,
             scale_tau,
             only_flux,
-            use_sparse_NLayers,
-            _is_compatible_with_autograd,
+            NLayers_to_use_sparse,
+            _autograd_compatible,
         )
         
         # TMS correction
@@ -324,16 +376,25 @@ def pydisort(
             # TODO: Despite being vectorized and despite many attempts at further optimization,
             # this block of code is slower than desired and is the bottleneck for the TMS correction.
             # --------------------------------------------------------------------------------------------------------------------------
-            if atmos_is_multilayered:
+            solution = (
+                mathscr_B[:, l, :] * np.vstack((TMS_correction_pos, TMS_correction_neg))[:, :, None]
+            )
+            
+            if is_atmos_multilayered:
                 layers_arr = np.arange(NLayers)[:, None]
                 pos_contribution_mask = (l[None, :] < layers_arr).ravel()
                 neg_contribution_mask = (l[None, :] > layers_arr).ravel()
+                any_pos_contribution = np.any(pos_contribution_mask)
+                any_neg_contribution = np.any(neg_contribution_mask)
+            
+                if not (any_pos_contribution or any_neg_contribution):
+                    return solution
 
                 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 np.any(pos_contribution_mask):
+                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[
@@ -365,7 +426,7 @@ def pydisort(
                         contribution_from_each_other_layer_pos.reshape((N, NLayers, Ntau, Nphi)),
                         axis=1,
                     )
-                if np.any(neg_contribution_mask):
+                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[
@@ -398,30 +459,26 @@ def pydisort(
                         axis=1,
                     )
 
-                if _is_compatible_with_autograd:
-                    solution = (
-                        mathscr_B[:, l, :]
-                        * np.vstack((TMS_correction_pos, TMS_correction_neg))[:, :, None]
-                    )
-                    solution[:N, :, :] += contribution_from_other_layers_pos
-                    solution[N:, :, :] += contribution_from_other_layers_neg
-                    return solution
-
-                else:
-                    return mathscr_B[:, l, :] * np.vstack((TMS_correction_pos, TMS_correction_neg))[
-                        :, :, None
-                    ] + np.concatenate(
+                if _autograd_compatible:
+                    if not any_pos_contribution:
+                        TMS_correction_pos = np.zeros((N, Ntau, Nphi))
+                    if not any_neg_contribution:
+                        TMS_correction_neg = np.zeros((N, Ntau, Nphi))
+                    return solution + np.concatenate(
                         (contribution_from_other_layers_pos, contribution_from_other_layers_neg),
                         axis=0,
                     )
+                else:
+                    if any_pos_contribution:
+                        solution[:N, :, :] += contribution_from_other_layers_pos
+                    if any_neg_contribution:
+                        solution[N:, :, :] += contribution_from_other_layers_neg
+                    return solution
                     
             # --------------------------------------------------------------------------------------------------------------------------
             
             else:
-                return (
-                    mathscr_B[:, l, :]
-                    * np.vstack((TMS_correction_pos, TMS_correction_neg))[:, :, None]
-                )
+                return solution
         # --------------------------------------------------------------------------------------------------------------------------
 
         # IMS correction
@@ -466,7 +523,7 @@ def pydisort(
             phi = np.atleast_1d(phi)
             NT_corrections = TMS_correction(tau, phi)
 
-            if _is_compatible_with_autograd:
+            if _autograd_compatible:
                 NT_corrections = NT_corrections + np.concatenate(
                     [np.zeros((N, len(tau), len(phi))), IMS_correction(tau, phi)], axis=0
                 )
@@ -488,15 +545,15 @@ def pydisort(
     else:
         if only_flux:
             NFourier = 1 # We only need to solve for the 0th Fourier mode to compute the flux
-        return (mu_arr,) + _assemble_solution_functions(
+        return (mu_arr,) + _assemble_intensity_and_fluxes(
             scaled_omega_arr,
             tau_arr,
             scaled_tau_arr_with_0,
             mu_arr_pos, mu_arr,
             M_inv, W,
-            N, NQuad, NLeg, NFourier,
-            NLayers, NBDRF,
-            atmos_is_multilayered,
+            N, NQuad, NLeg,
+            NFourier, NLayers, NBDRF,
+            is_atmos_multilayered,
             weighted_scaled_Leg_coeffs,
             BDRF_Fourier_modes,
             mu0, I0, I0_orig, phi0,
@@ -508,6 +565,6 @@ def pydisort(
             there_is_iso_source,
             scale_tau,
             only_flux,
-            use_sparse_NLayers,
-            _is_compatible_with_autograd,
+            NLayers_to_use_sparse,
+            _autograd_compatible,
         )

PythonicDISORT/subroutines.py

@@ -4,7 +4,7 @@ from math import pi
 from numpy.polynomial.legendre import leggauss
 
 
-def transform_interval(arr, c, d, a=-1, b=1):
+def transform_interval(arr, c, d, a, b):
     """Affine transformation of an array from interval [a, b] to [c, d].
 
     Parameters
@@ -29,7 +29,7 @@ def transform_interval(arr, c, d, a=-1, b=1):
     return (((arr - a) * (d - c)) / (b - a)) + c
 
 
-def transform_weights(weights, c, d, a=-1, b=1):
+def transform_weights(weights, c, d, a, b):
     """Transforms an array of quadrature weights from interval [a, b] to [c, d].
 
     Parameters
@@ -106,7 +106,7 @@ def Gauss_Legendre_quad(N, c=0, d=1):
     """
     mu_arr_pos, W = leggauss(int(N))
 
-    return transform_interval(mu_arr_pos, c, d), transform_weights(W, c, d)
+    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)):
@@ -142,7 +142,7 @@ def Clenshaw_Curtis_quad(Nphi, c=0, d=(2 * pi)):
     weights_phi_pos[0] /= 2
     full_weights_phi = np.hstack((weights_phi_pos, np.flip(weights_phi_pos[:-1])))
 
-    return transform_interval(phi_arr, c, d), transform_weights(full_weights_phi, c, d)
+    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):
@@ -256,18 +256,58 @@ def generate_diff_act_flux_funcs(u0):
     return flux_act_up, flux_act_down_diffuse
 
 
-def _mathscr_v(tau, l, s_poly_coeffs, Nscoeffs, G, K, G_inv, mu_arr, _is_compatible_with_autograd=False):
+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?
+                ):
     """Particular solution for isotropic internal sources.
-    It has many seemingly redundant arguments to maximize precomputation
-    in the `pydisort` function which calls it.
-
+    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.
+    
+    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                               |
+    
+    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`                      |
     """
     n = Nscoeffs - 1
     
-    if _is_compatible_with_autograd:
+    if _autograd_compatible:
         import autograd.numpy as np
     
         def mathscr_b(i):
+            """
+            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 = np.arange(i + 1)
             s_poly_coeffs_nj = s_poly_coeffs[:, n - j_arr]
             return np.sum(

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PythonicDISORT
-Version: 0.6.2
+Version: 0.7.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
@@ -46,7 +46,7 @@ We highly recommend reading the non-optional parts of sections 1 and 2 before us
 
 ## PyTest and examples of how to use PythonicDISORT
 
-Not only do we have verification tests in the Jupyter Notebook, we also used PyTest to recreate most of the test problems in Stamnes' `disotest.f90`; 
+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, 
@@ -56,13 +56,13 @@ to show how the tests were implemented. These notebooks double up as examples of
 
 * From PyPI: `pip install PythonicDISORT`
 * From Conda-forge: (TODO: we 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
+* 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
 * Python 3.8+
 * `numpy >= 1.8.0`
 * `scipy >= 1.8.0`
-* (OPTIONAL) `pytest >= 6.2.5` (Required for non-Notebook tests)
+* (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
 * `autograd >= 1.5`

PythonicDISORT-0.7.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+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,,

PythonicDISORT-0.6.2.dist-info/RECORD

@@ -1,11 +0,0 @@
-PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
-PythonicDISORT/_assemble_solution_functions.py,sha256=_VyCjn6WUXE6thYrf-KoGL0sYZI9dotGtlcaJPWo_fc,14849
-PythonicDISORT/_diagonalize.py,sha256=CT3pG1Hi3-iRJGYoJNiGYey4Mlp0yyGS6qYlUXHYYeA,9147
-PythonicDISORT/_solve_for_coefs.py,sha256=4LEcwveWC70-ALDQjwpnOKxb6yXd5BHqJYdxkkhGtvI,15014
-PythonicDISORT/pydisort.py,sha256=wjiIv0MZZ8QOXyYdZrd86YpzDdmIG12_HXAPNoO0P6s,24295
-PythonicDISORT/subroutines.py,sha256=KXaFUIBdt0o8Ag204oBGObcw1-JYUyxJF3_VX_zNEpI,13295
-PythonicDISORT-0.6.2.dist-info/LICENSE.md,sha256=q1i1QVh0OLzK0fj6kaNRNWsM_KZca66JwWgWwbuCPlA,1092
-PythonicDISORT-0.6.2.dist-info/METADATA,sha256=g6MguV17dX4teiOi0iLYSTq82HuOUwzrERa5jOiacPU,7135
-PythonicDISORT-0.6.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-PythonicDISORT-0.6.2.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
-PythonicDISORT-0.6.2.dist-info/RECORD,,