From 021a172672cb922477e0d1e3066fd1da049d4247 Mon Sep 17 00:00:00 2001 From: Stef Smeets Date: Thu, 1 Aug 2024 09:03:01 +0200 Subject: [PATCH] Refactor and speed up tests (#41) * Add fwl_data to .gitignore * Clean test files * Move env check to conftest * Avoid updating class variable * Clean paths and use tmp directory for output * Avoid running the same tests twice in a PR * Refactos greenhous test * Refactor instellation test * Clean imports and run code in tmpdir * Set absolute path for FWL_DATA * Fix download * Parametrize tests * Limit number of tests * Update tests/test_runaway_greenhouse.py Co-authored-by: Laurent Soucasse <135331272+lsoucasse@users.noreply.github.com> * Update tests/test_instellation.py Co-authored-by: Laurent Soucasse <135331272+lsoucasse@users.noreply.github.com> --------- Co-authored-by: Laurent Soucasse <135331272+lsoucasse@users.noreply.github.com> --- .gitignore | 1 + src/janus/data/tests/config_instellation.toml | 1 + src/janus/data/tests/data_instellation.csv | 8 - .../data/tests/data_runaway_greenhouse.csv | 21 --- src/janus/utils/data.py | 2 +- tests/conftest.py | 10 ++ tests/helpers/__init__.py | 73 +++++++++ tests/test_instellation.py | 118 ++++----------- tests/test_runaway_greenhouse.py | 142 +++++++----------- 9 files changed, 170 insertions(+), 206 deletions(-) delete mode 100644 src/janus/data/tests/data_instellation.csv delete mode 100644 src/janus/data/tests/data_runaway_greenhouse.csv create mode 100644 tests/conftest.py create mode 100644 tests/helpers/__init__.py diff --git a/.gitignore b/.gitignore index cc1d9a7..1d8e60d 100644 --- a/.gitignore +++ b/.gitignore @@ -29,6 +29,7 @@ src/janus/data/spectral_files/sp_b318_HITRAN_a16_2203 output/ utils/*.ipynb* +.coverage build/ fwl_data*/ dist/ diff --git a/src/janus/data/tests/config_instellation.toml b/src/janus/data/tests/config_instellation.toml index c095d71..68f06ea 100644 --- a/src/janus/data/tests/config_instellation.toml +++ b/src/janus/data/tests/config_instellation.toml @@ -6,6 +6,7 @@ pl_mass = 5.972e24 [star] time = 100e+6 # yr star_mass = 1.0 # M_sun, mass of star +mean_distance = 0.3 # au, orbital distance [atmos] T_surf = 2800.0 diff --git a/src/janus/data/tests/data_instellation.csv b/src/janus/data/tests/data_instellation.csv deleted file mode 100644 index c4aea42..0000000 --- a/src/janus/data/tests/data_instellation.csv +++ /dev/null @@ -1,8 +0,0 @@ -# r [AU], S_0 [W m-2], OLR [W m-2], net [W m-2], ts [K], tr[K] -3.00000e-01,2.34297e+03,1.94397e+03,-4.67185e+01,3.00023e+03,4.19568e+02 -4.83333e-01,9.02641e+02,8.43916e+02,7.70272e+01,2.99961e+03,3.30552e+02 -6.66667e-01,4.74451e+02,5.01731e+02,9.86425e+01,2.99951e+03,2.81454e+02 -8.50000e-01,2.91857e+02,4.21028e+02,1.73069e+02,2.99913e+03,2.49260e+02 -1.03333e+00,1.97482e+02,4.16319e+02,2.48540e+02,2.99876e+03,2.26070e+02 -1.21667e+00,1.42451e+02,4.15915e+02,2.94890e+02,2.99853e+03,2.08342e+02 -1.40000e+00,1.07585e+02,4.15768e+02,3.24365e+02,2.99838e+03,1.94222e+02 diff --git a/src/janus/data/tests/data_runaway_greenhouse.csv b/src/janus/data/tests/data_runaway_greenhouse.csv deleted file mode 100644 index 724e9e2..0000000 --- a/src/janus/data/tests/data_runaway_greenhouse.csv +++ /dev/null @@ -1,21 +0,0 @@ -# T_surf [K], OLR [W m-2] - 200.00000000,9.07314e+01 - 336.84210526,2.77416e+02 - 473.68421053,2.77349e+02 - 610.52631579,2.77138e+02 - 747.36842105,2.77217e+02 - 884.21052632,2.77049e+02 -1021.05263158,2.77241e+02 -1157.89473684,2.76997e+02 -1294.73684211,2.77013e+02 -1431.57894737,2.77166e+02 -1568.42105263,2.77504e+02 -1705.26315789,2.78884e+02 -1842.10526316,2.85329e+02 -1978.94736842,3.07788e+02 -2115.78947368,3.74184e+02 -2252.63157895,5.47014e+02 -2389.47368421,9.50390e+02 -2526.31578947,1.80287e+03 -2663.15789474,3.49537e+03 -2800.00000000,6.58174e+03 diff --git a/src/janus/utils/data.py b/src/janus/utils/data.py index df77967..6c5c7f0 100644 --- a/src/janus/utils/data.py +++ b/src/janus/utils/data.py @@ -92,7 +92,7 @@ def DownloadSpectralFiles(fname: str="",nband: int=256): elif fname in ("Kynesgrove","Legacy","Mallard","Oak","Reach","stellar_spectra"): folder_list = [fname] else: - print(f"Unrecognised folder name: {fname}") + raise ValueError(f"Unrecognised folder name: {fname}") folders = [folder for folder in folder_list if not (data_dir / folder).exists()] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..cf8fa79 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,10 @@ +import os + +if not os.environ.get('RAD_DIR'): + raise Exception( + 'Socrates environment variables not set! Have you installed Socrates and sourced set_rad_env?' + ) +if not os.environ.get('FWL_DATA'): + raise Exception( + 'The FWL_DATA environment variable where spectral and evolution tracks data will be downloaded needs to be set up!' + ) diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py new file mode 100644 index 0000000..32232c5 --- /dev/null +++ b/tests/helpers/__init__.py @@ -0,0 +1,73 @@ +import os +from contextlib import contextmanager +from importlib.resources import files +from pathlib import Path + +import mors +import toml +from janus.utils import ( + DownloadSpectralFiles, + DownloadStellarSpectra, + ReadBandEdges, + StellarSpectrum, + atmos, +) + +DATA_DRC = files('janus') / 'data' / 'tests' +FWL_DATA = os.environ.get('FWL_DATA') + + +@contextmanager +def work_directory(path: Path | str): + """Changes working directory and returns to previous on exit. + + Parameters + ---------- + path : Path | str + Temporarily change to this directory. + """ + prev_cwd = Path.cwd().resolve() + try: + os.chdir(path) + yield + finally: # In any case, no matter what happens, go back eventually + os.chdir(prev_cwd) + + +def get_spectrum_data(drc): + DownloadSpectralFiles('Oak') + DownloadStellarSpectra() + + spec = mors.Spectrum() + spec.LoadTSV(FWL_DATA + '/stellar_spectra/Named/sun.txt') + socstar = os.path.join(drc, 'socstar.txt') + StellarSpectrum.PrepareStellarSpectrum(spec.wl, spec.fl, socstar) + + StellarSpectrum.InsertStellarSpectrum( + FWL_DATA + '/spectral_files/Oak/318/Oak.sf', + socstar, + drc, + ) + band_edges = ReadBandEdges(drc + 'star.sf') + + return band_edges + + +def get_atmosphere_config(*, band_edges, cfg_name: str, distance: float): + cfg_file = DATA_DRC / cfg_name + + with open(cfg_file) as f: + cfg = toml.load(f) + + star_time = cfg['star']['time'] + star_mass = cfg['star']['star_mass'] + + vol_mixing = {'H2O': 1.0, 'CO2': 0.0, 'N2': 0.0} + + atm = atmos.from_file(cfg_file, band_edges, vol_mixing=vol_mixing, vol_partial={}) + + mors.DownloadEvolutionTracks('/Baraffe') + baraffe = mors.BaraffeTrack(star_mass) + atm.instellation = baraffe.BaraffeSolarConstant(star_time, distance) + + return atm, cfg diff --git a/tests/test_instellation.py b/tests/test_instellation.py index e092ebf..997e0b5 100755 --- a/tests/test_instellation.py +++ b/tests/test_instellation.py @@ -1,96 +1,42 @@ -#!/usr/bin/env python3 - -from importlib.resources import files -import os, shutil, toml -import numpy as np - +from numpy.testing import assert_allclose +from helpers import get_atmosphere_config, get_spectrum_data, work_directory from janus.modules import MCPA_CBL -from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, DownloadStellarSpectra, ReadBandEdges, StellarSpectrum -import mors - -def test_instellation(): - - # 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/" - } +import pytest - # Tidy directory - if os.path.exists(dirs["output"]): - shutil.rmtree(dirs["output"]) - os.mkdir(dirs["output"]) - #Download required spectral files - DownloadSpectralFiles("Oak") - DownloadStellarSpectra() +TEST_DATA = ( + (0.3, (2.34297e03, 1.94397e03, -4.67185e01, 3.00023e03, 4.19568e02)), +# (0.483333, (9.02641e02, 8.43916e02, 7.70272e01, 2.99961e03, 3.30552e02)), +# (0.666667, (4.74451e02, 5.01731e02, 9.86425e01, 2.99951e03, 2.81454e02)), +# (0.85, (2.91857e02, 4.21028e02, 1.73069e02, 2.99913e03, 2.49260e02)), +# (1.03333, (1.97482e02, 4.16319e02, 2.48540e02, 2.99876e03, 2.26070e02)), +# (1.21667, (1.42451e02, 4.15915e02, 2.94890e02, 2.99853e03, 2.08342e02)), + (1.4, (1.07585e02, 4.15768e02, 3.24365e02, 2.99838e03, 1.94222e02)), +) - # Read spectrum - spec = mors.Spectrum() - spec.LoadTSV(os.environ.get('FWL_DATA')+"/stellar_spectra/Named/sun.txt") - # Convert to SOCRATES format - socstar = os.path.join(dirs["output"], "socstar.txt") - StellarSpectrum.PrepareStellarSpectrum(spec.wl, spec.fl, socstar) +@pytest.mark.parametrize("inp,expected", TEST_DATA) +def test_instellation(tmpdir, inp, expected): + out_drc = str(tmpdir) + "/" - # Setup spectral file - print("Inserting stellar spectrum") - StellarSpectrum.InsertStellarSpectrum( - os.environ.get('FWL_DATA')+"/spectral_files/Oak/318/Oak.sf", - socstar, - dirs["output"] + band_edges = get_spectrum_data(out_drc) + atm, cfg = get_atmosphere_config( + band_edges=band_edges, + distance=inp, + cfg_name="config_instellation.toml", ) - band_edges = ReadBandEdges(dirs["output"]+"star.sf") - - # Open config file - cfg_file = dirs["janus"]+"data/tests/config_instellation.toml" - with open(cfg_file, 'r') as f: - cfg = toml.load(f) - - # Star luminosity - time = { "planet": cfg['planet']['time'], "star": cfg['star']['time']} - star_mass = cfg['star']['star_mass'] - mors.DownloadEvolutionTracks("/Baraffe") - baraffe = mors.BaraffeTrack(star_mass) - - # Define volatiles by mole fractions - vol_mixing = { - "H2O" : 1.0, - "CO2" : 0.0, - "N2" : 0.0 - } - - #Get reference data - ref = np.loadtxt(dirs["janus"]+"data/tests/data_instellation.csv", - dtype=float, skiprows=1, delimiter=',') - - # Create atmosphere object - atm = atmos.from_file(cfg_file, band_edges, vol_mixing=vol_mixing, vol_partial={}) - - r_arr = np.linspace(0.3, 1.4, 7) # orbital distance range [AU] - for i in range(7): - print("Orbital separation = %.2f AU" % 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) - - out = [atm.SW_flux_down[0], atm.LW_flux_up[0], atm.net_flux[0], atm.ts, atm.trppT] - - print_out = "%.5e,"%float(r_arr[i]) - print_out += ",".join(["%.5e"%o for o in out]) - print("Calculated:",print_out) + atm.setTropopauseTemperature() - # print_out = ",".join(["%.5e"%o for o in ref[i]]) - # print("Target:",print_out) + with work_directory(tmpdir): + atm = MCPA_CBL( + {"output": out_drc}, + atm, + False, + rscatter=True, + T_surf_max=9.0e99, + T_surf_guess=atm.trppT + 100, + ) - np.testing.assert_allclose(out, ref[i][1:6], rtol=1e-5, atol=0) + ret = (atm.SW_flux_down[0], atm.LW_flux_up[0], atm.net_flux[0], atm.ts, atm.trppT) - # Tidy - CleanOutputDir(os.getcwd()) - CleanOutputDir(dirs['output']) + assert_allclose(ret, expected, rtol=1e-5, atol=0) diff --git a/tests/test_runaway_greenhouse.py b/tests/test_runaway_greenhouse.py index 03eeebe..5816f4d 100755 --- a/tests/test_runaway_greenhouse.py +++ b/tests/test_runaway_greenhouse.py @@ -1,92 +1,54 @@ -#!/usr/bin/env python3 - -import os, shutil, toml import numpy as np -from matplotlib.ticker import MultipleLocator -from importlib.resources import files - +from helpers import get_atmosphere_config, get_spectrum_data, work_directory from janus.modules import RadConvEqm -from janus.utils import atmos, CleanOutputDir, DownloadSpectralFiles, DownloadStellarSpectra, ReadBandEdges, StellarSpectrum -import mors - -def test_runaway_greenhouse(): - - # 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/" - } - - # Tidy directory - if os.path.exists(dirs["output"]): - shutil.rmtree(dirs["output"]) - os.mkdir(dirs["output"]) - - #Download required spectral files - DownloadSpectralFiles("Oak") - DownloadStellarSpectra() - - # Read spectrum - spec = mors.Spectrum() - spec.LoadTSV(os.environ.get('FWL_DATA')+"/stellar_spectra/Named/sun.txt") - - # 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", - socstar, - dirs["output"] - ) - band_edges = ReadBandEdges(dirs["output"]+"star.sf") - - # Open config file - cfg_file = dirs["janus"]+"data/tests/config_runaway.toml" - with open(cfg_file, 'r') as f: - cfg = toml.load(f) - - # Planet - time = { "planet": cfg['planet']['time'], "star": cfg['star']['time']} - star_mass = cfg['star']['star_mass'] - mean_distance = cfg['star']['mean_distance'] - - vol_mixing = { - "H2O" : 1.0, - "CO2" : 0.0, - "N2" : 0.0 - } - - #Get reference values - OLR_ref = np.loadtxt(dirs["janus"]+"data/tests/data_runaway_greenhouse.csv", - dtype=float, skiprows=1, delimiter=',') - - # Create atmosphere object - atm = atmos.from_file(cfg_file, band_edges, vol_mixing=vol_mixing, vol_partial={}) - - # Compute stellar heating from Baraffe evolution tracks - mors.DownloadEvolutionTracks("/Baraffe") - baraffe = mors.BaraffeTrack(star_mass) - atm.instellation = baraffe.BaraffeSolarConstant(time['star'], mean_distance) - - #Run Janus - Ts_arr = np.linspace(200, 2800, 20) - for i in range(20): - - atmos.setSurfaceTemperature(atm, Ts_arr[i]) - - _, atm_moist = RadConvEqm(dirs, time, atm, standalone=True, cp_dry=False, trppD=False, rscatter=False) - - out = atm_moist.LW_flux_up[0] - print("Output %.5e; Reference %.5e" % (out, OLR_ref[i][1])) - np.testing.assert_allclose(out, OLR_ref[i][1], rtol=1e-5, atol=0) - - # Tidy - CleanOutputDir(os.getcwd()) - CleanOutputDir(dirs['output']) +import pytest + + +TEST_DATA = ( + (200.00000000,9.07314e+01), +# (336.84210526,2.77416e+02), +# (473.68421053,2.77349e+02), +# (610.52631579,2.77138e+02), +# (747.36842105,2.77217e+02), +# (884.21052632,2.77049e+02), +# (1021.05263158,2.77241e+02), +# (1157.89473684,2.76997e+02), +# (1294.73684211,2.77013e+02), +# (1431.57894737,2.77166e+02), +# (1568.42105263,2.77504e+02), + (1705.26315789,2.78884e+02), +# (1842.10526316,2.85329e+02), +# (1978.94736842,3.07788e+02), +# (2115.78947368,3.74184e+02), +# (2252.63157895,5.47014e+02), +# (2389.47368421,9.50390e+02), +# (2526.31578947,1.80287e+03), +# (2663.15789474,3.49537e+03), + (2800.00000000,6.58174e+03), +) + + +@pytest.mark.parametrize("inp,expected", TEST_DATA) +def test_runaway_greenhouse(tmpdir, inp, expected): + out_drc = str(tmpdir) + '/' + + band_edges = get_spectrum_data(out_drc) + atm, cfg = get_atmosphere_config(band_edges=band_edges, distance=0.3, cfg_name='config_runaway.toml') + + time = {'planet': cfg['planet']['time'], 'star': cfg['star']['time']} + + atm.setSurfaceTemperature(inp) + + with work_directory(tmpdir): + _, atm_moist = RadConvEqm( + {'output': out_drc}, + time, + atm, + standalone=True, + cp_dry=False, + trppD=False, + rscatter=False, + ) + + ret = atm_moist.LW_flux_up[0] + np.testing.assert_allclose(ret, expected, rtol=1e-5, atol=0)