PythonicDISORT 1.4__py3-none-any.whl → 1.6__py3-none-any.whl

PythonicDISORT/_assemble_intensity_and_fluxes.py

@@ -167,7 +167,7 @@ def _assemble_intensity_and_fluxes(
  
         # Construct the intensity function (refer to section 3.7 of the Comprehensive Documentation)
         # --------------------------------------------------------------------------------------------------------------------------
-        def u(tau, phi, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False, *, _return_l=False):
+        def u(tau, phi, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False):
             """
             Intensity function with arguments `(tau, phi)` of types `(array or float, array or float)`.
             Returns an ndarray with axes corresponding to variation with `mu, tau, phi` respectively.
@@ -236,7 +236,7 @@ def _assemble_intensity_and_fluxes(
             if there_is_iso_source:
                 l_uniq, l_inv = np.unique(l, return_inverse=True)
                 _mathscr_v_contribution = _mathscr_v(
-                    scaled_tau, l_inv, 
+                    scaled_tau, scale_tau, l_inv, 
                     Nscoeffs,
                     scaled_s_poly_coeffs[l_uniq],
                     G_collect_0[l_uniq],
@@ -318,13 +318,6 @@ def _assemble_intensity_and_fluxes(
             if return_tau_arr:
                 outputs += (tau_arr,)
                 output_is_tuple = True
-            
-            if _return_l: # This indicates to the function that NT corrections are activated
-                output_is_tuple = True
-                outputs += (l,)
-                if there_is_iso_source:
-                    # The TMS correction will use these
-                    outputs += (l_uniq, l_inv)
                 
             return outputs if output_is_tuple else outputs[0]
         # --------------------------------------------------------------------------------------------------------------------------
@@ -411,7 +404,7 @@ def _assemble_intensity_and_fluxes(
         if there_is_iso_source:
             l_uniq, l_inv = np.unique(l, return_inverse=True)
             _mathscr_v_contribution = _mathscr_v(
-                scaled_tau, l_inv,
+                scaled_tau, scale_tau, l_inv,
                 Nscoeffs,
                 scaled_s_poly_coeffs[l_uniq],
                 G_collect_0[l_uniq],
@@ -470,7 +463,7 @@ def _assemble_intensity_and_fluxes(
         if there_is_iso_source:
             l_uniq, l_inv = np.unique(l, return_inverse=True)
             _mathscr_v_contribution = _mathscr_v(
-                scaled_tau, l_inv,
+                scaled_tau, scale_tau, l_inv,
                 Nscoeffs,
                 scaled_s_poly_coeffs[l_uniq],
                 G_collect_0[l_uniq, :N, :],
@@ -552,7 +545,7 @@ def _assemble_intensity_and_fluxes(
         if there_is_iso_source:
             l_uniq, l_inv = np.unique(l, return_inverse=True)
             _mathscr_v_contribution = _mathscr_v(
-                scaled_tau, l_inv,
+                scaled_tau, scale_tau, l_inv,
                 Nscoeffs,
                 scaled_s_poly_coeffs[l_uniq],
                 G_collect_0[l_uniq, N:, :],

PythonicDISORT/_solve_for_coeffs.py

@@ -168,7 +168,8 @@ def _solve_for_coeffs(
         if m_equals_0 and there_is_iso_source:
             # Top boundary
             RHS[:N] -= _mathscr_v(
-                        np.array([0]), 
+                        np.array([0]),
+                        None,
                         np.array([0]),
                         Nscoeffs,
                         scaled_s_poly_coeffs[[0], :],
@@ -183,6 +184,7 @@ def _solve_for_coeffs(
                 RHS[N:-N] += (
                     _mathscr_v(
                         scaled_tau_arr_with_0[1:-1],
+                        None,
                         indices,
                         Nscoeffs,
                         scaled_s_poly_coeffs[indicesp1],
@@ -192,6 +194,7 @@ def _solve_for_coeffs(
                     )
                     - _mathscr_v(
                         scaled_tau_arr_with_0[1:-1],
+                        None,
                         indices,
                         Nscoeffs,
                         scaled_s_poly_coeffs[indices],
@@ -203,7 +206,8 @@ def _solve_for_coeffs(
                 
             # Bottom boundary
             RHS[-N:] -= _mathscr_v(
-                scaled_tau_arr_with_0[[-1]], 
+                scaled_tau_arr_with_0[[-1]],
+                None,
                 np.array([0]),
                 Nscoeffs,
                 scaled_s_poly_coeffs[[-1], :],
@@ -217,6 +221,7 @@ def _solve_for_coeffs(
                     R
                     @ _mathscr_v(
                         scaled_tau_arr_with_0[[-1]],
+                        None,
                         np.array([0]),
                         Nscoeffs,
                         scaled_s_poly_coeffs[[-1], :],

PythonicDISORT/_solve_for_gen_and_part_sols.py

@@ -198,7 +198,8 @@ def _solve_for_gen_and_part_sols(
         G_collect[no_shortcut_indices, N:, :N] = GpG_m_GmG_div_2
         
         K_collect[no_shortcut_indices, :] = K_arr
-        if len(no_shortcut_indices_0) > 0: # If there is no isotropic source this list will be empty
+        len_no_shortcut_indices_0 = len(no_shortcut_indices_0)
+        if len_no_shortcut_indices_0 > 0: # If there is no isotropic source this list will be empty
             G_inv_collect_0[no_shortcut_indices_0, :, :] = np.linalg.inv(
                 G_collect[no_shortcut_indices_0, :, :]
             )
@@ -209,10 +210,25 @@ def _solve_for_gen_and_part_sols(
         
             # Solve for particular solution w.r.t. direct beam source (refer to section 3.6.1 of the Comprehensive Documentation)
             # --------------------------------------------------------------------------------------------------------------------------
-            A_arr[:, :N, N:] = -A_arr[:, N:, :N]
-            A_arr[:, :N, :N] = -A_arr[:, N:, N:]
-            A_arr.reshape(-1, NQuad * NQuad)[:, ::NQuad + 1] += 1 / mu0
-            B_collect[no_shortcut_indices, :] = np.linalg.solve(A_arr, X_arr[no_shortcut_indices, :, None]).squeeze(-1)
+            if len_no_shortcut_indices_0 > 0: # In this case we have `G_inv_collect_0`, so use it
+                B_collect[no_shortcut_indices_0, :] = np.einsum(
+                    "bij, bjk, bk -> bi",
+                    # Note that the minus sign has been absorbed into `X_arr`
+                    G_collect[no_shortcut_indices_0, :, :]
+                    / (1 / mu0 + K_collect[no_shortcut_indices_0, None, :]),
+                    G_inv_collect_0[no_shortcut_indices_0, :, :],
+                    X_arr[no_shortcut_indices_0, :],
+                    optimize=True,
+                )
+                no_shortcut_indices = no_shortcut_indices[len_no_shortcut_indices_0:]
+            
+            if len(no_shortcut_indices) > 0:
+                A_arr = A_arr[len_no_shortcut_indices_0:]
+                A_arr[:, :N, N:] = -A_arr[:, N:, :N]
+                A_arr[:, :N, :N] = -A_arr[:, N:, N:]
+                A_arr.reshape(-1, NQuad * NQuad)[:, ::NQuad + 1] += 1 / mu0
+                # Note that the minus sign has been absorbed into `X_arr`
+                B_collect[no_shortcut_indices, :] = np.linalg.solve(A_arr, X_arr[no_shortcut_indices, :, None]).squeeze(-1)
             
             # --------------------------------------------------------------------------------------------------------------------------
 

PythonicDISORT/pydisort.py

@@ -98,7 +98,7 @@ def pydisort(
         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 this function with respect to ``tau``.
-        For example, ``Fp(tau_arr, is_antiderivative_wrt_tau = True) - Fp(np.insert(tau_arr[:-1] + 1e-15, 0, 0), is_antiderivative_wrt_tau = True)``
+        For example, ``Fp(tau_arr, True) - Fp(np.insert(tau_arr[:-1] + 1e-15, 0, 0), True)``
         will produce an array of the tau-integral over each layer.
         Pass ``return_tau_arr`` to return ``tau_arr`` (defaults to ``False``).
     Fm(tau) : function
@@ -406,19 +406,21 @@ def pydisort(
         
         # TMS correction for the intensity (see section 3.7.2)
         # --------------------------------------------------------------------------------------------------------------------------
-        def TMS_correction(tau, phi, is_antiderivative_wrt_tau, l, l_uniq, l_inv):
+        def TMS_correction(tau, phi, is_antiderivative_wrt_tau):
             Ntau = len(tau)
             Nphi = len(phi)
             # Atmospheric layer indices
-            if not there_is_iso_source:
-                l_uniq, l_inv = np.unique(l, return_inverse=True)
-            scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
+            l = np.argmax(tau[:, None] <= tau_arr[None, :], axis=1)
+            scaled_tau_arr_l = scaled_tau_arr_with_0[l + 1]
             scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
 
             # Delta-M scaling
-            tau_dist_from_top = tau_arr[l] - tau
-            scaled_tau_dist_from_top = tau_dist_from_top * scale_tau[l]
-            scaled_tau = scaled_tau_arr_l - scaled_tau_dist_from_top
+            if np.any(scale_tau != np.ones(NLayers)):
+                tau_dist_from_top = tau_arr[l] - tau
+                scaled_tau_dist_from_top = tau_dist_from_top * scale_tau[l]
+                scaled_tau = scaled_tau_arr_l - scaled_tau_dist_from_top
+            else:
+                scaled_tau = tau
             
             # mathscr_B
             # --------------------------------------------------------------------------------------------------------------------------
@@ -429,45 +431,44 @@ def pydisort(
                 [
                     f(nu)[:, None, :]
                     # Iterates over the 0th axis
-                    for f in map(Legendre, weighted_Leg_coeffs_all[l_uniq])
+                    for f in map(Legendre, weighted_Leg_coeffs_all)
                 ],
                 axis=1,
-            ) # (NQuad, NLayers_uniq, Nphi)
+            )
             p_trun = np.concatenate(
                 [
                     f(nu)[:, None, :]
                     # Iterates over the 0th axis
-                    for f in map(Legendre, weighted_scaled_Leg_coeffs[l_uniq])
+                    for f in map(Legendre, weighted_scaled_Leg_coeffs)
                 ],
                 axis=1,
-            ) # (NQuad, NLayers_uniq, Nphi)
-            
+            )
             mathscr_B = (
-                (scaled_omega_arr[l_uniq] * I0_div_4pi)[None, :, None]
+                (scaled_omega_arr * I0)[None, :, None]
+                / (4 * np.pi)
                 * (mu0 / (mu0 + mu_arr))[:, None, None]
-                * (p_true / (1 - f_arr)[None, l_uniq, None] - p_trun)
-            ) # (NQuad, NLayers_uniq, Nphi)
+                * (p_true / (1 - f_arr[None, :, None]) - p_trun)
+            )
             # --------------------------------------------------------------------------------------------------------------------------
             
             neg_scaled_tau_div_mu0 = -scaled_tau / mu0
             if is_antiderivative_wrt_tau:
 
-                scale_tau_l = scale_tau[l_uniq]
-                neg_scale_tau_l_div_mu0 = -scale_tau_l / mu0
-                scale_tau_l_div_mu_arr_pos = scale_tau_l[None, :] / mu_arr_pos[:, None]
+                neg_scale_tau_div_mu0 = -scale_tau / mu0
+                scale_tau_div_mu_arr_pos = scale_tau[None, :] / mu_arr_pos[:, None]
 
                 TMS_correction_pos = (
-                    np.exp(neg_scaled_tau_div_mu0) / neg_scale_tau_l_div_mu0[l_inv]
+                    np.exp(neg_scaled_tau_div_mu0) / neg_scale_tau_div_mu0[l]
                 )[None, :] - np.exp(
                     (scaled_tau - scaled_tau_arr_l)[None, :] / mu_arr_pos[:, None]
                     - scaled_tau_arr_l[None, :] / mu0
-                ) / scale_tau_l_div_mu_arr_pos[:, l_inv]
+                ) / scale_tau_div_mu_arr_pos[:, l]
                 TMS_correction_neg = (
-                    np.exp(neg_scaled_tau_div_mu0) / neg_scale_tau_l_div_mu0[l_inv]
+                    np.exp(neg_scaled_tau_div_mu0) / neg_scale_tau_div_mu0[l]
                 )[None, :] + np.exp(
                     (scaled_tau_arr_lm1 - scaled_tau)[None, :] / mu_arr_pos[:, None]
                     - scaled_tau_arr_lm1[None, :] / mu0
-                ) / scale_tau_l_div_mu_arr_pos[:, l_inv]
+                ) / scale_tau_div_mu_arr_pos[:, l]
             else:
                 TMS_correction_pos = np.exp(neg_scaled_tau_div_mu0)[None, :] - np.exp(
                     (scaled_tau - scaled_tau_arr_l)[None, :] / mu_arr_pos[:, None]
@@ -477,123 +478,111 @@ def pydisort(
                     (scaled_tau_arr_lm1 - scaled_tau)[None, :] / mu_arr_pos[:, None]
                     - scaled_tau_arr_lm1[None, :] / mu0
                 )
-
             
             ## Contribution from other layers
             # --------------------------------------------------------------------------------------------------------------------------
-            solution = (
-                mathscr_B[:, l_inv, :]
-                * np.concatenate((TMS_correction_pos, TMS_correction_neg), axis=0)[:, :, None]
-            )
+            Bpos = mathscr_B[:N, :, :]                                      # (N, NLayers, Nphi)
+            Bneg = mathscr_B[N:, :, :]                                      # (N, NLayers, Nphi)
             
-            if is_atmos_multilayered:
-                any_pos = (l.min() < NLayers - 1)
-                any_neg = (l.max() > 0)
-
-                mu0_inv = 1 / mu0
+            sol_pos = Bpos[:, l, :] * TMS_correction_pos[:, :, None]        # (N, Ntau, Nphi)
+            sol_neg = Bneg[:, l, :] * TMS_correction_neg[:, :, None]        # (N, Ntau, Nphi)
+            
+            if not is_atmos_multilayered:
+                return np.concatenate((sol_pos, sol_neg), axis=0)
 
-                scaled_tauwith0_front = scaled_tau_arr_with_0[:-1]          # (NLayers,)
-                scaled_tauwith0_back  = scaled_tau_arr_with_0[1:]           # (NLayers,); i.e. `scaled_tau_arr`
+            # ---------------- other-layer contributions ----------------
+            any_pos = (l.min() < NLayers - 1)  # exists tau with layers below
+            any_neg = (l.max() > 0)            # exists tau with layers above
 
-                exp_taufront_mu0_inv = np.concatenate(
-                    ([1], np.exp(-scaled_tauwith0_front[1:] * mu0_inv))
-                )  # (Nlayers, )
+            # Layer boundaries/thickness in scaled-tau coordinates
+            tau_front = scaled_tau_arr_with_0[:-1]                           # (NLayers,) = tau_r
+            tau_back  = scaled_tau_arr_with_0[1:]                            # (NLayers,) = tau_{r+1}
 
+            # decay[n,r] = exp(-scaled_thickness_arr[r]/mu[n])  (argument <= 0; underflow is safe)
+            decay = np.exp(-scaled_thickness_arr[None, :] * M_inv[:, None])                  # (N, NLayers)
 
-                Bpos = mathscr_B[:N, :, :]                       # (N, NLayers_uniq, Nphi)
-                Bneg = mathscr_B[N:, :, :]                       # (N, NLayers_uniq, Nphi)
+            exp_taufront_mu0_inv = np.exp(-tau_front / mu0)              # (NLayers,)
+            exp_tauback_mu0_inv  = np.exp(-tau_back / mu0)              # (NLayers,)
 
-                # decay[n,r] = exp(-scaled_thickness_arr[r]/mu[n]); always <= 1
-                log_decay = -scaled_thickness_arr[None, :] * M_inv[:, None]                    # (N, NLayers)
+            if is_antiderivative_wrt_tau:
+                # (N, NLayers): multiply by this instead of dividing by (scale_tau/mu)
+                mu_div_scale_layers = mu_arr_pos[:, None] / scale_tau[None, :]
 
-                # decay_prod[n,k] = prod_{m=0}^{k-1} decay[n,m], with decay_prod[:,0]=1
-                log_decay_prod = np.cumsum(log_decay, axis=1) # The cumsum requires ALL layers; `l_uniq` is inapplicable
-                decay_prod = np.concatenate((np.ones((N, 1)), np.exp(log_decay_prod)), axis=1) # (N, NLayers + 1)
+            # ===== POS: sum over r > l =====
+            if any_pos:
+                # fac_pos = 1 - exp(-scaled_thickness_arr*(1/mu + 1/mu0)) = -expm1(-scaled_thickness_arr*(...))
+                delta_pos = scaled_thickness_arr[None, :] * (M_inv + 1 / mu0)[:, None]        # (N, NLayers) >= 0
+                fac_pos = -np.expm1(-delta_pos)                               # (N, NLayers) in [0,1)
 
                 if is_antiderivative_wrt_tau:
-                    # (N, NLayers)
-                    integration_factor = mu_arr_pos[:, None] / scale_tau
-
-                # ====================================================================
-                # POS: pre-accumulate "contribution from below" sums Rpos[N, l, Nphi]
-                # ====================================================================
-                if any_pos:
-                    thickness_pos = scaled_thickness_arr[None, :] * (M_inv + mu0_inv)[:, None]  # (N, NLayers); always > 0
-                    # em1_pos = 1 - exp(-scaled_thickness_arr*(1/mu + 1/mu0)) = -expm1(-thickness_pos)
-                    em1_pos = -np.expm1(-thickness_pos)                                         # (N, NLayers)
-                    if is_antiderivative_wrt_tau:
-                        em1_pos = integration_factor * em1_pos
-
-                    layer_term_pos = em1_pos * exp_taufront_mu0_inv[None, :]                    # (N, NLayers)
-
-                    # Reverse-scan via cumsum:
-                    # Rpos[l] = (1/decay_prod[l+1]) * sum_{r=l+1..} (layer_term_pos[r] * decay_prod[r])
-                    weighted_layer_term_pos = layer_term_pos * decay_prod[:, :NLayers]          # (N, NLayers)
-                    cumsum_pos = np.cumsum(weighted_layer_term_pos[:, ::-1], axis=1)[:, ::-1]   # (N, NLayers)
-
-                    Rpos_rest = cumsum_pos[:, 1:] / decay_prod[:, 1 : NLayers]                  # (N, NLayers - 1)
-                    Rpos = np.concatenate([Rpos_rest, np.zeros((N, 1))], axis=1)                # (N, NLayers)
-
-                    # Apply to all `tau`
-                    # expfac_pos = exp(M_inv*tau) * exp(-M_inv*scaled_tauwith0_back[l])
-                    expfac_pos = np.exp(M_inv[:, None] * (scaled_tau - scaled_tauwith0_back[l])[None, :])  # (N, Ntau)
-                    tmp = np.take(Rpos, l, axis=1)                                                         # (N, Ntau)
-                    tmp = (tmp * expfac_pos)[:, :, None] * Bpos[:, l_inv, :]
-                    
-                    if autograd_compatible:
-                        solution = np.concatenate((solution[:N, :, :] + tmp, solution[N:, :, :]), axis=0)
-                    else:
-                        solution[:N, :, :] += tmp
-
-                # ====================================================================
-                # NEG: pre-accumulate "contribution from above" sums Rneg[N, l, Nphi]
-                # ====================================================================
-                if any_neg:
-                    decay = np.exp(log_decay)                                                    # (N, NLayers)
-                    
-                    # Ensure no positive exponents
-                    thickness_neg = scaled_thickness_arr[None, :] * (M_inv - mu0_inv)[:, None]   # (N, NLayers)
-                    mask = (thickness_neg >= 0)
+                    fac_pos = fac_pos * mu_div_scale_layers                  # (N, NLayers)
 
-                    exp_x1 = np.exp(-scaled_tauwith0_back * mu0_inv)[None, :]                    # (1, NLayers)
-                    exp_x0 = decay * exp_taufront_mu0_inv[None, :]                               # (N, NLayers)
+                # coef_pos[n,r] multiplies exp((tau - tau_front[r])/mu - tau_front[r]/mu0) after re-referencing
+                # includes exp(-tau_front[r]/mu0)
+                coef_pos = fac_pos * exp_taufront_mu0_inv[None, :]            # (N, NLayers)
 
-                    # Let a = |thickness_neg| >= 0. Then expm1(-a) uses only non-positive inputs.
-                    em1_neg = np.expm1(-np.abs(thickness_neg))                                       # (N, NLayers)
+                # Reverse scan to build Acc_l (referenced at tau_back[l]) INCLUDING Bpos:
+                # Acc_{r-1} = coef_pos[:,r]*Bpos[:,r,:] + decay[:,r]*Acc_r
+                acc0 = np.zeros((N, Nphi))
+                acc = acc0
+                cols_rev = [acc0]  # Acc_{NLayers-1} = 0
 
-                    # thickness_neg>=0: exp_x1 * (1 - exp(-a)) = exp_x1 * (-expm1(-a)) = exp_x1 * (-em1_neg)
-                    # thickness_neg<0 : exp_x0 * (exp(-a) - 1) = exp_x0 * expm1(-a)    = exp_x0 * ( em1_neg)
-                    layer_term_negthick = em1_neg * exp_x0                                       # (N, NLayers)
-                    layer_term_posthick = -em1_neg * exp_x1                                      # (N, NLayers)
-                    layer_term_neg = layer_term_posthick * mask + layer_term_negthick * (~mask)  # (N, NLayers)
+                for r in range(NLayers - 1, 0, -1):  # r = NLayers-1 ... 1
+                    term = coef_pos[:, r][:, None] * Bpos[:, r, :]            # (N, Nphi)
+                    acc = term + decay[:, r][:, None] * acc
+                    cols_rev.append(acc)
 
-                    if is_antiderivative_wrt_tau:
-                        layer_term_neg = -integration_factor * layer_term_neg
+                Acc_pos = np.stack(cols_rev[::-1], axis=1)                    # (N, NLayers, Nphi)
 
-                    # Forward-scan via cumsum:
-                    # Rneg[l] = decay_prod[l] * sum_{r<l} layer_term_neg[r] / decay_prod[r+1]
-                    weighted_layer_term_neg = layer_term_neg / decay_prod[:, 1:]                 # (N, NLayers)
-                    cumsum_neg = np.cumsum(weighted_layer_term_neg, axis=1)                      # (N, NLayers)
-                    
-                    Rneg_rest = decay_prod[:, 1 : NLayers] * cumsum_neg[:, :-1]                  # (N, NLayers - 1)
-                    Rneg = np.concatenate([np.zeros((N, 1)), Rneg_rest], axis=1)                 # (N, NLayers)
+                # Apply to tau: multiply by exp((scaled_tau - tau_back[l]) / mu) (argument <= 0 in exact arithmetic)
+                expfac_pos = np.exp(M_inv[:, None] * (scaled_tau - tau_back[l])[None, :])  # (N, Ntau)
+                contrib_pos = Acc_pos[:, l, :] * expfac_pos[:, :, None]         # (N, Ntau, Nphi)
 
-                    # Apply to all `tau`
-                    expfac_neg = np.exp(M_inv[:, None] * (scaled_tauwith0_front[l] - scaled_tau)[None, :])  # (N, Ntau)
-                    tmp = np.take(Rneg, l, axis=1)                                                          # (N, Ntau)
-                    tmp = (tmp * expfac_neg)[:, :, None] * Bneg[:, l_inv, :]
-                    
-                    if autograd_compatible:
-                        solution = np.concatenate((solution[:N, :, :], solution[N:, :, :] + tmp), axis=0)
-                    else:
-                        solution[N:, :, :] += tmp
+                sol_pos = sol_pos + contrib_pos
+
+            # ===== NEG: sum over r < l =====
+            if any_neg:
+                # We need term_base[n,r] = exp(-tau_back[r]/mu0) - exp(-scaled_thickness_arr[r]/mu)*exp(-tau_front[r]/mu0)
+                # computed stably via expm1 using only non-positive inputs.
+                delta_neg = scaled_thickness_arr[None, :] * (M_inv - 1 / mu0)[:, None]        # (N, NLayers), can be +/-.
+                abs_delta = np.abs(delta_neg)
+                em1 = np.expm1(-abs_delta)                                    # (N, NLayers) in [-1,0], argument <= 0
+
+                exp_x1 = exp_tauback_mu0_inv[None, :]                         # (1, NLayers)
+                exp_x0 = decay * exp_taufront_mu0_inv[None, :]                # (N, NLayers)
+
+                # if delta>=0: term = exp_x1*(1-exp(-delta)) = exp_x1*(-expm1(-delta)) = exp_x1*(-em1)
+                # if delta<0 : term = exp_x0*(exp(delta)-1)  = exp_x0*expm1(delta)     = exp_x0*( em1)  (since delta<0 -> em1=expm1(delta))
+                term_base = np.where(delta_neg >= 0.0, (-em1) * exp_x1, em1 * exp_x0)  # (N, NLayers)
+
+                if is_antiderivative_wrt_tau:
+                    # original divides by (-scale_tau/mu) == multiply by (-(mu/scale_tau))
+                    coef_neg = term_base * -mu_div_scale_layers               # (N, NLayers)
+                else:
+                    coef_neg = term_base                                      # (N, NLayers)
+
+                # Forward scan to build Acc_l (referenced at tau_front[l]) INCLUDING Bneg:
+                # Acc_{r+1} = coef_neg[:,r]*Bneg[:,r,:] + decay[:,r]*Acc_r
+                acc0 = np.zeros((N, Nphi))
+                acc = acc0
+                cols = [acc0]  # Acc_0 = 0
+
+                for r in range(0, NLayers - 1):  # r = 0 ... NLayers-2
+                    term = coef_neg[:, r][:, None] * Bneg[:, r, :]            # (N, Nphi)
+                    acc = term + decay[:, r][:, None] * acc
+                    cols.append(acc)
+
+                Acc_neg = np.stack(cols, axis=1)                              # (N, NLayers, Nphi)
+
+                # Apply to tau: multiply by exp((tau_front[l] - scaled_tau) / mu) (argument <= 0 in exact arithmetic)
+                expfac_neg = np.exp(M_inv[:, None] * (tau_front[l] - scaled_tau)[None, :])  # (N, Ntau)
+                contrib_neg = Acc_neg[:, l, :] * expfac_neg[:, :, None]          # (N, Ntau, Nphi)
 
-                return solution
+                sol_neg = sol_neg + contrib_neg
 
+            return np.concatenate((sol_pos, sol_neg), axis=0)
+                    
             # --------------------------------------------------------------------------------------------------------------------------
-            
-            else:
-                return solution
         # --------------------------------------------------------------------------------------------------------------------------
 
         # IMS correction for the intensity (see section 3.7.2)
@@ -628,7 +617,8 @@ def pydisort(
                 ) / (mu_arr_pos * scaled_mu0 * x)[:, None]
                 
             return (
-                I0_div_4pi
+                I0
+                / (4 * pi)
                 * (omega_avg * f_avg) ** 2
                 / (1 - omega_avg * f_avg)
                 * Legendre(
@@ -640,59 +630,25 @@ def pydisort(
 
         # The corrected intensity
         # --------------------------------------------------------------------------------------------------------------------------
-        def u_corrected(
-            tau,
-            phi,
-            is_antiderivative_wrt_tau=False,
-            return_Fourier_error=False,
-            return_tau_arr=False,
-        ):
+        def u_corrected(tau, phi, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False):
             tau = np.atleast_1d(tau)
             phi = np.atleast_1d(phi)
-
-            u_star_outputs = u_star(
-                tau,
-                phi,
-                is_antiderivative_wrt_tau,
-                return_Fourier_error,
-                return_tau_arr,
-                _return_l=True,
-            )
-            if there_is_iso_source:
-                NT_corrections = TMS_correction(
-                    tau,
-                    phi,
-                    is_antiderivative_wrt_tau,
-                    u_star_outputs[-3],
-                    u_star_outputs[-2],
-                    u_star_outputs[-1],
-                )
-                u_star_outputs = u_star_outputs[:-3]
-            else:
-                NT_corrections = TMS_correction(
-                    tau, phi, is_antiderivative_wrt_tau, u_star_outputs[-1], None, None
-                )
-                u_star_outputs = u_star_outputs[:-1]
-
+            NT_corrections = TMS_correction(tau, phi, is_antiderivative_wrt_tau)
+            
             if autograd_compatible:
-                NT_corrections = np.concatenate(
-                    (
-                        NT_corrections[:N, :, :],
-                        NT_corrections[N:, :, :]
-                        + IMS_correction(tau, phi, is_antiderivative_wrt_tau),
-                    ),
-                    axis=0,
+                NT_corrections = NT_corrections + np.concatenate(
+                    [np.zeros((N, len(tau), len(phi))), IMS_correction(tau, phi, is_antiderivative_wrt_tau)], axis=0
                 )
             else:
                 NT_corrections[N:, :, :] += IMS_correction(tau, phi, is_antiderivative_wrt_tau)
-
+            
             if return_Fourier_error or return_tau_arr:
+                u_star_outputs = u_star(tau, phi, is_antiderivative_wrt_tau, return_Fourier_error, return_tau_arr)
                 return (
                     u_star_outputs[0] + rescale_factor * np.squeeze(NT_corrections),
                 ) + u_star_outputs[1:]
             else:
-                return u_star_outputs[0] + rescale_factor * np.squeeze(NT_corrections)
-
+                return u_star(tau, phi, is_antiderivative_wrt_tau) + rescale_factor * np.squeeze(NT_corrections)
         # --------------------------------------------------------------------------------------------------------------------------
 
         return mu_arr, flux_up, flux_down, u0, u_corrected

PythonicDISORT/subroutines.py

@@ -744,6 +744,7 @@ def to_diag_ordered_form(A, Nsuperdiags, Nsubdiags):
 
     
 def _mathscr_v(tau,                             # Input optical depths
+               scale_tau,                       # Delta-scale factor for tau
                l,                               # Layer index of each input optical depth
                Nscoeffs,                        # Number of isotropic source polynomial coefficients
                s_poly_coeffs,                   # Polynomial coefficients of isotropic source
@@ -763,6 +764,7 @@ def _mathscr_v(tau,                             # Input optical depths
     |          Variable             |                   Shape                   |
     | ----------------------------- | ----------------------------------------- |
     | ``tau``                       | ``Ntau``                                  |
+    | ``scale_tau``                 | `NLayers`                                 |
     | ``l``                         | ``Ntau``                                  |
     | ``Nscoeffs``                  | scalar                                    |
     | ``s_poly_coeffs``             | ``NL x Nscoeffs``                         |
@@ -784,13 +786,13 @@ def _mathscr_v(tau,                             # Input optical depths
     if Nscoeffs == 1:
         s0 = s_poly_coeffs[:, 0:1]
         c0w = s0 * K_inv * G_inv_mu_inv  # (NL, 1) * (NL, NQuad) -> (NL, NQuad)
+        
+        d0 = np.einsum("lij, lj -> il", G, c0w, optimize=True)
 
         if is_antiderivative_wrt_tau:
-            p0 = tau                 
+            p0 = tau / scale_tau                 
         else:
             p0 = np.ones_like(tau)       
-            
-        d0 = np.einsum("lij, lj -> il", G, c0w, optimize=True)
 
         return d0[:, l] * p0[None, :]
 
@@ -803,19 +805,19 @@ def _mathscr_v(tau,                             # Input optical depths
 
         c0w = s1 * K_inv * G_inv_mu_inv
         c1w = (s0 * K_inv + s1 * K_inv2) * G_inv_mu_inv
+        
+        d0 = np.einsum("lij, lj -> il", G, c0w, optimize=True)  # (NQuad, NL)
+        d1 = np.einsum("lij, lj -> il", G, c1w, optimize=True)  # (NQuad, NL)
 
         if is_antiderivative_wrt_tau:
             p0 = tau * tau / 2
             p1 = tau
+            return (d0[:, l] * p0[None, :] + d1[:, l] * p1[None, :]) / scale_tau
         else:
             p0 = tau
             p1 = np.ones_like(tau) 
-
-        d0 = np.einsum("lij, lj -> il", G, c0w, optimize=True)  # (NQuad, NL)
-        d1 = np.einsum("lij, lj -> il", G, c1w, optimize=True)  # (NQuad, NL)
-
-        return d0[:, l] * p0[None, :] + d1[:, l] * p1[None, :]
-
+            return d0[:, l] * p0[None, :] + d1[:, l] * p1[None, :]
+        
     # --------------------- Nscoeffs > 2 ----------------------
     else:
         if Nscoeffs > 10:
@@ -853,7 +855,7 @@ def _mathscr_v(tau,                             # Input optical depths
         powers = np.arange(Nscoeffs - 1, -1, -1)[None, :]
         if is_antiderivative_wrt_tau:
             p = powers + 1
-            tau_poly = (tau[:, None] ** p) / p
+            tau_poly = (tau[:, None] ** p) / (p * scale_tau)
         else:
             tau_poly = tau[:, None] ** powers
 

{pythonicdisort-1.4.dist-info → pythonicdisort-1.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PythonicDISORT
-Version: 1.4
+Version: 1.6
 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
@@ -16,6 +16,7 @@ Provides-Extra: notebook_dependencies
 Requires-Dist: autograd >=1.5 ; extra == 'notebook_dependencies'
 Requires-Dist: jupyter >1.0.0 ; extra == 'notebook_dependencies'
 Requires-Dist: notebook >6.5.2 ; extra == 'notebook_dependencies'
+Requires-Dist: matplotlib >=3.6.0 ; extra == 'notebook_dependencies'
 Provides-Extra: pytest
 Requires-Dist: pytest >=6.2.5 ; extra == 'pytest'
 
@@ -40,7 +41,7 @@ You may contact me, Dion, through dh3065@columbia.edu.
 
 The **GitHub repository** is https://github.com/LDEO-CREW/Pythonic-DISORT.
 
-Accompanying Journal of Open Source Software paper: https://joss.theoj.org/papers/10.21105/joss.06442.
+Accompanying Journal of Open Source Software paper: https://joss.theoj.org/papers/10.21105/joss.06442. **Please cite this paper if you use PythonicDISORT**.
 
 # Documentation
 https://pythonic-disort.readthedocs.io/en/latest/
@@ -70,18 +71,19 @@ These notebooks double up as examples of how to use PythonicDISORT. The tests wh
 * Test Problem 8: Absorbing / Isotropic-Scattering Medium (multiple layers)
 * Test Problem 9: General Emitting / Absorbing / Scattering Medium (multiple layers)
 * Test Problem 11: Single-Layer vs. Multiple Layers (no corresponding Jupyter Notebook)
+* Test Problem I: Antiderivative / integration functionality (no corresponding Jupyter Notebook)
 
 # Installation
 
 * From PyPI: `pip install PythonicDISORT`
-* From Conda-forge: (TODO: need to first publish on Conda-forge)
+* From Conda-Forge: `conda install -c conda-forge pythonicdisort` or `mamba install -c conda-forge pythonicdisort`
 * 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) `autograd >= 1.5`
+* (OPTIONAL) `autograd >= 1.5` (`autograd` can be used to compute gradients of the output functions)
 * (OPTIONAL) `pytest >= 6.2.5` (Required to use the command `pytest`, see *PyTest and examples of how to use PythonicDISORT*)
 
 ## (OPTIONAL) Additional requirements to run the Jupyter Notebook
@@ -89,7 +91,8 @@ These notebooks double up as examples of how to use PythonicDISORT. The tests wh
 * `notebook > 6.5.2`
 * `matplotlib >= 3.6.0`
 
-In addition, a F2PY-wrapped Stamnes' DISORT, or equivalent, is required to properly run the last section (section 6).
+In addition, a F2PY-wrapped Stamnes' DISORT or equivalent (https://github.com/kconnour/pyRT_DISORT may work) 
+is required to properly run the last section (section 6).
 
 ## Compatibility
 

pythonicdisort-1.6.dist-info/RECORD

@@ -0,0 +1,11 @@
+PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
+PythonicDISORT/_assemble_intensity_and_fluxes.py,sha256=nN2Kn7UCB6tYlVjKj1Hk_IwNO420uibsSD6o9Lx9X5o,29520
+PythonicDISORT/_solve_for_coeffs.py,sha256=5hXhKi_wkEBAcRmoeauGWDygmH3aB0g9QPZN7I2gaIM,19529
+PythonicDISORT/_solve_for_gen_and_part_sols.py,sha256=5p5AlosW2ALjUBn6aEcPj_1pDpfxMn4MwwGorf9WLUY,14177
+PythonicDISORT/pydisort.py,sha256=GRVU2cEzM1ySuPtSEeYtmRdEPnLa3vLzfA8xd1vb95w,37860
+PythonicDISORT/subroutines.py,sha256=YscXrzzT2pKexe4S-30dfpMX69oUqyvQPnaMu8C1uGc,35223
+pythonicdisort-1.6.dist-info/LICENSE.md,sha256=q1i1QVh0OLzK0fj6kaNRNWsM_KZca66JwWgWwbuCPlA,1092
+pythonicdisort-1.6.dist-info/METADATA,sha256=uc2na6BfMyUgBFE1GSQe3RmekJa4EqIV2t-EK8ZPOAk,8699
+pythonicdisort-1.6.dist-info/WHEEL,sha256=WnJ8fYhv8N4SYVK2lLYNI6N0kVATA7b0piVUNvqIIJE,91
+pythonicdisort-1.6.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
+pythonicdisort-1.6.dist-info/RECORD,,

pythonicdisort-1.4.dist-info/RECORD

@@ -1,11 +0,0 @@
-PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
-PythonicDISORT/_assemble_intensity_and_fluxes.py,sha256=mNhV6KjXWiZwPFCH8nuI6hPssl4PXznOn2nkdqOmS9g,29823
-PythonicDISORT/_solve_for_coeffs.py,sha256=8ok2xS6emv4Qtnj_5CTQkYRH8jYzKE06ypmZErQWTPg,19384
-PythonicDISORT/_solve_for_gen_and_part_sols.py,sha256=OjMn1JrhyC7VqZUUMjJVloNvSRWd8l-dQ1MxGt2QZbc,13216
-PythonicDISORT/pydisort.py,sha256=hHN4u6eG7w7QftlgS7fy-EKwU1a_aJQgk502S5kRdbo,39682
-PythonicDISORT/subroutines.py,sha256=Yv81a57KBBijz5VFWKxQ99mnogJo2TllazL5IVZOTJU,34940
-pythonicdisort-1.4.dist-info/LICENSE.md,sha256=q1i1QVh0OLzK0fj6kaNRNWsM_KZca66JwWgWwbuCPlA,1092
-pythonicdisort-1.4.dist-info/METADATA,sha256=TkrQNF7KQEkfFAtkb5B2LYSqcmG6Sy29eMgVOyel-zY,8309
-pythonicdisort-1.4.dist-info/WHEEL,sha256=WnJ8fYhv8N4SYVK2lLYNI6N0kVATA7b0piVUNvqIIJE,91
-pythonicdisort-1.4.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
-pythonicdisort-1.4.dist-info/RECORD,,