PythonicDISORT 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

PythonicDISORT/_loop_and_assemble_results.py

@@ -23,7 +23,8 @@ def _loop_and_assemble_results(
     Nscoeffs,
     scale_tau,
     only_flux,
-    use_sparse_NLayers
+    use_sparse_NLayers,
+    n_jobs
 ):  
     """This function is wrapped by the `pydisort` function.
     It should be called through `pydisort` and never directly.
@@ -77,10 +78,35 @@ def _loop_and_assemble_results(
         if I0 > 0:
             B_collect = np.empty((NLoops, NLayers, NQuad))
             B_collect[0, :, :] = B_collect_0
-        # TODO: Look into the "xarray.apply_ufunc" method or "dask delayed / bag" to parallelize this
-        # May need to code an xarray wrapper. Dask works with Python code.
-        for m in range(1, NLoops): 
-            outputs = _one_Fourier_mode(
+                
+        if n_jobs == 1:
+            for m in range(1, NLoops): 
+                outputs = _one_Fourier_mode(
+                    m,
+                    scaled_omega_arr,
+                    tau_arr,
+                    scaled_tau_arr_with_0,
+                    mu_arr_pos, mu_arr,
+                    M_inv, W,
+                    N, NQuad, NLeg,
+                    NLayers, NBDRF,
+                    weighted_scaled_Leg_coeffs,
+                    weighted_Leg_coeffs_BDRF,
+                    mu0, I0,
+                    b_pos, b_neg,
+                    scalar_b_pos, scalar_b_neg,
+                    s_poly_coeffs,
+                    Nscoeffs,
+                    use_sparse_NLayers
+                )
+                if I0 > 0:
+                    GC_collect[m, :, :, :], K_collect[m, :, :], B_collect[m, :, :] = outputs
+                else:
+                    GC_collect[m, :, :, :], K_collect[m, :, :] = outputs
+        else:
+            from joblib import Parallel, delayed
+
+            _one_Fourier_mode_m = lambda m: _one_Fourier_mode(
                 m,
                 scaled_omega_arr,
                 tau_arr,
@@ -98,10 +124,16 @@ def _loop_and_assemble_results(
                 Nscoeffs,
                 use_sparse_NLayers
             )
-            if I0 > 0:
-                GC_collect[m, :, :, :], K_collect[m, :, :], B_collect[m, :, :] = outputs
-            else:
-                GC_collect[m, :, :, :], K_collect[m, :, :] = outputs
+            outputs = Parallel(n_jobs=n_jobs)(
+                delayed(_one_Fourier_mode_m)(m) for m in range(1, NLoops)
+            )
+            
+            for m in range(1, NLoops):
+                outputs_m = outputs[m - 1]
+                GC_collect[m, :, :, :] = outputs_m[0]
+                K_collect[m, :, :] = outputs_m[1]
+                if I0 > 0:
+                    B_collect[m, :, :] = outputs_m[2]
                 
         # --------------------------------------------------------------------------------------------------------------------------
         

PythonicDISORT/pydisort.py

@@ -24,7 +24,8 @@ def pydisort(
     NT_cor=False,
     Leg_coeffs_BDRF=np.array([]),
     s_poly_coeffs=np.array([[]]),
-    use_sparse_NLayers=6
+    use_sparse_NLayers=6,
+    n_jobs=1
 ):
     """Solves the 1D RTE for the fluxes, and optionally intensity,
     of a multi-layer atmosphere with the specified optical properties, boundary conditions
@@ -73,6 +74,11 @@ def pydisort(
         Arrange coefficients from lowest order term to highest.
     use_sparse_NLayers : optional, int
         At or above how many atmospheric layers should SciPy's sparse matrix framework be used?
+    n_jobs: optional, int
+        Maximum number of concurrently running jobs during parallelization.
+        This is exactly the `n_jobs` argument in `joblib.Parallel`, see 
+        https://joblib.readthedocs.io/en/stable/generated/joblib.Parallel.html
+        and https://joblib.readthedocs.io/en/stable/parallel.html.
 
     Returns
     -------
@@ -224,7 +230,8 @@ def pydisort(
             Nscoeffs,
             scale_tau,
             False,
-            use_sparse_NLayers
+            use_sparse_NLayers,
+            n_jobs
         )
         
         # TMS correction
@@ -297,12 +304,12 @@ def pydisort(
 
                 def Contribution_from_layer(j):
                     contribution = np.zeros((NQuad, Ntau, Nphi))
-                    pos_contribution_bools = l < j
-                    neg_contribution_bools = l > j
-                    scaled_tau_l_pos = scaled_tau[None, pos_contribution_bools]
-                    scaled_tau_l_neg = scaled_tau[None, neg_contribution_bools]
-                    if np.any(pos_contribution_bools):
-                        contribution[:N, pos_contribution_bools, :] = (
+                    pos_contribution_mask = l < j
+                    neg_contribution_mask = l > j
+                    scaled_tau_l_pos = scaled_tau[None, pos_contribution_mask]
+                    scaled_tau_l_neg = scaled_tau[None, neg_contribution_mask]
+                    if np.any(pos_contribution_mask):
+                        contribution[:N, pos_contribution_mask, :] = (
                             mathscr_B_pos[:, [j], :]
                             * (
                                 np.exp(
@@ -317,8 +324,8 @@ def pydisort(
                                 )
                             )[:, :, None]
                         )
-                    if np.any(neg_contribution_bools):
-                        contribution[N:, neg_contribution_bools, :] = (
+                    if np.any(neg_contribution_mask):
+                        contribution[N:, neg_contribution_mask, :] = (
                             mathscr_B_neg[:, [j], :]
                             * (
                                 np.exp(
@@ -433,5 +440,6 @@ def pydisort(
             Nscoeffs,
             scale_tau,
             only_flux,
-            use_sparse_NLayers
+            use_sparse_NLayers,
+            n_jobs
         )

{PythonicDISORT-0.1.1.dist-info → PythonicDISORT-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PythonicDISORT
-Version: 0.1.1
+Version: 0.2.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/PythonicDISORT
@@ -11,11 +11,12 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 Requires-Dist: numpy (>=1.17.3)
 Requires-Dist: scipy (>=1.8.0)
+Provides-Extra: joblib
+Requires-Dist: joblib (>=1.0.0) ; extra == 'joblib'
 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: ipympl (>=0.8.8) ; extra == 'notebook_dependencies'
 Provides-Extra: pytest
 Requires-Dist: pytest (>=6.2.5) ; extra == 'pytest'
 
@@ -41,7 +42,7 @@ We highly recommend reading the non-optional parts of sections 1 and 2 before us
 
 ## PyTest
 
-Separate from the notebook verification tests, we used PyTest to recreate most of the test problems from Stamnes et. al.'s `disotest.f90`.
+Separate from the verification tests in the Jupyter Notebook, we used PyTest to recreate most of the test problems from Stamnes et. al.'s `disotest.f90`.
 With PyTest installed, execute the console command `pytest` in the `pydisotest` directory to run these tests.
 
 # Installation
@@ -54,13 +55,13 @@ With PyTest installed, execute the console command `pytest` in the `pydisotest`
 * Python 3.8+
 * `numpy >= 1.17.3`
 * `scipy >= 1.8.0`
-* (OPTIONAL) `pytest >= 6.2.5`
+* (OPTIONAL) `joblib >= 1.0.0` (Required for parallelization)
+* (OPTIONAL) `pytest >= 6.2.5` (Required for non-Notebook tests)
 
 ## Additional requirements to run the Jupyter Notebook
 * `autograd >= 1.5`
 * `jupyter > 1.0.0`
 * `notebook > 6.5.2`
-* (OPTIONAL) `ipympl >= 0.8.8`
 
 In addition, our F2PY-wrapped Stamnes' DISORT (in the `disort4.0.99_f2py` directory) must be set up to run the last section (section 6).
 

PythonicDISORT-0.2.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
+PythonicDISORT/_loop_and_assemble_results.py,sha256=qzgVrVULUhFVZx2wXo3lgzmJmICoQ4RG7UrAVqhl1TQ,13382
+PythonicDISORT/_one_Fourier_mode.py,sha256=OWvrwL5T8I9Xc9bps1FZSCgn3RPV7PfdIHapdqqv5Rw,15022
+PythonicDISORT/pydisort.py,sha256=x7bgPV5Fueoi1d1j2_ULGytTOg4oeG7pmCluCHuZJMw,20502
+PythonicDISORT/subroutines.py,sha256=AyBZz24nuZaK89ShI1Z5-YmTUgNmNsp-JSV7VImngtE,10195
+PythonicDISORT-0.2.0.dist-info/METADATA,sha256=zvYRkIt99k7lX1eLwIa1jl9UHGe2nA0tLEbTCNPrb0o,5844
+PythonicDISORT-0.2.0.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
+PythonicDISORT-0.2.0.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
+PythonicDISORT-0.2.0.dist-info/RECORD,,

PythonicDISORT-0.1.1.dist-info/RECORD

@@ -1,9 +0,0 @@
-PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
-PythonicDISORT/_loop_and_assemble_results.py,sha256=nmTyxSmoraSED8SnWuczegr2rY9iNIZMiiGUael2bWQ,12306
-PythonicDISORT/_one_Fourier_mode.py,sha256=OWvrwL5T8I9Xc9bps1FZSCgn3RPV7PfdIHapdqqv5Rw,15022
-PythonicDISORT/pydisort.py,sha256=9J2GkRs7YoRt3m9i9CXUvDCq7zgx0NuMcHZfBmslvkU,20127
-PythonicDISORT/subroutines.py,sha256=AyBZz24nuZaK89ShI1Z5-YmTUgNmNsp-JSV7VImngtE,10195
-PythonicDISORT-0.1.1.dist-info/METADATA,sha256=WEF4IDFTkdb2lwI6beGPWOY5m6G20e1wZ_kV0rMpuPo,5756
-PythonicDISORT-0.1.1.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-PythonicDISORT-0.1.1.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
-PythonicDISORT-0.1.1.dist-info/RECORD,,