Add the option to save all temperature profiles in a climate run

.gitignore

@@ -10,6 +10,7 @@
 *.swo
 *.csv
 *.npy
+*.h5
 
 # Jupyter Notebooks
 .ipynb_checkpoints

HISTORY.rst

@@ -2,6 +2,9 @@
 
 History
 -------
+3.2 (2024-7-8)
+~~~~~~~~~~~~~~
+* Disequilibrium climate modeling
 
 3.1 (2023-3-23)
 ~~~~~~~~~~~~~~~

picaso/climate.py

@@ -327,7 +327,7 @@ def lu_backsubs(a, n, ntot, indx, b):
     return b
 
 # @logger.catch # Add this to track errors
-@jit(nopython=True, cache=True)
+@jit(nopython=True,cache=True)
 def t_start(nofczns,nstr,it_max,conv,x_max_mult, 
             rfaci, rfacv, nlevel, temp, pressure, p_table, t_table, 
             grad, cp, tidal, tmin,tmax, dwni , bb , y2, tp, DTAU, TAU, W0, COSB, 
@@ -971,7 +971,7 @@ def t_start(nofczns,nstr,it_max,conv,x_max_mult,
         if verbose: print("Iteration number ", its,", min , max temp ", min(temp),max(temp), ", flux balance ", flux_net[0]/abs(tidal[0]))
 
         if save_profile == 1:
-            all_profiles = np.append(all_profiles,temp_old)
+            all_profiles = np.vstack((all_profiles, temp_old.reshape((1,len(temp_old)))))
         if flag_converge == 2 : # converged
             # calculate  lapse rate
             dtdp=np.zeros(shape=(nlevel-1))

picaso/justdoit.py

@@ -38,19 +38,20 @@
 import math
 import xarray as xr
 from joblib import Parallel, delayed, cpu_count
+import h5py
 
 # #testing error tracker
 # from loguru import logger 
 __refdata__ = os.environ.get('picaso_refdata')
-__version__ = 3.2
+__version__ = '3.2.2'
 
 
 if not os.path.exists(__refdata__): 
     raise Exception("You have not downloaded the PICASO reference data. You can find it on github here: https://github.com/natashabatalha/picaso/tree/master/reference . If you think you have already downloaded it then you likely just need to set your environment variable. See instructions here: https://natashabatalha.github.io/picaso/installation.html#download-and-link-reference-documentation . You can use `os.environ['PYSYN_CDBS']=<yourpath>` directly in python if you run the line of code before you import PICASO.")
 else: 
     ref_v = json.load(open(os.path.join(__refdata__,'config.json'))).get('version',2.3)
 
-    if __version__ > ref_v: 
+    if __version__ != str(ref_v): 
         warnings.warn(f"Your code version is {__version__} but your reference data version is {ref_v}. For some functionality you may experience Keyword errors. Please download the newest ref version or update your code: https://github.com/natashabatalha/picaso/tree/master/reference")
 
 
@@ -3660,8 +3661,9 @@ def climate(self, opacityclass, save_all_profiles = False, as_dict=True,with_spe
         -----------
         opacityclass : class
             Opacity class from `justdoit.opannection`
-        save_all_profiles : bool
-            If you want to save and return all iterations in the T(P) profile,True/False
+        save_all_profiles : bool or str
+            If you want to save and return all iterations in the T(P) profile, True/False.
+            If str, specifies a path to which all iterations are written as an HDF5 file.
         with_spec : bool 
             Runs picaso spectrum at the end to get the full converged outputs, Default=False
         save_all_kzz : bool
@@ -3723,14 +3725,11 @@ def climate(self, opacityclass, save_all_profiles = False, as_dict=True,with_spe
         if rfacv==0:compute_reflected=False
         else:compute_reflected=True
 
-        all_profiles= []
-        if save_all_profiles:
-            save_profile = 1
-        else :
-            save_profile = 0
+        save_profile = bool(save_all_profiles)
 
         TEMP1 = self.inputs['climate']['guess_temp']
-        all_profiles=np.append(all_profiles,TEMP1)
+        all_profiles = np.empty((1,len(TEMP1)))
+        all_profiles[0,:] = TEMP1
         pressure = self.inputs['climate']['pressure']
         t_table = self.inputs['climate']['t_table']
         p_table = self.inputs['climate']['p_table']
@@ -4088,6 +4087,16 @@ def climate(self, opacityclass, save_all_profiles = False, as_dict=True,with_spe
             df_cld = vj.picaso_format(opd_now, w0_now, g0_now, pressure = cld_out['pressure'], wavenumber=1e4/cld_out['wave'])
             all_out['cld_output_picaso'] = df_cld
             all_out['virga_output'] = cld_out
+        if isinstance(save_all_profiles, str):
+            if verbose:
+                print(f"Saving all T(P) profiles to {save_all_profiles}")
+            profiles_file = h5py.File(save_all_profiles, 'w')
+            gp = profiles_file.create_group("all_profiles")
+            num_digits = len(str(len(all_profiles)))
+            for (i, profile_to_save) in enumerate(all_profiles):
+                i_str = str(i).zfill(num_digits)
+                gp.create_dataset(f"pt_{i_str}", data=profile_to_save)
+            profiles_file.close()
         if as_dict: 
             return all_out
         else: 
@@ -4581,7 +4590,7 @@ def profile(mieff_dir, it_max, itmx, conv, convt, nofczns,nstr,x_max_mult,
             temp[j1]= exp(log(temp[j1-1]) + grad_x*(log(pressure[j1]) - log(pressure[j1-1])))
 
     temp_old= np.copy(temp)
-
+    
 
 
     bundle = inputs(calculation='brown')
@@ -4591,7 +4600,8 @@ def profile(mieff_dir, it_max, itmx, conv, convt, nofczns,nstr,x_max_mult,
 
     bundle.premix_atmosphere(opacityclass, df = bundle.inputs['atmosphere']['profile'].loc[:,['pressure','temperature']])
     if save_profile == 1:
-            all_profiles = np.append(all_profiles,temp_old)
+        all_profiles = np.vstack((all_profiles, temp_old.reshape((1,len(temp_old)))))
+
 
     if first_call_ever == False:
         if cloudy == 1 :
@@ -4656,7 +4666,6 @@ def profile(mieff_dir, it_max, itmx, conv, convt, nofczns,nstr,x_max_mult,
 
     ## begin bigger loop which gets opacities
     for iii in range(itmx):
-
         temp, dtdp, flag_converge, flux_net_ir_layer, flux_plus_ir_attop, all_profiles = t_start(
                     nofczns,nstr,it_max,conv,x_max_mult, 
                     rfaci, rfacv, nlevel, temp, pressure, p_table, t_table, 

picaso/justplotit.py

@@ -1927,12 +1927,12 @@ def animate_convergence(clima_out, picaso_bundle, opacity, calculation='thermal'
                 np.copy(clima_out['all_profiles']))
 
     nlevel = len(t_eq)
-    mols_to_plot = {i:np.zeros(len(all_profiles_eq)) for i in molecules}
-    spec = np.zeros(shape =(int(len(all_profiles_eq)/nlevel),opacity.nwno))
+    nstep = all_profiles_eq.shape[0]
+    mols_to_plot = {i:np.zeros(all_profiles_eq.size) for i in molecules}
+    spec = np.zeros(shape =(nstep,opacity.nwno))
 
-    for i in range(int(len(all_profiles_eq)/nlevel)):
-
-        picaso_bundle.add_pt(all_profiles_eq[i*nlevel:(i+1)*nlevel], 
+    for i in range(nstep):
+        picaso_bundle.add_pt(all_profiles_eq[i], 
                              p_eq)
 
         picaso_bundle.premix_atmosphere(opacity,picaso_bundle.inputs['atmosphere']['profile'])
@@ -1963,7 +1963,7 @@ def animate_convergence(clima_out, picaso_bundle, opacity, calculation='thermal'
             # set the width ratios between the columns
             "width_ratios": [1,1,0.1,1,1,0.1,1,1]})
 
-    temp = all_profiles_eq[0*nlevel:(0+1)*nlevel]
+    temp = all_profiles_eq[0]
     lines = {}
     for imol,col in zip(molecules,Colorblind8):
         lines[imol], = ax['B'].loglog(mols_to_plot[imol][0:nlevel], p_eq,linewidth=3,color=col, label=imol)
@@ -2007,15 +2007,15 @@ def init():
         return lines
 
     def animate(i):                       
-        lines['temp'].set_xdata(all_profiles_eq[i*nlevel:(i+1)*nlevel])
+        lines['temp'].set_xdata(all_profiles_eq[i])
 
         for imol in molecules:
             lines[imol].set_xdata(mols_to_plot[imol][i*nlevel:(i+1)*nlevel])
 
         lines['spec'].set_ydata(spec[i,:][wh])
         return lines
 
-    ani = animation.FuncAnimation(fig, animate, frames=int(len(all_profiles_eq)/nlevel),init_func=init,interval=50, blit=False)
+    ani = animation.FuncAnimation(fig, animate, frames=nstep,init_func=init,interval=50, blit=False)
     plt.close()
     return ani
 

reference/config.json

@@ -1,5 +1,5 @@
 {
-	"version":3.2,
+	"version":"3.2.2",
 	"calculation":"spectrum_only",
 	"phase":0,
 	"planet":{

reference/version.md

@@ -1,4 +1,8 @@
-Version 3.1 (current beta)
+Version 3.2
+-----------
+- Disequilibrium Climate (Mukherjee et al. 2024) Elf-OWL models
+
+Version 3.1 
 --------------------------
 - Spherical harmonics with reflected light (Rooney et al. 2023a)
 - Spherical harmonics with thermal emission (Rooney et al. 2023b) 

setup.py

@@ -41,7 +41,7 @@
 # to this sample package.
 setup(
     name='picaso', 
-    version = '3.2',
+    version = '3.2.2',
     description = 'planetary intesity code for atmospheric scattering observations',
     long_description = 'README.md',
     author = 'Natasha E. Batalha',