diff --git a/input/chili/intercomp/_base.grid.toml b/input/chili/intercomp/_base.grid.toml index 2483e9835..dc7f06383 100644 --- a/input/chili/intercomp/_base.grid.toml +++ b/input/chili/intercomp/_base.grid.toml @@ -1,7 +1,7 @@ output = "chili_base_grid/" symlink = "" ref_config = "input/chili/intercomp/_base.toml" -use_slurm = false +use_slurm = true max_jobs = 9 # maximum number of concurrent tasks (e.g. 500 on Habrok) max_days = 1 # maximum number of days to run (e.g. 1) diff --git a/input/chili/intercomp/_base.toml b/input/chili/intercomp/_base.toml index 15adab058..ef195be23 100644 --- a/input/chili/intercomp/_base.toml +++ b/input/chili/intercomp/_base.toml @@ -6,8 +6,8 @@ version = "2.0" path = "chili_base" logging = "INFO" plot_fmt = "pdf" -write_mod = 2 -plot_mod = 30 +write_mod = 4 +plot_mod = 40 archive_mod = "none" remove_sf = true @@ -15,14 +15,14 @@ remove_sf = true starspec = 1e9 starinst = 10.0 method = "adaptive" -minimum = 1000.0 -minimum_rel = 1e-05 +minimum = 100.0 +minimum_rel = 5e-3 maximum = 10000000.0 -initial = 30.0 +initial = 10.0 [params.dt.adaptive] -atol = 0.02 -rtol = 0.09 +atol = 0.04 +rtol = 0.11 [params.stop] strict = false @@ -133,7 +133,7 @@ ls_default = 1 max_steps = 70 perturb_all = true mlt_criterion = "s" -fastchem_floor = 500.0 +fastchem_floor = 700.0 ini_profile = "isothermal" [atmos_clim.dummy] @@ -146,7 +146,7 @@ reservoir = "outgas" [escape.zephyrus] Pxuv = 5e-05 -efficiency = 0.1 +efficiency = 0.3 tidal = true [interior] diff --git a/src/proteus/atmos_clim/common.py b/src/proteus/atmos_clim/common.py index ea82060e0..5bbb3e1c0 100644 --- a/src/proteus/atmos_clim/common.py +++ b/src/proteus/atmos_clim/common.py @@ -40,12 +40,14 @@ def ncdf_flag_to_bool(var) -> bool: raise ValueError(f'Could not parse NetCDF atmos flag variable \n {var}') -def read_ncdf_profile(nc_fpath: str, extra_keys: list = []): +def read_ncdf_profile(nc_fpath: str, extra_keys: list = [], combine_edges: bool = True) -> dict: """Read data from atmosphere NetCDF output file. - Automatically reads pressure (p), temperature (t), radius (z) arrays with - cell-centre (N) and cell-edge (N+1) values interleaved into a single combined array of - length (2*N+1). + All variables in SI units, same as NetCDF file content. + + Automatically reads pressure (p), temperature (t), radius (z) arrays. + If `combine_edges` is True, cell-centre (N) and cell-edge (N+1) values + are interleaved into a single combined array of length (2*N+1). Extra keys can be read-in using the extra_keys parameter. These will be stored with the same dimensions as in the NetCDF file. @@ -54,9 +56,10 @@ def read_ncdf_profile(nc_fpath: str, extra_keys: list = []): ---------- nc_fpath : str Path to NetCDF file. - extra_keys : list List of extra keys (strings) to read from the file. + combine_edges : bool + Whether to combine cell-centre and cell-edge values into a single array. Returns ---------- @@ -96,22 +99,32 @@ def read_ncdf_profile(nc_fpath: str, extra_keys: list = []): # read pressure, temperature, height data into dictionary values out = {} - out['p'] = [pl[0]] - out['t'] = [tl[0]] - out['z'] = [zl[0]] - out['r'] = [rl[0]] - for i in range(nlev_c): - out['p'].append(p[i]) - out['p'].append(pl[i + 1]) - - out['t'].append(t[i]) - out['t'].append(tl[i + 1]) - - out['z'].append(z[i]) - out['z'].append(zl[i + 1]) - - out['r'].append(r[i]) - out['r'].append(rl[i + 1]) + if combine_edges: + out['p'] = [pl[0]] + out['t'] = [tl[0]] + out['z'] = [zl[0]] + out['r'] = [rl[0]] + for i in range(nlev_c): + out['p'].append(p[i]) + out['p'].append(pl[i + 1]) + + out['t'].append(t[i]) + out['t'].append(tl[i + 1]) + + out['z'].append(z[i]) + out['z'].append(zl[i + 1]) + + out['r'].append(r[i]) + out['r'].append(rl[i + 1]) + else: + out['p'] = p + out['t'] = t + out['z'] = z + out['r'] = r + out['pl'] = pl + out['tmpl'] = tl + out['zl'] = zl + out['rl'] = rl # flags for fk in ('transparent', 'solved', 'converged'): diff --git a/src/proteus/interior/timestep.py b/src/proteus/interior/timestep.py index d8684d907..dccec52ec 100644 --- a/src/proteus/interior/timestep.py +++ b/src/proteus/interior/timestep.py @@ -230,7 +230,7 @@ def next_step( # Min step size dtminimum = config.params.dt.minimum # absolute - dtminimum += config.params.dt.minimum_rel * hf_row['Time'] * 0.01 # allow small steps + dtminimum += config.params.dt.minimum_rel * hf_row['Time'] # allow small steps dtswitch = max(dtswitch, dtminimum) log.info('New time-step target is %.2e years' % dtswitch) diff --git a/tests/atmos_clim/test_common.py b/tests/atmos_clim/test_common.py index dae7a5617..eea7ca8b0 100644 --- a/tests/atmos_clim/test_common.py +++ b/tests/atmos_clim/test_common.py @@ -113,6 +113,48 @@ def test_read_ncdf_profile(mock_ds, mock_isfile): mock_ds.assert_called_with('dummy.nc') +@pytest.mark.unit +@patch('proteus.atmos_clim.common.os.path.isfile') +@patch('netCDF4.Dataset') +def test_read_ncdf_profile_without_combining_edges(mock_ds, mock_isfile): + """Read centre and edge arrays separately when ``combine_edges`` is false. + + This covers the branch added for callers that need native NetCDF layering + (N centre levels and N+1 edge levels as separate arrays). + """ + mock_isfile.return_value = True + + ds_instance = MagicMock() + mock_ds.return_value = ds_instance + + # Use JANUS-style height variables to ensure this branch also works with z/zl input. + ds_instance.variables = { + 'p': np.array([100.0, 80.0]), + 'pl': np.array([110.0, 90.0, 70.0]), + 'tmp': np.array([500.0, 450.0]), + 'tmpl': np.array([520.0, 470.0, 430.0]), + 'z': np.array([1.0e4, 2.0e4]), + 'zl': np.array([0.0, 1.5e4, 2.5e4]), + 'planet_radius': [6.0e6], + 'solved': np.array([b'n'], dtype='S1'), + } + + result = read_ncdf_profile('dummy.nc', combine_edges=False) + + np.testing.assert_allclose(result['p'], np.array([100.0, 80.0])) + np.testing.assert_allclose(result['pl'], np.array([110.0, 90.0, 70.0])) + np.testing.assert_allclose(result['t'], np.array([500.0, 450.0])) + np.testing.assert_allclose(result['tmpl'], np.array([520.0, 470.0, 430.0])) + + # JANUS path: r = z + planet_radius and rl = zl + planet_radius. + np.testing.assert_allclose(result['r'], np.array([6.01e6, 6.02e6])) + np.testing.assert_allclose(result['rl'], np.array([6.0e6, 6.015e6, 6.025e6])) + + assert result['solved'] == 0.0 + assert result['transparent'] == 0.0 + assert result['converged'] == 0.0 + + @pytest.mark.unit @patch('proteus.atmos_clim.common.read_ncdf_profile') def test_read_atmosphere_data(mock_read): diff --git a/tests/tools/test_chili_postproc.py b/tests/tools/test_chili_postproc.py new file mode 100644 index 000000000..52db8a08c --- /dev/null +++ b/tests/tools/test_chili_postproc.py @@ -0,0 +1,140 @@ +""" +Unit tests for CHILI postprocessing helper script. + +These tests cover the scalar postprocessing path in ``tools/chili_postproc.py`` +without invoking heavy atmosphere/profile readers or plotting. + +See also: +- docs/How-to/test_infrastructure.md +- docs/How-to/test_categorization.md +- docs/How-to/test_building.md +""" + +from __future__ import annotations + +import builtins +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest + + +def _load_chili_postproc_module(): + """Load ``tools/chili_postproc.py`` as a module for direct function tests.""" + repo_root = Path(__file__).resolve().parents[2] + script_path = repo_root / 'tools' / 'chili_postproc.py' + spec = importlib.util.spec_from_file_location('chili_postproc_under_test', script_path) + module = importlib.util.module_from_spec(spec) + assert spec is not None + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def _make_runtime_helpfile(path: Path, gas_list: list[str]): + """Create a minimal runtime helpfile with physically valid positive values.""" + data = { + 'Time': [1.0], + 'T_surf': [300.0], + 'T_pot': [1500.0], + 'F_int': [100.0], + 'F_olr': [220.0], + 'F_ins': [340.0], + 'Phi_global_vol': [0.2], + 'O2_bar': [1.0e-3], + 'C_kg_solid': [1.0e17], + 'C_kg_liquid': [2.0e17], + 'C_kg_atm': [3.0e17], + 'H_kg_solid': [4.0e17], + 'H_kg_liquid': [5.0e17], + 'H_kg_atm': [6.0e17], + 'O_kg_atm': [7.0e17], + 'P_surf': [1.0], + 'atm_kg_per_mol': [2.8e-2], + 'R_obs': [6.5e6], + 'R_int': [6.3e6], + 'RF_depth': [0.01], + 'Phi_global': [0.15], + } + for gas in gas_list: + data[f'{gas}_bar'] = [1.0e-6] + + df = pd.DataFrame(data) + df.to_csv(path, sep=' ', index=False) + + +@pytest.mark.unit +def test_postproc_once_writes_scalar_csv_and_orders_log_message(tmp_path): + """Writes scalar CSV and logs scalar-write message before config parsing.""" + chili_postproc = _load_chili_postproc_module() + + simdir = tmp_path / 'run_chili_tr1b' + simdir.mkdir() + + init_path = simdir / 'init_coupler.toml' + init_path.write_text('placeholder = true\n', encoding='utf-8') + + all_gases = sorted(set(chili_postproc.vol_list)) + _make_runtime_helpfile(simdir / 'runtime_helpfile.csv', all_gases) + + class _FakeInteriorData: + def get_dict_values(self, keys): + if keys == ['data', 'temp_b']: + return np.array([1000.0, 1500.0, 2000.0]) + if keys == ['data', 'visc_b']: + return np.array([1.0e20, 2.0e20, 3.0e20]) + raise KeyError(keys) + + recorded_messages = [] + original_print = builtins.print + + def fake_print(*args, **kwargs): + msg = ' '.join(str(a) for a in args) + recorded_messages.append(msg) + + def fake_read_config(_config_path): + recorded_messages.append('READ_CONFIG_CALLED') + return SimpleNamespace( + orbit=SimpleNamespace(s0_factor=1.0), + atmos_clim=SimpleNamespace(albedo_pl=0.3, surface_d=1000.0), + ) + + chili_postproc.read_config_object = fake_read_config + chili_postproc.read_jsons = lambda _simdir, _times: [_FakeInteriorData()] + chili_postproc.read_ncdf_profile = lambda *_args, **_kwargs: None + + try: + builtins.print = fake_print + name = chili_postproc.postproc_once(str(simdir), plot=False) + finally: + builtins.print = original_print + + assert name == 'trappist1b' + assert ' write CSV file for scalars' in recorded_messages + assert 'READ_CONFIG_CALLED' in recorded_messages + assert recorded_messages.index(' write CSV file for scalars') < recorded_messages.index( + 'READ_CONFIG_CALLED' + ) + + out_csv = simdir / 'chili' / 'evolution-proteus-trappist1b-data.csv' + assert out_csv.is_file() + + out_df = pd.read_csv(out_csv) + assert 'flux_ASR(W/m2)' in out_df.columns + assert 'viscosity(Pa.s)' in out_df.columns + assert out_df.loc[0, 'viscosity(Pa.s)'] == pytest.approx(2.0e20) + + +@pytest.mark.unit +def test_postproc_once_raises_when_helpfile_missing(tmp_path): + """Raises ``FileNotFoundError`` when runtime helpfile is absent.""" + chili_postproc = _load_chili_postproc_module() + + simdir = tmp_path / 'case_no_helpfile' + simdir.mkdir() + + with pytest.raises(FileNotFoundError, match='runtime_helpfile.csv'): + chili_postproc.postproc_once(str(simdir), plot=False) diff --git a/tools/chili_generate.py b/tools/chili_generate.py index c3da08d12..e544ddc08 100755 --- a/tools/chili_generate.py +++ b/tools/chili_generate.py @@ -12,7 +12,7 @@ import toml # Options -use_scratch = False # Store output in scratch folder +use_scratch = True # Store output in scratch folder # ----------------------------------------------- @@ -46,7 +46,8 @@ cfg[p] = deepcopy(cfg['base']) # output - cfg[p]['params']['out']['path'] = f'chili_{p}' + cfg[p]['params']['out']['path'] = 'scratch/' if use_scratch else '' + cfg[p]['params']['out']['path'] += f'chili_{p}/' # star cfg[p]['star']['mass'] = 1.0 # Msun diff --git a/tools/chili_postproc.py b/tools/chili_postproc.py index db6d277c2..58e3fd964 100755 --- a/tools/chili_postproc.py +++ b/tools/chili_postproc.py @@ -16,17 +16,23 @@ import tomllib import matplotlib.pyplot as plt -import netCDF4 as nc import numpy as np import pandas as pd +from proteus import Proteus +from proteus.atmos_clim.common import read_ncdf_profile from proteus.config import read_config_object +from proteus.interior.spider import read_jsons from proteus.utils.constants import R_earth, vol_list from proteus.utils.plot import get_colour, latexify # Target times for sampling profile [log years] tau_target = [3, 4, 5, 7, 9] +# Gas lists +chili_gases = ['H2O', 'CO2', 'CO', 'H2', 'CH4', 'O2'] +extra_gases = [gas for gas in vol_list if gas not in chili_gases] + # Potential planet names pl_names = { 'tr1b': 'trappist1b', @@ -67,37 +73,65 @@ def postproc_once(simdir: str, plot: bool = True): # Read config config = read_config_object(config_path) + # Extract data from simulation + print(' extract data from simulation') + handler = Proteus(config_path=config_path) + if handler is None: + raise RuntimeError(f'Cannot create Proteus handler for {simdir}') + handler.extract_archives() + # Write to expected format + print(' write CSV file for scalars') + # https://github.com/projectcuisines/chili/blob/main/intercomparison/README.md#evolution-models out = {} - out['t(yr)'] = np.array(hf_all['Time'].iloc[:]) - out['T_surf(K)'] = np.array(hf_all['T_surf'].iloc[:]) - out['T_pot(K)'] = np.array(hf_all['T_pot'].iloc[:]) - out['phi(vol_frac)'] = np.array(hf_all['Phi_global_vol'].iloc[:]) - out['massC_solid(kg)'] = np.array(hf_all['C_kg_solid'].iloc[:]) - out['massC_melt(kg)'] = np.array(hf_all['C_kg_liquid'].iloc[:]) - out['massC_atm(kg)'] = np.array(hf_all['C_kg_atm'].iloc[:]) - out['massH_solid(kg)'] = np.array(hf_all['H_kg_solid'].iloc[:]) - out['massH_melt(kg)'] = np.array(hf_all['H_kg_liquid'].iloc[:]) - out['massH_atm(kg)'] = np.array(hf_all['H_kg_atm'].iloc[:]) - out['massO_atm(kg)'] = np.array(hf_all['O_kg_atm'].iloc[:]) - out['mmw(kg/mol)'] = np.array(hf_all['atm_kg_per_mol'].iloc[:]) - out['p_surf(bar)'] = np.array(hf_all['P_surf'].iloc[:]) - for gas in vol_list: - out[f'p_{gas}(bar)'] = np.array(hf_all[f'{gas}_bar'].iloc[:]) - out['fO2_liquid(bar)'] = out['p_O2(bar)'] - out['fO2_solid(bar)'] = np.zeros_like(out['t(yr)']) - out['R_trans(m)'] = np.array(hf_all['R_obs'].iloc[:]) - out['flux_surf(W/m2)'] = np.array(hf_all['F_int'].iloc[:]) - out['flux_OLR(W/m2)'] = np.array(hf_all['F_olr'].iloc[:]) + # required columns + out['t(yr)'] = np.array(hf_all['Time']) + out['T_surf(K)'] = np.array(hf_all['T_surf']) + out['T_pot(K)'] = np.array(hf_all['T_pot']) + + out['flux_surf(W/m2)'] = np.array(hf_all['F_int']) + out['flux_OLR(W/m2)'] = np.array(hf_all['F_olr']) out['flux_ASR(W/m2)'] = ( - np.array(hf_all['F_ins'].iloc[:]) - * config.orbit.s0_factor - * (1 - config.atmos_clim.albedo_pl) + np.array(hf_all['F_ins']) * config.orbit.s0_factor * (1 - config.atmos_clim.albedo_pl) ) + + out['phi(vol_frac)'] = np.array(hf_all['Phi_global_vol']) + out['fO2_melt(bar)'] = np.array(hf_all['O2_bar']) + out['fO2_solid(bar)'] = out['fO2_melt(bar)'] # assumes homogeneous fO2 out['thick_surf_bl(m)'] = np.ones_like(out['t(yr)']) * config.atmos_clim.surface_d + out['massC_solid(kg)'] = np.array(hf_all['C_kg_solid']) + out['massC_melt(kg)'] = np.array(hf_all['C_kg_liquid']) + out['massC_atm(kg)'] = np.array(hf_all['C_kg_atm']) + out['massH_solid(kg)'] = np.array(hf_all['H_kg_solid']) + out['massH_melt(kg)'] = np.array(hf_all['H_kg_liquid']) + out['massH_atm(kg)'] = np.array(hf_all['H_kg_atm']) + out['massO_atm(kg)'] = np.array(hf_all['O_kg_atm']) + + out['p_surf(bar)'] = np.array(hf_all['P_surf']) + for gas in chili_gases: + out[f'p_{gas}(bar)'] = np.array(hf_all[f'{gas}_bar']) + out['mmw(kg/mol)'] = np.array(hf_all['atm_kg_per_mol']) + + out['R_trans(m)'] = np.array(hf_all['R_obs']) + out['R_solid(m)'] = np.array(hf_all['R_int']) * (1 - hf_all['RF_depth']) + + visc_arr = [] + for idx_t, t in enumerate(out['t(yr)']): + int_data = read_jsons(simdir, [t])[0] + int_tmp = int_data.get_dict_values(['data', 'temp_b']) + int_eta = int_data.get_dict_values(['data', 'visc_b']) + # get viscosity of mantle layer equal to potential temperature + idx_z = np.argmin(np.abs(int_tmp - hf_all['T_pot'].iloc[idx_t])) + visc_arr.append(int_eta[idx_z]) + out['viscosity(Pa.s)'] = np.array(visc_arr) + + # Additional columns + out['phi(mass_frac)'] = hf_all['Phi_global'] + for gas in extra_gases: + out[f'p_{gas}(bar)'] = np.array(hf_all[f'{gas}_bar']) + # Write to scalar CSV - print(' write CSV file for scalars') outpath = os.path.join(chilidir, f'evolution-proteus-{name}-data.csv') pd.DataFrame(out).to_csv(outpath, sep=',', index=False, float_format='%.10e') @@ -114,15 +148,20 @@ def postproc_once(simdir: str, plot: bool = True): print(f' tau{tau} -> {tclose:.2e} yr') ncfile = os.path.join(simdir, 'data', f'{tclose:.0f}_atm.nc') - with nc.Dataset(ncfile) as ds: - tarr = np.array(ds.variables['tmp'][:]) - parr = np.array(ds.variables['p'][:]) * 1e-5 - zarr = np.array(ds.variables['r'][:]) - np.amin(ds.variables['r']) - X = np.array([tarr, parr, zarr]).T - H = 'T(K),p(bar),z(m)' + atm = read_ncdf_profile(ncfile, extra_keys=['x_gas'], combine_edges=False) + + # climate profile + X = [atm['z'], atm['p'] / 1e5, atm['t']] + H = 'z(m),p_tot(bar),T(K)' + + # chemistry profile + for gas in vol_list: + X.append(atm[f'{gas}_vmr'] * atm['p'] / 1e5) # partial pressure profiles + H += f',p_{gas}(bar)' + csvname = f'evolution-proteus-{name}-tau{tau}-data.csv' f = os.path.join(chilidir, csvname) - np.savetxt(f, X, fmt='%.10e', delimiter=',', header=H) + np.savetxt(f, np.array(X).T, fmt='%.10e', delimiter=',', header=H, comments='') # Make plot... if plot: @@ -238,6 +277,10 @@ def postproc_once(simdir: str, plot: bool = True): fig.tight_layout() fig.savefig(os.path.join(chilidir, 'chili.pdf'), dpi=300, bbox_inches='tight') + # Create archive + print(' create archive of output data') + handler.create_archives() + return name @@ -318,8 +361,8 @@ def postproc_grid(griddir: str): if 'tau' in fpath: continue hf_all = pd.read_csv(fpath, delimiter=',') - x = np.array(hf_all['t(yr)'].iloc[:]) - y = np.array(hf_all['T_surf(K)'].iloc[:]) + x = np.array(hf_all['t(yr)']) + y = np.array(hf_all['T_surf(K)']) l = os.path.basename(fpath).split('.')[0] l = ' '.join(l.split('-')[4:6]) ax.plot(x, y, label=l, lw=1.5, alpha=0.8, zorder=4)