Save all profiles for PICASO 4

Author: aditya-sengupta

picaso/climate.py

@@ -12,8 +12,7 @@
 from .optics import compute_opacity
 from .disco import compress_thermal
 from .deq_chem import get_quench_levels
-
-import os
+from .io_utils import write_all_profiles
 
 from collections import namedtuple
 
@@ -72,7 +71,7 @@ def run_diseq_climate_workflow(bundle, nofczns, nstr, temp, pressure,
             rfaci,rfacv,tidal,
             Opagrid,
             CloudParameters,
-            save_profile,all_profiles,all_opd,
+            save_all_profiles,all_profiles,all_opd,
             verbose=True, moist = None,
             save_kzz=False, self_consistent_kzz=True):
     """ Can deprecate all of this since we moved to profile 
@@ -163,7 +162,7 @@ def run_diseq_climate_workflow(bundle, nofczns, nstr, temp, pressure,
             rfaci,rfacv,tidal,
             Opagrid,
             CloudParameters,
-            save_profile,all_profiles,all_opd,
+            save_all_profiles,all_profiles,all_opd,
             convergence_criteria, final,
             flux_net_ir_layer=None, flux_plus_ir_attop=None,first_call_ever=False,
             verbose=verbose, moist = moist,
@@ -177,7 +176,7 @@ def run_diseq_climate_workflow(bundle, nofczns, nstr, temp, pressure,
             rfaci, rfacv, tidal ,
             Opagrid,
             CloudParameters,
-            save_profile, all_profiles, all_opd,
+            save_all_profiles, all_profiles, all_opd,
             flux_net_ir_layer, flux_plus_ir_attop,
             verbose=verbose,  moist = moist,
             save_kzz=save_kzz,self_consistent_kzz=self_consistent_kzz,diseq=True, all_kzz=all_kzz)
@@ -196,7 +195,7 @@ def run_chemeq_climate_workflow(bundle, nofczns, nstr, temp, pressure,
             rfaci,rfacv,tidal,
             Opagrid,
             CloudParameters,
-            save_profile,all_profiles,all_opd,
+            save_all_profiles,all_profiles,all_opd,
             verbose=True, moist = None, 
             save_kzz = True, self_consistent_kzz=True): 
 
@@ -215,7 +214,7 @@ def run_chemeq_climate_workflow(bundle, nofczns, nstr, temp, pressure,
             rfaci, rfacv,  tidal, #energy balance 
             Opagrid, #delta_wno, tmin, tmax, 
             CloudParameters,#cloudy,cld_species,mh,fsed,beta,param_flag,mieff_dir ,opd_cld_climate,g0_cld_climate,w0_cld_climate, #scattering/cloud properties 
-            save_profile,all_profiles, all_opd,
+            save_all_profiles,all_profiles, all_opd,
             convergence_criteria, final , 
             first_call_ever=True, verbose=verbose, moist = moist,
             save_kzz=save_kzz,all_kzz=[],self_consistent_kzz=self_consistent_kzz)
@@ -238,11 +237,11 @@ def run_chemeq_climate_workflow(bundle, nofczns, nstr, temp, pressure,
             rfaci, rfacv,  tidal, #energy balance 
             Opagrid, #delta_wno, tmin, tmax, 
             CloudParameters,#cloudy,cld_species,mh,fsed,beta,param_flag,mieff_dir ,opd_cld_climate,g0_cld_climate,w0_cld_climate, #scattering/cloud properties 
-            save_profile,all_profiles, all_opd,               
+            save_all_profiles,all_profiles, all_opd,               
             convergence_criteria,final ,      
             flux_net_ir_layer=flux_net_ir_layer, flux_plus_ir_attop=flux_plus_ir_attop, 
             verbose=verbose,moist = moist,
-            save_kzz=True,all_kzz=all_kzz,self_consistent_kzz=self_consistent_kzz)   
+            save_kzz=True,all_kzz=all_kzz,self_consistent_kzz=self_consistent_kzz)  
 
     #STEP 3) find strat that will now run profile several times, each time updating the opacities and chemistry 
     #and also refine the convective zone guess while it does this. 
@@ -254,7 +253,7 @@ def run_chemeq_climate_workflow(bundle, nofczns, nstr, temp, pressure,
             rfaci, rfacv, tidal ,
             Opagrid,
             CloudParameters,
-            save_profile, all_profiles, all_opd,
+            save_all_profiles, all_profiles, all_opd,
             flux_net_ir_layer, flux_plus_ir_attop,
             verbose=verbose, moist = moist,self_consistent_kzz=self_consistent_kzz)
 
@@ -731,7 +730,7 @@ def t_start(nofczns,nstr,convergence_criteria,#
             rfaci, rfacv, tidal,
             Atmosphere, OpacityWEd, OpacityNoEd,ScatteringPhase, Disco,Opagrid, AdiabatBundle,
             F0PI,
-            save_profile, all_profiles, 
+            save_all_profiles, all_profiles, 
             fhole=None, hole_OpacityWEd=None,hole_OpacityNoEd=None, 
             verbose=1, do_holes=None, 
             moist = False, egp_stepmax = False):
@@ -777,7 +776,7 @@ def t_start(nofczns,nstr,convergence_criteria,#
         Any info for the adiabat calculations such as the precomputed tables from Didie 
     F0PI : ndarray
         Stellar spectrum if it exists otherwise this is just 1s array
-    save_profile : bool 
+    save_all_profiles : bool 
         bool to specify if all intermediate profiles will be saved 
     all_profiles : ndarray 
         All saved profiles if it is requested 
@@ -1454,7 +1453,7 @@ def t_start(nofczns,nstr,convergence_criteria,#
 
         if verbose: print("Iteration number ", its,", min , max temp ", min(temp),max(temp), ", flux balance ", flux_net[0]/abs(tidal[0])) #f/abs(tidal[0])**2) this other output here is slightly less straightforward with the square terms for exoplanets so making this just fnet/tidal for now
 
-        if save_profile == 1:
+        if save_all_profiles == 1:
             all_profiles = np.append(all_profiles,temp_old)
             cldsave_count += 1
         if flag_converge == 2 : # converged
@@ -2395,7 +2394,7 @@ def find_strat(bundle, nofczns,nstr,
         rfaci, rfacv, tidal ,
         Opagrid,
         CloudParameters,
-        save_profile, all_profiles, all_opd,
+        save_all_profiles, all_profiles, all_opd,
         flux_net_ir_layer, flux_plus_ir_attop,
         verbose=1, moist = None,
         save_kzz=False,self_consistent_kzz=True,diseq=False, all_kzz=[]):
@@ -2457,7 +2456,7 @@ def find_strat(bundle, nofczns,nstr,
             rfaci,rfacv,tidal,
             Opagrid,
             CloudParameters,
-            save_profile,all_profiles,all_opd,
+            save_all_profiles,all_profiles,all_opd,
             convergence_criteria, final,
             flux_net_ir_layer=flux_net_ir_layer, flux_plus_ir_attop=flux_plus_ir_attop,
             verbose=verbose,moist = moist,
@@ -2504,7 +2503,7 @@ def find_strat(bundle, nofczns,nstr,
             rfaci,rfacv,tidal,
             Opagrid,
             CloudParameters,
-            save_profile,all_profiles,all_opd,
+            save_all_profiles,all_profiles,all_opd,
             convergence_criteria, final,
             flux_net_ir_layer=flux_net_ir_layer, flux_plus_ir_attop=flux_plus_ir_attop,
             verbose=verbose,moist = moist,
@@ -2547,7 +2546,7 @@ def find_strat(bundle, nofczns,nstr,
                                 rfaci,rfacv,tidal,
                                 Opagrid,
                                 CloudParameters,
-                                save_profile,all_profiles,all_opd,
+                                save_all_profiles,all_profiles,all_opd,
                                 convergence_criteria, final,
                                 flux_net_ir_layer=flux_net_ir_layer, flux_plus_ir_attop=flux_plus_ir_attop,
                                 verbose=verbose, moist = moist,
@@ -2577,7 +2576,7 @@ def find_strat(bundle, nofczns,nstr,
                                 rfaci,rfacv,tidal,
                                 Opagrid,
                                 CloudParameters,
-                                save_profile,all_profiles,all_opd,
+                                save_all_profiles,all_profiles,all_opd,
                                 convergence_criteria, final,
                                 flux_net_ir_layer=flux_net_ir_layer, flux_plus_ir_attop=flux_plus_ir_attop,
                                 verbose=verbose, moist = moist,
@@ -2602,15 +2601,15 @@ def find_strat(bundle, nofczns,nstr,
                 rfaci,rfacv,tidal,
                 Opagrid,
                 CloudParameters,
-                save_profile,all_profiles,all_opd,
+                save_all_profiles,all_profiles,all_opd,
                 convergence_criteria, final,
                 flux_net_ir_layer=flux_net_ir_layer, flux_plus_ir_attop=flux_plus_ir_attop,
                 verbose=verbose,moist = moist,
             save_kzz=save_kzz,self_consistent_kzz=self_consistent_kzz,diseq=diseq, all_kzz=all_kzz)
                 #(mieff_dir, it_max_strat, itmx_strat, conv_strat, convt_strat, nofczns,nstr,x_max_mult,
                 #temp,pressure, F0PI, t_table, p_table, grad, cp,opacityclass, grav, 
                 #rfaci, rfacv, nlevel, tidal, tmin, tmax, dwni, bb , y2 , tp, final, 
-                #cloudy, cld_species,mh,fsed,flag_hack,save_profile, all_profiles, all_opd,
+                #cloudy, cld_species,mh,fsed,flag_hack,save_all_profiles, all_profiles, all_opd,
                 #opd_cld_climate,g0_cld_climate,w0_cld_climate,beta, param_flag,flux_net_ir_layer, 
                 #flux_plus_ir_attop, verbose=verbose,
             #fhole=fhole, fthin_cld=fthin_cld, do_holes = do_holes, moist = moist, cold_trap = cold_trap)
@@ -2630,7 +2629,7 @@ def find_strat(bundle, nofczns,nstr,
 
 
 def update_clouds(bundle, CloudParameters, Atmosphere, kzz,virga_kwargs,
-                   verbose=False,save_profile=True,all_opd=[]):
+                   verbose=False,save_all_profiles=True,all_opd=[]):
     """
     Updates cloud parameters and returns the cloud output.
     """
@@ -2669,7 +2668,7 @@ def update_clouds(bundle, CloudParameters, Atmosphere, kzz,virga_kwargs,
     taudif = np.max(np.abs(diff))
     taudif_tol = 0.4 * np.max(0.5 * (opd_clmt + opd_prev_cld_step))
 
-    if save_profile == 1:
+    if save_all_profiles == 1:
         all_opd = np.append(all_opd, df_cld['opd'].values[55::196])
 
     if verbose:
@@ -2684,7 +2683,7 @@ def profile(bundle, nofczns, nstr, temp, pressure,
             rfaci,rfacv,tidal,
             Opagrid,
             CloudParameters,
-            save_profile,all_profiles,all_opd,
+            save_all_profiles,all_profiles,all_opd,
             convergence_criteria, final,
             flux_net_ir_layer=None, flux_plus_ir_attop=None,first_call_ever=False,
             verbose=True, moist = None,
@@ -2762,7 +2761,7 @@ def profile(bundle, nofczns, nstr, temp, pressure,
                 temp[j1]= exp(log(temp[j1-1]) + grad_x*(log(pressure[j1]) - log(pressure[j1-1])))
 
     temp_old= np.copy(temp)
-    if save_profile == 1: all_profiles = np.append(all_profiles,temp_old)
+    if save_all_profiles == 1: all_profiles = np.append(all_profiles,temp_old)
 
     ### 1) ALWAYS UPDATE PT, CHEM, OPACITIES
     bundle.add_pt( temp, pressure)
@@ -2801,7 +2800,7 @@ def profile(bundle, nofczns, nstr, temp, pressure,
     ### 4) IF: COMPUTE CLOUDS 
     if cloudy :
         cld_out,df_cld, taudif, taudif_tol, all_opd, CloudParameters=update_clouds(bundle, CloudParameters,Atmosphere,
-                                                                          kz,virga_kwargs,save_profile=save_profile,
+                                                                          kz,virga_kwargs,save_all_profiles=save_all_profiles,
                                                                           all_opd=all_opd,verbose=verbose)
         bundle.clouds(df=df_cld,**hole_kwargs)
 
@@ -2818,7 +2817,7 @@ def profile(bundle, nofczns, nstr, temp, pressure,
     #    dill.dump([bundle, nofczns,nstr,convergence_criteria, rfaci, rfacv, tidal,
     #            Atmosphere, OpacityWEd, OpacityNoEd,ScatteringPhase, Disco,Opagrid, AdiabatBundle,
     #            F0PI,
-    #            save_profile, all_profiles, 
+    #            save_all_profiles, all_profiles, 
     #            verbose, moist , egp_stepmax ],file)
 
     ## begin bigger loop which gets opacities
@@ -2829,15 +2828,15 @@ def profile(bundle, nofczns, nstr, temp, pressure,
                 nofczns,nstr,convergence_criteria, rfaci, rfacv, tidal,
                 Atmosphere, OpacityWEd, OpacityNoEd,ScatteringPhase, Disco,Opagrid, AdiabatBundle,
                 F0PI,
-                save_profile, all_profiles, 
+                save_all_profiles, all_profiles, 
                 verbose=verbose, moist = moist, egp_stepmax = egp_stepmax, 
                 do_holes=do_holes, fhole=fhole, hole_OpacityWEd=OpacityWEd_clear,hole_OpacityNoEd=OpacityNoEd_clear)
         else:
             temp, dtdp, all_profiles,  flux_net_ir_layer,flux_net_v_layer, flux_plus_ir_attop = t_start(
                     nofczns,nstr,convergence_criteria, rfaci, rfacv, tidal,
                     Atmosphere, OpacityWEd, OpacityNoEd,ScatteringPhase, Disco,Opagrid, AdiabatBundle,
                     F0PI,
-                    save_profile, all_profiles, 
+                    save_all_profiles, all_profiles, 
                     verbose=verbose, moist = moist, egp_stepmax = egp_stepmax)
 
         ### 1) ALWAYS UPDATE PT, CHEM, OPACITIES
@@ -2880,13 +2879,13 @@ def profile(bundle, nofczns, nstr, temp, pressure,
         ### 4) IF: COMPUTE CLOUDS 
         if cloudy:
             cld_out,df_cld, taudif, taudif_tol, all_opd, CloudParameters=update_clouds(bundle, CloudParameters,Atmosphere,
-                                                                          kz,virga_kwargs,save_profile=save_profile,
+                                                                          kz,virga_kwargs,save_all_profiles=save_all_profiles,
                                                                           all_opd=all_opd,verbose=verbose)
             bundle.clouds(df=df_cld,**hole_kwargs)
         else: 
             cld_out=np.nan
 
-        if save_profile and cloudy:
+        if save_all_profiles and cloudy:
             all_opd = np.append(all_opd,df_cld['opd'].values[55::196]) #save opd at 4 micron
 
         ### 5) IF NEEDED: COMPUTE OPACITIES 
@@ -2928,6 +2927,7 @@ def profile(bundle, nofczns, nstr, temp, pressure,
 
         if verbose: print("Big iteration is ",min(temp), iii)
 
+    write_all_profiles(save_all_profiles, all_profiles)
     if conv_flag == 0:
         if verbose: print("Not converged")
     else :

picaso/io_utils.py

@@ -1,4 +1,3 @@
-import os 
 import json 
 import pandas as pd
 import warnings 

picaso/justdoit.py

@@ -4743,10 +4743,6 @@ def climate(self, opacityclass, save_all_profiles = False, with_spec=False,
 
         all_profiles= []
         all_opd = []
-        if save_all_profiles:
-            save_profile = 1
-        else :
-            save_profile = 0
 
         #initial guess 
         pressure = self.inputs['climate']['pressure']
@@ -4829,7 +4825,7 @@ def climate(self, opacityclass, save_all_profiles = False, with_spec=False,
                     rfaci, rfacv,  tidal, #energy balance 
                     Opagrid, #delta_wno, tmin, tmax, 
                     CloudParameters,#cloudy,cld_species,mh,fsed,beta,param_flag,mieff_dir ,opd_cld_climate,g0_cld_climate,w0_cld_climate, #scattering/cloud properties 
-                    save_profile,all_profiles, all_opd,
+                    save_all_profiles, all_profiles, all_opd,
                     verbose=verbose, moist = moist,
                     save_kzz=save_all_kzz, self_consistent_kzz=self_consistent_kzz)
 
@@ -4842,7 +4838,7 @@ def climate(self, opacityclass, save_all_profiles = False, with_spec=False,
                         rfaci,rfacv,tidal,
                         Opagrid,
                         CloudParameters,
-                        save_profile,all_profiles,all_opd,
+                        save_all_profiles, all_profiles,all_opd,
                         verbose=verbose, moist = moist, 
                         save_kzz=save_all_kzz, self_consistent_kzz=self_consistent_kzz)
         #all output to user