Save data to standard directory (#53)

* Set standard directory for data * Add CLI command, `janus env` to show env variables/data locations * Update documentation * Print destination * Uppercase variable
FormingWorlds · Aug 19, 2024 · d289631 · d289631
1 parent cff6ddf
commit d289631
Show file tree

Hide file tree

Showing 11 changed files with 100 additions and 70 deletions.
diff --git a/docs/getting_started.md b/docs/getting_started.md
@@ -27,6 +27,12 @@ janus download stellar
 
 ## Environment variables
 
+To see all environment variables and locations, use:
+
+```console
+janus env
+```
+
 ### `SOCRATES`
 
 By default, SOCRATES is installed to the default location based on the [XDG specification](https://specifications.freedesktop.org/basedir-spec/latest/).
@@ -37,3 +43,12 @@ you can override the path using the `SOCRATES` environment variable, e.g.
 ```console
 SOCRATES=/home/user/path/to/SOCRATES pytest
 ```
+
+### `FWL_DATA`
+
+Set this variable to modify where janus stores its stellar and spectral data. By default this is based on the [XDG specification](https://specifications.freedesktop.org/basedir-spec/latest/).
+You can override the path using the `FWL_DATA` environment variable, e.g.
+
+```console
+FWL_DATA=/home/user/path/to/fwl_data pytest
+```
diff --git a/examples/SocRadConv.py b/examples/SocRadConv.py
@@ -3,9 +3,9 @@
 """
 Created on Mon Jan 23 12:20:27 2023
 
-@authors: 
+@authors:
 Mark Hammond (MH)
-Tim Lichtenberg (TL)    
+Tim Lichtenberg (TL)
 Ryan Boukrouche (RB)
 Harrison Nicholls (HN)
 
@@ -24,6 +24,10 @@
 from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, DownloadStellarSpectra, plot_adiabats, ReadBandEdges, StellarSpectrum
 import mors
 
+import logging
+
+from janus.utils.data import FWL_DATA_DIR
+
 ####################################
 ##### Stand-alone initial conditions
 ####################################
@@ -34,8 +38,6 @@
     # Set up dirs
     if os.environ.get('RAD_DIR') == None:
         raise Exception("Socrates environment variables not set! Have you installed Socrates and sourced set_rad_env?")
-    if os.environ.get('FWL_DATA') == None:
-        raise Exception("The FWL_DATA environment variable where spectral and evolution tracks data will be downloaded needs to be set up!")
     dirs = {
             "janus": str(files("janus"))+"/",
             "output": os.path.abspath(os.getcwd())+"/output/"
@@ -45,13 +47,13 @@
     ##### Settings
     cfg_file =  dirs["janus"]+"data/tests/config_janus.toml"
     with open(cfg_file, 'r'):
-          cfg = toml.load(cfg_file)  
+          cfg = toml.load(cfg_file)
 
     # Planet
     time = { "planet": cfg['planet']['time'], "star": cfg['star']['time']}
     star_mass = cfg['star']['star_mass']
     mean_distance = cfg['star']['mean_distance']
- 
+
     # Define volatiles by partial pressures
     vol_mixing = {}
     vol_partial = {
@@ -74,16 +76,16 @@
 
     # Read spectrum
     spec = mors.Spectrum()
-    spec.LoadTSV(os.environ.get('FWL_DATA')+"/stellar_spectra/Named/sun.txt")
+    spec.LoadTSV(str(FWL_DATA_DIR / 'stellar_spectra' / 'Named' / 'sun.txt'))
 
-    # Convert to SOCRATES format 
+    # Convert to SOCRATES format
     socstar = os.path.join(dirs["output"], "socstar.txt")
     StellarSpectrum.PrepareStellarSpectrum(spec.wl, spec.fl, socstar)
 
     # Move/prepare spectral file
     print("Inserting stellar spectrum")
     StellarSpectrum.InsertStellarSpectrum(
-        os.environ.get('FWL_DATA')+"/spectral_files/Dayspring/256/Dayspring.sf",
+        str(FWL_DATA_DIR / 'spectral_files'/'Oak'/'318'/'Oak.sf'),
         socstar,
         dirs["output"]
     )
@@ -94,12 +96,12 @@
     atm = atmos.from_file(cfg_file, band_edges, vol_mixing=vol_mixing, vol_partial=vol_partial)
 
     # Set stellar heating on or off
-    if cfg['star']['stellar_heating'] == False: 
+    if cfg['star']['stellar_heating'] == False:
         atm.instellation = 0.
     else:
         mors.DownloadEvolutionTracks("/Baraffe")
         baraffe = mors.BaraffeTrack(star_mass)
-        atm.instellation = baraffe.BaraffeSolarConstant(time['star'], mean_distance) 
+        atm.instellation = baraffe.BaraffeSolarConstant(time['star'], mean_distance)
         print("Instellation:", round(atm.instellation), "W/m^2")
 
     # Set up atmosphere with general adiabat
@@ -131,4 +133,3 @@
 
     end = t.time()
     print("Runtime:", round(end - start,2), "s")
-
diff --git a/examples/demo_instellation.py b/examples/demo_instellation.py
@@ -16,15 +16,16 @@
 
 import mors
 
+from janus.utils.data import FWL_DATA_DIR
+
 if __name__=='__main__':
 
     print("Start")
 
     # Set up dirs
     if os.environ.get('RAD_DIR') == None:
         raise Exception("Socrates environment variables not set! Have you installed Socrates and sourced set_rad_env?")
-    if os.environ.get('FWL_DATA') == None:
-        raise Exception("The FWL_DATA environment variable where spectral and evolution tracks data will be downloaded needs to be set up!")
+
     dirs = {
             "janus": str(files("janus"))+"/",
             "output": os.path.abspath(os.getcwd())+"/output/"
@@ -41,17 +42,17 @@
 
     # Read spectrum
     spec = mors.Spectrum()
-    spec.LoadTSV(os.environ.get('FWL_DATA')+"/stellar_spectra/Named/sun.txt")
+    spec.LoadTSV(str(FWL_DATA_DIR / 'stellar_spectra' / 'Named' / 'sun.txt'))
 
-    # Convert to SOCRATES format 
+    # Convert to SOCRATES format
     socstar = os.path.join(dirs["output"], "socstar.txt")
     StellarSpectrum.PrepareStellarSpectrum(spec.wl, spec.fl, socstar)
 
 
     # Setup spectral file
     print("Inserting stellar spectrum")
     StellarSpectrum.InsertStellarSpectrum(
-        os.environ.get('FWL_DATA')+"/spectral_files/Oak/318/Oak.sf",
+        str(FWL_DATA_DIR / 'spectral_files'/'Oak'/'318'/'Oak.sf'),
         socstar,
         dirs["output"]
     )
@@ -86,7 +87,7 @@
     T_magma = atm.tmp_magma #get default value 3000 K but this could be a config option
 
     r_arr = np.linspace(0.3, 1.4, 7) # orbital distance range [AU]
-    
+
     asf_arr  = []   # ASF
     OLR_arr = []    # OLR
     net_arr = []    # net flux at TOA
@@ -95,7 +96,7 @@
     for i in range(7):
       print("Orbital separation = %.2f AU" % r_arr[i])
 
-      atm.instellation = baraffe.BaraffeSolarConstant(time['star'], r_arr[i]) 
+      atm.instellation = baraffe.BaraffeSolarConstant(time['star'], r_arr[i])
       atmos.setTropopauseTemperature(atm)
 
       atm = MCPA_CBL(dirs, atm, False, rscatter = True, T_surf_max=9.0e99, T_surf_guess = atm.trppT+100)
@@ -106,7 +107,7 @@
       ts_arr.append(atm.ts)
       tr_arr.append(atm.trppT)
 
-      # Plot case 
+      # Plot case
       plt.ioff()
       fig,ax = plt.subplots(1,1)
       ax.plot(atm.tmpl, atm.pl, color='black', lw=2)
@@ -123,10 +124,10 @@
       print(" ")
 
     save_arr = [r_arr, asf_arr, OLR_arr, net_arr, ts_arr, tr_arr]
-    np.savetxt(dirs["output"]+"/data_%dK.csv"%T_magma, 
+    np.savetxt(dirs["output"]+"/data_%dK.csv"%T_magma,
                np.array(save_arr).T, fmt="%.5e", delimiter=",",
                header="r [AU],  S_0 [W m-2], OLR [W m-2], net [W m-2], ts [K], tr[K] ")
-    
+
     print("Making plots")
 
     plt.ioff()
@@ -170,7 +171,7 @@
 
     arr_magma = np.ones(len(r_arr))*T_magma
     ax2.plot(r_arr, np.zeros(len(r_arr)), zorder=6, color='silver', lw=2.5, label=r"$\tilde{T}_s$") # magma temperature
-    ax2.plot(r_arr, ts_arr - arr_magma,   zorder=6, color='black',  lw=2.5, label=r"$T_s$") # surface solution temperature 
+    ax2.plot(r_arr, ts_arr - arr_magma,   zorder=6, color='black',  lw=2.5, label=r"$T_s$") # surface solution temperature
 
     ax2.legend(loc='center right', framealpha=1.0)
     ax2.set_ylabel(r"$T - \tilde{T_s}$ [K]")
@@ -180,12 +181,11 @@
     ax2.set_xlabel("Orbital separation [AU]")
     fig.subplots_adjust(hspace=0.08)
     fig.savefig(dirs["output"]+"inst_%dK.pdf"%T_magma, bbox_inches='tight')
-    
+
 
     # Tidy
     CleanOutputDir(os.getcwd())
     CleanOutputDir(dirs['output'])
 
     # Done
     print("Done!")
-
diff --git a/examples/demo_runaway_greenhouse.py b/examples/demo_runaway_greenhouse.py
@@ -14,6 +14,8 @@
 from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, DownloadStellarSpectra, ReadBandEdges, StellarSpectrum
 import mors
 
+from janus.utils.data import FWL_DATA_DIR
+
 if __name__=='__main__':
 
     print("Start")
@@ -22,8 +24,7 @@
     # Set up dirs
     if os.environ.get('RAD_DIR') == None:
         raise Exception("Socrates environment variables not set! Have you installed Socrates and sourced set_rad_env?")
-    if os.environ.get('FWL_DATA') == None:
-        raise Exception("The FWL_DATA environment variable where spectral and evolution tracks data will be downloaded needs to be set up!")
+
     dirs = {
             "janus": str(files("janus"))+"/",
             "output": os.path.abspath(os.getcwd())+"/output/"
@@ -40,17 +41,17 @@
 
     # Read spectrum
     spec = mors.Spectrum()
-    spec.LoadTSV(os.environ.get('FWL_DATA')+"/stellar_spectra/Named/sun.txt")
+    spec.LoadTSV(str(FWL_DATA_DIR / 'stellar_spectra' / 'Named' / 'sun.txt'))
 
-    # Convert to SOCRATES format 
+    # Convert to SOCRATES format
     socstar = os.path.join(dirs["output"], "socstar.txt")
     StellarSpectrum.PrepareStellarSpectrum(spec.wl, spec.fl, socstar)
 
 
     # Setup spectral file
     print("Inserting stellar spectrum")
     StellarSpectrum.InsertStellarSpectrum(
-        os.environ.get('FWL_DATA')+"/spectral_files/Oak/318/Oak.sf",
+        str(FWL_DATA_DIR / 'spectral_files'/'Oak'/'318'/'Oak.sf'),
         socstar,
         dirs["output"]
     )
@@ -78,10 +79,10 @@
     # Compute stellar heating
     mors.DownloadEvolutionTracks("/Baraffe")
     baraffe = mors.BaraffeTrack(star_mass)
-    atm.instellation = baraffe.BaraffeSolarConstant(time['star'], mean_distance) 
+    atm.instellation = baraffe.BaraffeSolarConstant(time['star'], mean_distance)
 
     #Run Janus
- 
+
     # Run JANUS in a loop to generate runaway curve
     print("Running JANUS...")
     Ts_arr = np.linspace(200, 2800, 20)
@@ -100,13 +101,13 @@
 
     # Get literature data
     g2013 = np.loadtxt(dirs["janus"]+"data/comparison_data/Goldblatt13_data.txt",
-                          dtype=float, skiprows=2, delimiter=',').T 
+                          dtype=float, skiprows=2, delimiter=',').T
     k2013 = np.loadtxt(dirs["janus"]+"data/comparison_data/Kopparapu13_data.txt",
-                          dtype=float, skiprows=2, delimiter=',').T 
+                          dtype=float, skiprows=2, delimiter=',').T
     h2015 = np.loadtxt(dirs["janus"]+"data/comparison_data/Hamano15_data.txt",
-                          dtype=float, skiprows=2, delimiter=',').T 
+                          dtype=float, skiprows=2, delimiter=',').T
     s2023 = np.loadtxt(dirs["janus"]+"data/comparison_data/Selsis23_convective.txt",
-                          dtype=float, skiprows=2, delimiter=',').T 
+                          dtype=float, skiprows=2, delimiter=',').T
 
     # Setup plot
     print("Making plot")
@@ -125,11 +126,11 @@
 
     ax.set_xlabel("Surface temperature [K]")
     ax.xaxis.set_minor_locator(MultipleLocator(100.0))
-    ax.set_xlim(200.0,  2700.0)  
+    ax.set_xlim(200.0,  2700.0)
 
     ax.set_ylabel("OLR [W m$^{-2}$]")
     ax.set_ylim(np.amin(OLR_arr) - 10.0, 500.0)
-    ax.yaxis.set_minor_locator(MultipleLocator(25.0))  
+    ax.yaxis.set_minor_locator(MultipleLocator(25.0))
 
     fig.savefig(dirs["output"]+"runaway_demo.pdf", bbox_inches='tight')
     fig.savefig(dirs["output"]+"runaway_demo.png", bbox_inches='tight', dpi=190)
@@ -141,4 +142,3 @@
 
     # Done
     print("Done!")
-
diff --git a/src/janus/cli.py b/src/janus/cli.py
@@ -42,10 +42,21 @@ def socrates(**kwargs):
     download_socrates(**kwargs)
 
 
+@click.command()
+def env():
+    """Show environment variables and locations"""
+    from janus.socrates import SOCRATES_DIR
+    from janus.utils.data import FWL_DATA_DIR
+
+    click.echo(f'RAD_DIR location: {SOCRATES_DIR}')
+    click.echo(f'FWL_DATA location: {FWL_DATA_DIR}')
+
+
 cli.add_command(download)
 download.add_command(spectral)
 download.add_command(stellar)
 download.add_command(socrates)
+cli.add_command(env)
 
 if __name__ == '__main__':
     cli()
diff --git a/src/janus/set_socrates_env.py b/src/janus/set_socrates_env.py
@@ -12,19 +12,22 @@
 import sys
 import zipfile
 from pathlib import Path
+import logging
 
 import click
 import platformdirs
 import requests
 
+logger = logging.getLogger(__name__)
+
 if not SOCRATES_DIR.exists():
     raise RuntimeError(f'Cannot find SOCRATES in this location: {SOCRATES_DIR}')
 
 with open(SOCRATES_DIR / 'version') as f:
-    version = f.readline()
+    SOCRATES_VERSION = f.readline()
 
-print(f'socrates location: {SOCRATES_DIR}')
-print(f'socrates version: {version}')
+logger.info(f'socrates location: %s', SOCRATES_DIR)
+logger.info('socrates version: %s', SOCRATES_VERSION)
 
 sep = os.pathsep