Skip to content

Commit

Permalink
Merge sulfur fixes and plotting updates. Downloads stellar spectra fr…
Browse files Browse the repository at this point in the history
…om OSF
  • Loading branch information
nichollsh committed Jul 18, 2024
1 parent 4bcb47f commit 2f8bdb6
Show file tree
Hide file tree
Showing 10 changed files with 66 additions and 193 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
## JANUS (temperature structure generator)
## JANUS (1D convective atmosphere model)

Generates a temperature profile using the generalised moist pseudoadiabat and a prescribed stratosphere. Calculates radiative fluxes using SOCRATES.

Expand Down
4 changes: 2 additions & 2 deletions examples/SocRadConv.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from importlib.resources import files

from janus.modules import RadConvEqm, plot_fluxes, plot_emission
from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, plot_adiabats, ReadBandEdges, StellarSpectrum
from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, DownloadStellarSpectra, plot_adiabats, ReadBandEdges, StellarSpectrum
import mors

####################################
Expand Down Expand Up @@ -70,7 +70,7 @@

#Download required spectral files
DownloadSpectralFiles("/Dayspring")
DownloadSpectralFiles("/stellar_spectra")
DownloadStellarSpectra()

# Move/prepare spectral file
print("Inserting stellar spectrum")
Expand Down
4 changes: 2 additions & 2 deletions examples/demo_instellation.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
import numpy as np

from janus.modules import MCPA_CBL
from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, ReadBandEdges, StellarSpectrum
from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, DownloadStellarSpectra, ReadBandEdges, StellarSpectrum

import mors

Expand All @@ -37,7 +37,7 @@

#Download required spectral files
DownloadSpectralFiles("/Oak")
DownloadSpectralFiles("/stellar_spectra")
DownloadStellarSpectra()

# Setup spectral file
print("Inserting stellar spectrum")
Expand Down
4 changes: 2 additions & 2 deletions examples/demo_runaway_greenhouse.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
from importlib.resources import files

from janus.modules import RadConvEqm
from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, ReadBandEdges, StellarSpectrum
from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, DownloadStellarSpectra, ReadBandEdges, StellarSpectrum
import mors

if __name__=='__main__':
Expand All @@ -36,7 +36,7 @@

#Download required spectral files
DownloadSpectralFiles("/Oak")
DownloadSpectralFiles("/stellar_spectra")
DownloadStellarSpectra()

# Setup spectral file
print("Inserting stellar spectrum")
Expand Down
168 changes: 9 additions & 159 deletions src/janus/modules/plot_flux_balance.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,180 +16,30 @@
from janus.modules.spectral_planck_surface import surf_Planck_nu
import janus.utils.GeneralAdiabat as ga # Moist adiabat with multiple condensibles


def plot_flux_balance(atm_dry, atm_moist, cp_dry, time, dirs):

fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(13,10))
# sns.set_style("ticks")
# sns.despine()

# Line settings
col_idx = 3
col_vol1 = "H2O"
col_vol2 = "N2"
col_vol3 = "H2"
col_vol4 = "O2"

# Temperature vs. pressure
ax1.semilogy(atm_moist.tmp,atm_moist.p, color=ga.vol_colors[col_vol1][col_idx+1], ls="-", label=r'Moist adiabat')
if cp_dry == True: ax1.semilogy(atm_dry.tmp,atm_dry.p, color=ga.vol_colors[col_vol3][col_idx+1], ls="-", label=r'Dry adiabat')
ax1.legend()
ax1.invert_yaxis()
ax1.set_xlabel(r'Temperature $T$ (K)')
ax1.set_ylabel(r'Pressure $P$ (Pa)')
# ax1.set_ylim(bottom=atm_moist.ps*1.01)
ax1.set_ylim(top=atm_moist.ptop, bottom=atm_moist.ps)

# Print active species
active_species = r""
for vol in atm_moist.vol_list:
if atm_moist.vol_list[vol] > 1e-5:
active_species = active_species + ga.vol_latex[vol] + ", "
active_species = active_species[:-2]
ax1.text(0.02, 0.02, r"Active species: "+active_species, va="bottom", ha="left", fontsize=10, transform=ax1.transAxes, bbox=dict(fc='white', ec="white", alpha=0.5, boxstyle='round', pad=0.1), color=ga.vol_colors["black_1"] )

# Fluxes vs. pressure

# Zero line
ax2.axvline(0, color=ga.vol_colors["qgray_light"], lw=0.5)
# ax2.axvline(-1e+3, color=ga.vol_colors["qgray_light"], lw=0.5)
# ax2.axvline(1e+3, color=ga.vol_colors["qgray_light"], lw=0.5)

# SW down / instellation flux
if cp_dry == True: ax2.semilogy(atm_dry.SW_flux_down*(-1),atm_dry.pl, color=ga.vol_colors[col_vol3][col_idx], ls=":")
ax2.semilogy(atm_moist.SW_flux_down*(-1),atm_moist.pl, color=ga.vol_colors[col_vol2][col_idx], ls=":", label=r'$F_{\odot}^{\downarrow}$')

# LW down / thermal downward flux
if cp_dry == True: ax2.semilogy(atm_dry.LW_flux_down*(-1),atm_dry.pl, color=ga.vol_colors[col_vol3][col_idx], ls="--")
ax2.semilogy(atm_moist.LW_flux_down*(-1),atm_moist.pl, color=ga.vol_colors[col_vol2][col_idx+1], ls="--", label=r'$F_\mathrm{t}^{\downarrow}$')
# ls=(0, (3, 1, 1, 1))

# Thermal upward flux, total
if cp_dry == True: ax2.semilogy(atm_dry.flux_up_total,atm_dry.pl, color=ga.vol_colors[col_vol3][col_idx], ls="--")
ax2.semilogy(atm_moist.flux_up_total,atm_moist.pl, color=ga.vol_colors[col_vol1][col_idx], ls="--", label=r'$F_\mathrm{t}^{\uparrow}$')

# Net flux
if cp_dry == True: ax2.semilogy(atm_dry.net_flux,atm_dry.pl, color=ga.vol_colors[col_vol3][6], ls="-", lw=2)
ax2.semilogy(atm_moist.net_flux,atm_moist.pl, color=ga.vol_colors[col_vol1][6], ls="-", lw=2, label=r'$F_\mathrm{net}^{\uparrow}$')

# # SW up
# if cp_dry == True: ax2.semilogy(atm_dry.SW_flux_up,atm_dry.pl, color=ga.vol_colors[col_vol3][col_idx], ls=":")
# ax2.semilogy(atm_moist.SW_flux_up,atm_moist.pl, color=ga.vol_colors[col_vol1][col_idx], ls=":", label=r'$F_\mathrm{SW}^{\uparrow}$')

# # LW up
# if cp_dry == True: ax2.semilogy(atm_dry.LW_flux_up,atm_dry.pl, color=ga.vol_colors[col_vol3][col_idx], ls=(0, (5, 1)))
# ax2.semilogy(atm_moist.LW_flux_up,atm_moist.pl, color=ga.vol_colors[col_vol1][col_idx], ls=(0, (5, 1)), label=r'$F_\mathrm{LW}^{\uparrow}$')

ax2.legend(ncol=6, fontsize=10, loc=3)
ax2.invert_yaxis()
ax2.set_xscale("symlog") # https://stackoverflow.com/questions/3305865/what-is-the-difference-between-log-and-symlog
ax2.set_xlabel(r'Outgoing flux $F^{\uparrow}$ (W m$^{-2}$)')
ax2.set_ylabel(r'Pressure $P$ (Pa)')
ax2.set_ylim(top=atm_moist.ptop, bottom=atm_moist.ps)

# Wavenumber vs. OLR
ax3.plot(atm_moist.band_centres, surf_Planck_nu(atm_moist)/atm_moist.band_widths, color="gray",ls='--',label=str(round(atm_moist.ts))+' K blackbody')
if cp_dry == True: ax3.plot(atm_dry.band_centres, atm_dry.net_spectral_flux[:,0]/atm_dry.band_widths, color=ga.vol_colors[col_vol3][col_idx])
ax3.plot(atm_moist.band_centres, atm_moist.net_spectral_flux[:,0]/atm_moist.band_widths, color=ga.vol_colors[col_vol1][col_idx+1])
ax3.set_ylabel(r'Spectral flux density (W m$^{-2}$ cm$^{-1}$)')
ax3.set_xlabel(r'Wavenumber (cm$^{-1}$)')
ax3.legend(loc=1)
ymax_plot = 1.2*np.max(atm_moist.net_spectral_flux[:,0]/atm_moist.band_widths)
ax3.set_ylim(bottom=0, top=ymax_plot)
ax3.set_xlim(left=0, right=np.max(np.where(atm_moist.net_spectral_flux[:,0]/atm_moist.band_widths > ymax_plot/1000., atm_moist.band_centres, 0.)))
# ax3.set_xlim(left=0, right=30000)


# Heating versus pressure
ax4.axvline(0, color=ga.vol_colors["qgray_light"], lw=0.5)

if cp_dry == True:
ax4.plot(atm_dry.LW_heating, atm_dry.p, ls="--", color=ga.vol_colors[col_vol3][col_idx+1])
ax4.plot(atm_dry.net_heating, atm_dry.p, lw=2, color=ga.vol_colors[col_vol3][col_idx+1])
ax4.plot(atm_dry.SW_heating, atm_dry.p, ls=":", color=ga.vol_colors[col_vol3][col_idx+1])

ax4.plot(atm_moist.LW_heating, atm_moist.p, ls="--", color=ga.vol_colors[col_vol1][col_idx+1], label=r'LW')
ax4.plot(atm_moist.net_heating, atm_moist.p, lw=2, color=ga.vol_colors[col_vol1][col_idx+1], label=r'Net')
ax4.plot(atm_moist.SW_heating, atm_moist.p, ls=":", color=ga.vol_colors[col_vol1][col_idx+1], label=r'SW')

# Plot tropopause
trpp_idx = int(atm_moist.trppidx)
if trpp_idx > 0:
ax4.axhline(atm_moist.pl[trpp_idx], color=ga.vol_colors[col_vol2][col_idx], lw=1.0, ls="-.", label=r'Tropopause')

ax4.invert_yaxis()
ax4.legend(ncol=4, fontsize=10, loc=3)
ax4.set_ylabel(r'Pressure $P$ (Pa)')
ax4.set_xlabel(r'Heating (K/day)')
# ax4.set_xscale("log")
ax4.set_yscale("log")
x_minmax = np.max([abs(np.min(atm_moist.net_heating[10:])), abs(np.max(atm_moist.net_heating[10:]))])
x_minmax = np.max([ 20, x_minmax ])
if not math.isnan(x_minmax):
ax4.set_xlim(left=-x_minmax, right=x_minmax)
ax4.set_ylim(top=atm_moist.ptop, bottom=atm_moist.ps)
# ax4.set_yticks([1e-10, 1e-5, 1e0, 1e5])
# ax4.set_xticks([0.1, 0.3, 1, 3, 10, 30, 100])
# ax4.set_xticklabels(["0.1", "0.3", "1", "3", "10", "30", "100"])

# # Wavelength versus OLR log plot
# OLR_cm_moist = atm_moist.LW_spectral_flux_up[:,0]/atm_moist.band_widths
# wavelength_moist = [ 1e+4/i for i in atm_moist.band_centres ] # microns
# OLR_micron_moist = [ 1e+4*i for i in OLR_cm_moist ] # microns
# if cp_dry == True:
# OLR_cm_dry = atm_dry.LW_spectral_flux_up[:,0]/atm_dry.band_widths
# wavelength_dry = [ 1e+4/i for i in atm_dry.band_centres ] # microns
# OLR_micron_dry = [ 1e+4*i for i in OLR_cm_dry ] # microns
# ax4.plot(wavelength_dry, OLR_micron_dry, color=ga.vol_colors[col_vol3][col_idx+1])

# ax4.plot(wavelength_moist, OLR_micron_moist, color=ga.vol_colors[col_vol1][col_idx+1])
# ax4.set_ylabel(r'Spectral flux density (W m$^{-2}$ $\mu$m$^{-1}$)')
# ax4.set_xlabel(r'Wavelength $\lambda$ ($\mu$m)')
# ax4.set_xscale("log")
# ax4.set_yscale("log")
# ax4.set_xlim(left=0.1, right=100)
# ax4.set_ylim(bottom=1e-20, top=1e5)
# # ax4.set_yticks([1e-10, 1e-5, 1e0, 1e5])
# ax4.set_xticks([0.1, 0.3, 1, 3, 10, 30, 100])
# ax4.set_xticklabels(["0.1", "0.3", "1", "3", "10", "30", "100"])

fig.savefig(dirs["output"]+"/"+"TP_"+str(round(time["planet"]))+'.pdf', bbox_inches="tight")
plt.close(fig)

# with open(dirs["output"]+"/"+str(int(time["planet"]))+"_atm.pkl", "wb") as atm_file:
# pkl.dump(atm, atm_file, protocol=pkl.HIGHEST_PROTOCOL)

# # Save atm object to .json file
# json_atm = json.dumps(atm.__dict__)
# with open(dirs["output"]+"/"+str(int(time_current))+"_atm.json", "wb") as atm_file:
# json.dump(json_atm, atm_file)


# Plotting
def plot_fluxes(atm,filename='output/fluxes.pdf'):

fig,ax = plt.subplots()

ax.axvline(0,color='black',lw=0.8)

# for p in np.logspace(-6, 3, 80)*1.0e5:
# ax.axhline(y=p, color='purple', lw=1, alpha=0.3)
arr_p = atm.pl*1e-5

ax.plot(atm.flux_up_total,atm.pl,color='red',label='UP',lw=1)
ax.plot(atm.SW_flux_up ,atm.pl,color='red',label='UP SW',linestyle='dotted',lw=2)
ax.plot(atm.LW_flux_up ,atm.pl,color='red',label='UP LW',linestyle='dashed',lw=1)
ax.plot(atm.flux_up_total,arr_p,color='red',label='UP',lw=1)
ax.plot(atm.SW_flux_up ,arr_p,color='red',label='UP SW',linestyle='dotted',lw=2)
ax.plot(atm.LW_flux_up ,arr_p,color='red',label='UP LW',linestyle='dashed',lw=1)

ax.plot(-1.0*atm.flux_down_total,atm.pl,color='green',label='DN',lw=2)
ax.plot(-1.0*atm.SW_flux_down ,atm.pl,color='green',label='DN SW',linestyle='dotted',lw=3)
ax.plot(-1.0*atm.LW_flux_down ,atm.pl,color='green',label='DN LW',linestyle='dashed',lw=2)
ax.plot(-1.0*atm.flux_down_total,arr_p,color='green',label='DN',lw=2)
ax.plot(-1.0*atm.SW_flux_down ,arr_p,color='green',label='DN SW',linestyle='dotted',lw=3)
ax.plot(-1.0*atm.LW_flux_down ,arr_p,color='green',label='DN LW',linestyle='dashed',lw=2)

ax.plot(atm.net_flux ,atm.pl,color='black',label='NET')
ax.plot(atm.net_flux ,arr_p,color='black',label='NET')

ax.set_xlabel("Upward-directed flux [W m-2]")
ax.set_xscale("symlog")
ax.set_ylabel("Pressure [Pa]")
ax.set_yscale("log")
ax.set_ylim([atm.pl[-1],atm.pl[0]])
ax.set_ylim(top=np.amin(arr_p)/1.5, bottom=np.amax(arr_p)*1.5)

ax.legend(loc='upper left')

Expand Down
26 changes: 14 additions & 12 deletions src/janus/utils/GeneralAdiabat.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,17 +33,19 @@
# https://chrisalbon.com/python/data_visualization/seaborn_color_palettes/
no_colors = 7
vol_colors = {
"H2O" : [mpl.colormaps["PuBu"](i) for i in np.linspace(0,1.0,no_colors)],
"CO2" : [mpl.colormaps["Reds"](i) for i in np.linspace(0,1.0,no_colors)],
"H2" : [mpl.colormaps["Greens"](i) for i in np.linspace(0,1.0,no_colors)],
"N2" : [mpl.colormaps["Purples"](i) for i in np.linspace(0,1.0,no_colors)],
"O2" : [mpl.colormaps["Wistia"](i) for i in np.linspace(0,1.0,no_colors+2)],
"O3" : [mpl.colormaps["spring"](i) for i in np.linspace(0,1.0,no_colors+2)],
"CH4" : [mpl.colormaps["RdPu"](i) for i in np.linspace(0,1.0,no_colors)],
"CO" : [mpl.colormaps["pink_r"](i) for i in np.linspace(0,1.0,no_colors)],
"S" : [mpl.colormaps["YlOrBr"](i) for i in np.linspace(0,1.0,no_colors)],
"He" : [mpl.colormaps["Greys"](i) for i in np.linspace(0,1.0,no_colors)],
"NH3" : [mpl.colormaps["summer"](i) for i in np.linspace(0,1.0,no_colors)],
"H2O": "#027FB1",
"CO2": "#D24901",
"H2" : "#008C01",
"CH4": "#C720DD",
"CO" : "#D1AC02",
"N2" : "#870036",
"S2" : "#FF8FA1",
"SO2": "#00008B",
"He" : "#30FF71",
"NH3": "#675200",
"mixtures" : [mpl.colormaps["Set3"](i) for i in np.linspace(0,1.0,no_colors)],
"H2O-CO2" : mpl.colormaps["Set3"](1.0/no_colors),
"CO2-H2O" : mpl.colormaps["Set3"](1.0/no_colors),
Expand Down Expand Up @@ -943,17 +945,17 @@ def plot_adiabats(atm,filename='output/general_adiabat.pdf'):

# Saturation vapor pressure for given temperature
Psat_array = [ p_sat(vol, T,water_lookup=atm.water_lookup) for T in T_sat_array ]
ax1.semilogy( T_sat_array, Psat_array, label=r'$p_\mathrm{sat}$'+vol_latex[vol], lw=ls_ind, ls=":", color=vol_colors[vol][4])
ax1.semilogy( T_sat_array, Psat_array, label=r'$p_\mathrm{sat}$'+vol_latex[vol], lw=ls_ind, ls=":", color=vol_colors[vol])

# Plot partial pressures
ax1.semilogy(atm.tmp, atm.p_vol[vol], color=vol_colors[vol][4], lw=ls_ind, ls="-", label=r'$p$'+vol_latex[vol],alpha=0.99)
ax1.semilogy(atm.tmp, atm.p_vol[vol], color=vol_colors[vol], lw=ls_ind, ls="-", label=r'$p$'+vol_latex[vol],alpha=0.99)

# Sum up partial pressures
p_partial_sum += atm.p_vol[vol]

# Plot individual molar concentrations
ax2.semilogy(atm.x_cond[vol],atm.p, color=vol_colors[vol][4], lw=ls_ind, ls="--", label=vol_latex[vol]+" cond.")
ax2.semilogy(atm.x_gas[vol],atm.p, color=vol_colors[vol][4], lw=ls_ind, ls="-", label=vol_latex[vol]+" gas")
ax2.semilogy(atm.x_cond[vol],atm.p, color=vol_colors[vol], lw=ls_ind, ls="--", label=vol_latex[vol]+" cond.")
ax2.semilogy(atm.x_gas[vol],atm.p, color=vol_colors[vol], lw=ls_ind, ls="-", label=vol_latex[vol]+" gas")

# # Plot sum of partial pressures as check
ax1.semilogy(atm.tmp, p_partial_sum, color="green", lw=ls_dry, ls="--", label=r'$\sum p^\mathrm{i}$',alpha=0.99)
Expand Down
39 changes: 29 additions & 10 deletions src/janus/utils/data.py
Original file line number Diff line number Diff line change
@@ -1,17 +1,13 @@
import os
from osfclient.api import OSF

#project ID of the stellar evolution tracks folder in the OSF
project_id = 'vehxg'

basic_list =[
"/Dayspring/256",
"/Frostflow/256",
"/Legacy",
"/Mallard",
"/Oak",
"/Reach",
"/stellar_spectra"
"/Reach"
]

def download_folder(storage, folder_name, local_path):
Expand All @@ -33,6 +29,31 @@ def download_folder(storage, folder_name, local_path):
file.write_to(local_file)
return


def GetFWLData():
fwl_data_dir = os.getenv('FWL_DATA')
if os.environ.get("FWL_DATA") == None:
raise Exception("The FWL_DATA environment variable where spectral data will be downloaded needs to be set up!")
return fwl_data_dir

def DownloadStellarSpectra():
#project ID of the stellar spectra on OSF
project_id = '8r2sw'

# Link with OSF project repository
osf = OSF()
project = osf.project(project_id)
storage = project.storage('osfstorage')

# Folder
data_dir = GetFWLData() + "/stellar_spectra"
if not os.path.exists(data_dir):
os.makedirs(data_dir)

# Get all named spectra
download_folder(storage,"/Named",data_dir)


def DownloadSpectralFiles(fname="",nband=256):
''''
Download spectral files data
Expand All @@ -44,13 +65,11 @@ def DownloadSpectralFiles(fname="",nband=256):
(only relevant for Dayspring, Frostflow and Honeyside)
'''

#Check if data environment variable is set up
fwl_data_dir = os.getenv('FWL_DATA')
if os.environ.get("FWL_DATA") == None:
raise Exception("The FWL_DATA environment variable where spectral data will be downloaded needs to be set up!")
#project ID of the spectral files on OSF
project_id = 'vehxg'

#Create spectral file data repository if not existing
data_dir = fwl_data_dir + "/spectral_files"
data_dir = GetFWLData() + "/spectral_files"
if not os.path.exists(data_dir):
os.makedirs(data_dir)

Expand Down
4 changes: 3 additions & 1 deletion src/janus/utils/phys.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,9 @@
"NO" : 0.0300061,
"NO2" : 0.0460055,
"HNO3": 0.06301284,
"OCS" : 0.060075
"OCS" : 0.060075,
"S2" : 0.0641, # kg mol-1

}

#This class allows convenient access
Expand Down
4 changes: 2 additions & 2 deletions tests/test_instellation.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import numpy as np

from janus.modules import MCPA_CBL
from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, ReadBandEdges, StellarSpectrum
from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, DownloadStellarSpectra, ReadBandEdges, StellarSpectrum
import mors

def test_instellation():
Expand All @@ -27,7 +27,7 @@ def test_instellation():

#Download required spectral files
DownloadSpectralFiles("/Oak")
DownloadSpectralFiles("/stellar_spectra")
DownloadStellarSpectra()

# Setup spectral file
print("Inserting stellar spectrum")
Expand Down
Loading

0 comments on commit 2f8bdb6

Please sign in to comment.