From 496fbe0ea806ba0a407f0696da6d1b1d20525740 Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 15:46:32 +0000 Subject: [PATCH 01/17] Add aersols_enabled config option and update agni.toml file. Add validators. --- README.md | 12 +++++------ input/demos/agni.toml | 3 +++ src/proteus/atmos_clim/agni.py | 36 +++++++++++++++++++++++++++++++ src/proteus/config/_atmos_clim.py | 5 ++++- 4 files changed, 49 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index bb1880ff4..c10f3a259 100644 --- a/README.md +++ b/README.md @@ -32,11 +32,11 @@ More information can be found on the [documentation](https://proteus-framework.org/PROTEUS/) pages: -* [model description](https://proteus-framework.org/PROTEUS/model.html) -* [installation guide](https://proteus-framework.org/PROTEUS/installation.html) -* [usage guide](https://proteus-framework.org/PROTEUS/usage.html) -* [contributing guide](https://proteus-framework.org/PROTEUS/CONTRIBUTING.html) +* [model description](https://proteus-framework.org/PROTEUS/Explanations/model.html) +* [installation guide](https://proteus-framework.org/PROTEUS/How-to/installation.html) +* [usage guide](https://proteus-framework.org/PROTEUS/How-to/usage.html) +* [contributing guide](https://proteus-framework.org/PROTEUS/Community/CONTRIBUTING.html) -You can find help on the [discussions page](https://github.com/orgs/FormingWorlds/discussions) or by [contacting the developers](https://proteus-framework.org/PROTEUS/contact.html) directly. +You can find help on the [discussions page](https://github.com/orgs/FormingWorlds/discussions) or by [contacting the developers](https://proteus-framework.org/PROTEUS/Community/contact.html) directly. -If you make use of PROTEUS, please reference the papers outlined in the [bibliography](https://proteus-framework.org/PROTEUS/bibliography.html). +If you make use of PROTEUS, please reference the papers outlined in the [bibliography](https://proteus-framework.org/PROTEUS/Reference/bibliography.html). diff --git a/input/demos/agni.toml b/input/demos/agni.toml index 769642c0b..02185b9aa 100644 --- a/input/demos/agni.toml +++ b/input/demos/agni.toml @@ -104,6 +104,9 @@ version = "2.0" overlap_method = "ee" # gas overlap method surf_roughness = 1e-3 # characteristic surface roughness [m] surf_windspeed = 2.0 # characteristic surface wind speed [m/s]. + cloud_enabled = false # include water cloud radiative effects? + cloud_alpha = 0.5 # water condensate retention fraction + aerosols_enabled = false # include aerosol radiative effects? rainout = true # include volatile condensation/evaporation aloft latent_heat = false # include latent heat release when `rainout=true`? oceans = true # form liquid oceans at planet surface? diff --git a/src/proteus/atmos_clim/agni.py b/src/proteus/atmos_clim/agni.py index 517f51fd9..a19bac008 100644 --- a/src/proteus/atmos_clim/agni.py +++ b/src/proteus/atmos_clim/agni.py @@ -117,6 +117,33 @@ def _determine_condensates(vol_list: list): # all dry gases... return [v for v in vol_list if v not in ALWAYS_DRY] +def _determine_aerosols(dirs: dict) -> list: + """ + Determine which aerosols are available. + + Parameters + ---------- + dirs : dict + Dictionary containing paths to directories + + Returns + ---------- + aerosols : list + List of available aerosols + """ + + aerosol_dir = os.path.join(dirs['fwl'], 'aerosols') + if not os.path.isdir(aerosol_dir): + log.warning(f'Aerosol directory not found: {aerosol_dir}') + return [] + + aerosols = [] + for f in os.listdir(aerosol_dir): + if f.endswith('.mon'): + aerosols.append(f.split('.')[0]) + + log.debug(f'Available aerosols: {aerosols}') + return aerosols def init_agni_atmos(dirs: dict, config: Config, hf_row: dict): """Initialise atmosphere struct for use by AGNI. @@ -197,6 +224,13 @@ def init_agni_atmos(dirs: dict, config: Config, hf_row: dict): p_top = config.atmos_clim.agni.p_top p_surf = max(p_surf, p_top * 1.1) # this will happen if the atmosphere is stripped + # Aerosol species dictionary (set MMR to zero initially) + aerosol_species = dict() + if config.atmos_clim.aerosols_enabled: + aerosol_species = {a: 0.0 for a in _determine_aerosols(dirs)} + if len(aerosol_species) == 0: + log.warning('No data found for aerosol species') + # Setup struct succ = jl.AGNI.atmosphere.setup_b( atmos, @@ -218,6 +252,8 @@ def init_agni_atmos(dirs: dict, config: Config, hf_row: dict): IO_DIR=io_dir, flag_rayleigh=config.atmos_clim.rayleigh, flag_cloud=config.atmos_clim.cloud_enabled, + flag_aerosol=config.atmos_clim.aerosols_enabled, + aerosol_species=convert(jl.Dict, aerosol_species), overlap_method=config.atmos_clim.agni.overlap_method, albedo_s=config.atmos_clim.surf_greyalbedo, surface_material=surface_material, diff --git a/src/proteus/config/_atmos_clim.py b/src/proteus/config/_atmos_clim.py index bb3cf694e..afa81b60b 100644 --- a/src/proteus/config/_atmos_clim.py +++ b/src/proteus/config/_atmos_clim.py @@ -310,6 +310,8 @@ class AtmosClim: Conductive skin thickness [m], surface_k: float Conductive skin thermal conductivity [W m-1 K-1]. + aerosols_enabled: bool + Enable aerosol radiative effects. cloud_enabled: bool Enable water cloud radiative effects. cloud_alpha: float @@ -346,7 +348,8 @@ class AtmosClim: prevent_warming: bool = field(default=False) surface_d: float = field(default=0.01, validator=gt(0)) surface_k: float = field(default=2.0, validator=gt(0)) - cloud_enabled: bool = field(default=False) + aerosols_enabled: bool = field(default=False, validator=warn_if_dummy) + cloud_enabled: bool = field(default=False, validator=warn_if_dummy) cloud_alpha: float = field(default=0.0, validator=(ge(0), le(1))) surf_greyalbedo: float = field(default=0.2, validator=(ge(0), le(1))) albedo_pl = field(default=0.0, validator=valid_albedo) From 13ef226f2bd9a55a1976c44e3a533690ac23f7c3 Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 16:22:07 +0000 Subject: [PATCH 02/17] Update CLI to list available plots, on error. Read aerosol data from netcdf. Update chemistry plot. Add colours --- src/proteus/atmos_clim/agni.py | 8 +- src/proteus/atmos_clim/common.py | 15 +++ src/proteus/cli.py | 1 + src/proteus/plot/cpl_chem_atmosphere.py | 126 +++++++++++++++++++----- src/proteus/utils/plot.py | 10 +- 5 files changed, 132 insertions(+), 28 deletions(-) diff --git a/src/proteus/atmos_clim/agni.py b/src/proteus/atmos_clim/agni.py index a19bac008..ebbe2df25 100644 --- a/src/proteus/atmos_clim/agni.py +++ b/src/proteus/atmos_clim/agni.py @@ -132,13 +132,13 @@ def _determine_aerosols(dirs: dict) -> list: List of available aerosols """ - aerosol_dir = os.path.join(dirs['fwl'], 'aerosols') - if not os.path.isdir(aerosol_dir): - log.warning(f'Aerosol directory not found: {aerosol_dir}') + scattering_dir = os.path.join(dirs['fwl'], 'scattering') + if not os.path.isdir(scattering_dir): + log.warning(f'Scattering data directory not found: {scattering_dir}') return [] aerosols = [] - for f in os.listdir(aerosol_dir): + for f in os.listdir(scattering_dir): if f.endswith('.mon'): aerosols.append(f.split('.')[0]) diff --git a/src/proteus/atmos_clim/common.py b/src/proteus/atmos_clim/common.py index 5bbb3e1c0..5dcb8b264 100644 --- a/src/proteus/atmos_clim/common.py +++ b/src/proteus/atmos_clim/common.py @@ -150,6 +150,21 @@ def read_ncdf_profile(nc_fpath: str, extra_keys: list = [], combine_edges: bool gas_lbl = ''.join([c.decode(encoding='utf-8') for c in gas]).strip() out[gas_lbl + '_vmr'] = np.array(gas_x[:, igas]) + # Reading aerosol mass mixing ratios + elif key == 'aer_mmr': + if 'aerosols' in ds.variables.keys(): + aer_l = ds.variables['aerosols'][:] # names (bytes matrix) + aer_x = ds.variables['aer_mmr'][:] # mmrs (float matrix) + + # get data for each aerosol + for iaer, aer in enumerate(aer_l): + if aer[0] is not np.bytes_: + continue + aer_lbl = ''.join([c.decode(encoding='utf-8') for c in aer]).strip() + out[aer_lbl + '_mmr'] = np.array(aer_x[:, iaer]) + # Also store raw array + out[key] = np.array(ds.variables[key][:]) + else: out[key] = np.array(ds.variables[key][:]) diff --git a/src/proteus/cli.py b/src/proteus/cli.py index 0eff28d6d..d71ec0d04 100644 --- a/src/proteus/cli.py +++ b/src/proteus/cli.py @@ -101,6 +101,7 @@ def plot(plots, config_path: Path): for plot in plots: if plot not in plot_dispatch.keys(): click.echo(f'Invalid plot: {plot}') + click.echo('Available plots: ' + ', '.join(plot_dispatch.keys())) else: click.echo(f'Plotting: {plot}') plot_func = plot_dispatch[plot] diff --git a/src/proteus/plot/cpl_chem_atmosphere.py b/src/proteus/plot/cpl_chem_atmosphere.py index 3851bc2d7..0617a88a9 100644 --- a/src/proteus/plot/cpl_chem_atmosphere.py +++ b/src/proteus/plot/cpl_chem_atmosphere.py @@ -91,9 +91,12 @@ def plot_chem_atmosphere( log.warning('No atmosphere NetCDF files found in output folder') return nc_fpath = natural_sort(files)[-1] - atm_profile = read_ncdf_profile(nc_fpath, extra_keys=['pl', 'tmpl', 'x_gas']) + atm_profile = read_ncdf_profile( + nc_fpath, extra_keys=['pl', 'tmpl', 'x_gas', 'cloud_mmr', 'cloud_area', 'aer_mmr'] + ) parr = atm_profile['pl'] * 1e-5 # convert to bar + tarr = atm_profile['tmpl'] # temperature profile # Get year year = float(nc_fpath.split('/')[-1].split('_atm')[0]) @@ -107,10 +110,13 @@ def plot_chem_atmosphere( else: has_offchem = False - # init plot - scale = 1.2 - fig, ax = plt.subplots(1, 1, figsize=(6 * scale, 5 * scale)) + # init plot with two panels + scale = 1.15 + fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(5 * scale, 6 * scale)) + # ==================== + # Panel 1: Gas Species + # ==================== # plot species profiles lw = 0.9 al = 0.8 @@ -131,48 +137,122 @@ def plot_chem_atmosphere( xarr = [xarr[0]] + xarr if np.amax(xarr) >= xmin: vmr = float(xarr[-1]) - ax.plot(xarr, parr, ls='dashed', color=col, lw=_lw, alpha=al) + ax1.plot(xarr, parr, ls='dashed', color=col, lw=_lw, alpha=al) # plot from offline chemistry, if available (solid lines) if has_offchem and (gas in atm_offchem.keys()): xarr = list(atm_offchem[gas].values) if np.amax(xarr) >= xmin: vmr = float(xarr[-1]) # prefer vmr from offline chemistry - ax.plot(xarr, parr, ls='solid', color=col, lw=_lw, alpha=al) + ax1.plot(xarr, parr, ls='solid', color=col, lw=_lw, alpha=al) # create legend entry and store surface vmr if vmr > 0.0: - ax.plot(1e30, 1e30, ls='solid', color=col, lw=_lw, alpha=al, label=lbl) + ax1.plot(1e30, 1e30, ls='solid', color=col, lw=_lw, alpha=al, label=lbl) vmr_surf.append(vmr) - # Decorate - ax.set_xscale('log') - ax.set_xlim(left=xmin, right=1.1) - ax.set_xlabel('Volume mixing ratio, at t=%.2e yr' % year) - ax.xaxis.set_major_locator(LogLocator(numticks=1000)) + # Add temperature profile to Panel 1 (secondary x-axis at top) + ax1_temp = ax1.twiny() + ax1_temp.plot(tarr, parr, 'k', lw=1.5, alpha=0.6, label='Temperature') + ax1_temp.set_xlabel('Temperature [K]') + ax1_temp.set_xlim(left=0) + + # Decorate Panel 1 + ax1.set_xscale('symlog', linthresh=xmin) + ax1.set_xlim(left=0, right=1.1) + ax1.set_xlabel('Gas volume mixing ratio') - ax.set_ylabel('Pressure [bar]') - ax.set_yscale('log') - ax.set_ylim(bottom=np.amax(parr), top=np.amin(parr)) - ax.yaxis.set_major_locator(LogLocator(numticks=1000)) + ax1.set_ylabel('Pressure [bar]') + ax1.set_yscale('log') + ax1.set_ylim(bottom=np.amax(parr), top=np.amin(parr)) + ax1.yaxis.set_major_locator(LogLocator(numticks=1000)) # Legend (handles sorted by surface vmr) - handles, labels = ax.get_legend_handles_labels() + handles, labels = ax1.get_legend_handles_labels() order = np.argsort(vmr_surf)[::-1] - ax.legend( + ax1.legend( [handles[idx] for idx in order], [labels[idx] for idx in order], - loc='lower center', - bbox_to_anchor=(0.5, 1), - ncols=11, + loc='center left', + bbox_to_anchor=(1.00, 0.5), + ncols=1, + fontsize=9, + borderpad=0.0, + labelspacing=0.3, + columnspacing=1.0, + handlelength=1.0, + handletextpad=0.3, + frameon=False, + ) + + # ============================= + # Panel 2: Aerosols and Clouds + # ============================= + # Cloud profiles + if 'cloud_mmr' in atm_profile.keys(): + cloud_mmr = atm_profile['cloud_mmr'] + # Extend to match parr length (edge values) + if np.max(cloud_mmr) > 0: + ax2.plot( + [cloud_mmr[0]] + list(cloud_mmr), + parr, + ls='solid', + color=get_colour('cloud'), + lw=1.5, + alpha=0.7, + label='cloud', + ) + + # Aerosol profiles + # Check for individual aerosol species + for key in atm_profile.keys(): + if key.endswith('_mmr') and "cloud" not in key: + aer_mmr = atm_profile[key] + if np.max(aer_mmr) > 0: + aer_name = key.split('_')[0] + ax2.plot( + [aer_mmr[0]] + list(aer_mmr), + parr, + ls='solid', + lw=1.5, + alpha=0.7, + color=get_colour(aer_name), + label=aer_name, + ) + + # Add temperature profile to Panel 2 (secondary x-axis at top) + ax2_temp = ax2.twiny() + ax2_temp.plot(tarr, parr, 'k', lw=1.5, alpha=0.6) + ax2_temp.set_xticks([]) # hide x-tick labels on secondary axis + ax2_temp.set_xlim(left=0) + + # Decorate Panel 2 + ax2.set_xscale('symlog', linthresh=xmin) + ax2.set_xlim(left=0, right=1.1) + ax2.set_xlabel('Aerosol mass mixing ratio') + + ax2.set_ylabel('Pressure [bar]') + ax2.set_yscale('log') + ax2.set_ylim(bottom=np.amax(parr), top=np.amin(parr)) + ax2.yaxis.set_major_locator(LogLocator(numticks=1000)) + + ax2.legend( + loc='center left', + bbox_to_anchor=(1.00, 0.5), + ncols=1, fontsize=9, - borderpad=0.4, + borderpad=0.0, labelspacing=0.3, columnspacing=1.0, - handlelength=1.5, + handlelength=1.0, handletextpad=0.3, + frameon=False, ) + # Add time annotation + fig.text(0.02, 0.98, f't = {year:.2e} yr', transform=ax1.transAxes, + ha='left', va='top',fontsize=11, weight='bold', zorder=999) + # Save file fpath = os.path.join(output_dir, 'plots', 'plot_chem_atmosphere.%s' % plot_format) fig.savefig(fpath, dpi=200, bbox_inches='tight') diff --git a/src/proteus/utils/plot.py b/src/proteus/utils/plot.py index c9238baba..57475b53f 100644 --- a/src/proteus/utils/plot.py +++ b/src/proteus/utils/plot.py @@ -38,7 +38,7 @@ 'O': '#88bb22', 'N': '#ffaa00', 'S': '#ff22ff', - 'P': '#33ccff', + 'P': "#283134", # refractory elements 'Fe': '#888888', 'Si': '#aa2277', @@ -65,6 +65,14 @@ 'atm_bkg': '#f2faff', 'int_bkg': '#fffaf2', 'cor_bkg': '#efefef', + # Clouds and aerosols + 'cloud': "#027FB1", + 'freshsoot': "#876e44", + 'soot': "#63401a", + 'agedsoot': "#251607", + 'nitrate': "#ffaa00", + 'ash': "#970505", + 'sulph': "#ff22ff" } From 9dcf7a92270cf83eacdf05548ffdd3055a345a47 Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 16:22:14 +0000 Subject: [PATCH 03/17] restore P colour --- src/proteus/utils/plot.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/proteus/utils/plot.py b/src/proteus/utils/plot.py index 57475b53f..359eba42a 100644 --- a/src/proteus/utils/plot.py +++ b/src/proteus/utils/plot.py @@ -38,7 +38,7 @@ 'O': '#88bb22', 'N': '#ffaa00', 'S': '#ff22ff', - 'P': "#283134", + 'P': "#33ccff", # refractory elements 'Fe': '#888888', 'Si': '#aa2277', From 1b9f58093c0482e4d17cae09fe4fd881e12fa680 Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 16:29:17 +0000 Subject: [PATCH 04/17] Get scattering data from Zenodo --- input/all_options.toml | 4 ++++ src/proteus/atmos_clim/agni.py | 2 +- src/proteus/cli.py | 9 ++++++++- src/proteus/utils/data.py | 23 +++++++++++++++++++++++ 4 files changed, 36 insertions(+), 2 deletions(-) diff --git a/input/all_options.toml b/input/all_options.toml index bcf054c36..2ba5c3014 100644 --- a/input/all_options.toml +++ b/input/all_options.toml @@ -207,6 +207,7 @@ version = "2.0" prevent_warming = true # do not allow the planet to heat up surface_d = 0.01 # m, conductive skin thickness surface_k = 2.0 # W m-1 K-1, conductive skin thermal conductivity + aerosols_enabled = false # enable aerosol radiative effects cloud_enabled = false # enable water cloud radiative effects cloud_alpha = 0.0 # condensate retention fraction (1 -> fully retained) surf_state = "fixed" # surface scheme: "mixed_layer" | "fixed" | "skin" @@ -233,6 +234,9 @@ version = "2.0" overlap_method = "ee" # gas overlap method surf_roughness = 1e-3 # characteristic surface roughness [m] surf_windspeed = 2.0 # characteristic surface wind speed [m/s]. + cloud_enabled = false # include water cloud radiative effects? + cloud_alpha = 0.5 # water condensate retention fraction + aerosols_enabled = false # include aerosol radiative effects? rainout = true # include volatile condensation/evaporation aloft latent_heat = false # include latent heat release when `rainout=true`? oceans = true # form liquid oceans at planet surface? diff --git a/src/proteus/atmos_clim/agni.py b/src/proteus/atmos_clim/agni.py index ebbe2df25..4b2b113be 100644 --- a/src/proteus/atmos_clim/agni.py +++ b/src/proteus/atmos_clim/agni.py @@ -132,7 +132,7 @@ def _determine_aerosols(dirs: dict) -> list: List of available aerosols """ - scattering_dir = os.path.join(dirs['fwl'], 'scattering') + scattering_dir = os.path.join(dirs['fwl'], 'scattering', 'scattering') if not os.path.isdir(scattering_dir): log.warning(f'Scattering data directory not found: {scattering_dir}') return [] diff --git a/src/proteus/cli.py b/src/proteus/cli.py index d71ec0d04..825d25cf7 100644 --- a/src/proteus/cli.py +++ b/src/proteus/cli.py @@ -22,7 +22,7 @@ from proteus import Proteus from proteus import __version__ as proteus_version from proteus.config import read_config_object -from proteus.utils.data import download_sufficient_data +from proteus.utils.data import download_scattering, download_sufficient_data from proteus.utils.logs import setup_logger config_option = click.option( @@ -409,6 +409,12 @@ def surfaces(): download_surface_albedos() +@click.command() +def scattering(): + """Get scattering radiative properties""" + from .utils.data import download_scattering + + download_scattering() @click.command() def reference(): @@ -463,6 +469,7 @@ def spider(): cli.add_command(get) get.add_command(spectral) get.add_command(surfaces) +get.add_command(scattering) get.add_command(muscles) get.add_command(phoenix) get.add_command(solar) diff --git a/src/proteus/utils/data.py b/src/proteus/utils/data.py index 60a715582..bac409bf4 100644 --- a/src/proteus/utils/data.py +++ b/src/proteus/utils/data.py @@ -478,6 +478,8 @@ def validate_zenodo_folder(zenodo_id: str, folder_dir: Path, hash_maxfilesize=10 'Population': {'zenodo_id': '15727998', 'osf_id': 'dpkjb', 'osf_project': 'dpkjb'}, # EOS material properties (OSF project: dpkjb) 'EOS_Seager2007': {'zenodo_id': '15727998', 'osf_id': 'dpkjb', 'osf_project': 'dpkjb'}, + # Aerosol scattering data (no OSF project) + 'scattering': {'zenodo_id': '19232950'}, } @@ -999,6 +1001,23 @@ def download_surface_albedos(): ) +def download_scattering(): + """ + Download scattering radiative properties data + """ + folder = 'scattering' + source_info = get_data_source_info(folder) + if not source_info: + raise ValueError(f'No data source mapping found for folder: {folder}') + + download( + folder=folder, + target='scattering', + zenodo_id=source_info['zenodo_id'], + desc='radiative properties scattering data', + ) + + def download_spectral_file(name: str, bands: str): """ Download spectral file. @@ -1417,6 +1436,10 @@ def _get_sufficient(config: Config, clean: bool = False): if config.atmos_clim.module == 'agni': download_surface_albedos() + # Aerosol scattering data + if config.atmos_clim.module == 'agni' and config.atmos_clim.aerosols_enabled: + download_scattering() + # Exoplanet population data download_exoplanet_data() From d078c1a34f6fde628bb34d431189edd4397d7df2 Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 16:30:09 +0000 Subject: [PATCH 05/17] Ruff --- src/proteus/atmos_clim/agni.py | 2 ++ src/proteus/cli.py | 4 +++- src/proteus/plot/cpl_chem_atmosphere.py | 15 ++++++++++++--- src/proteus/utils/plot.py | 16 ++++++++-------- 4 files changed, 25 insertions(+), 12 deletions(-) diff --git a/src/proteus/atmos_clim/agni.py b/src/proteus/atmos_clim/agni.py index 4b2b113be..e5d3286dc 100644 --- a/src/proteus/atmos_clim/agni.py +++ b/src/proteus/atmos_clim/agni.py @@ -117,6 +117,7 @@ def _determine_condensates(vol_list: list): # all dry gases... return [v for v in vol_list if v not in ALWAYS_DRY] + def _determine_aerosols(dirs: dict) -> list: """ Determine which aerosols are available. @@ -145,6 +146,7 @@ def _determine_aerosols(dirs: dict) -> list: log.debug(f'Available aerosols: {aerosols}') return aerosols + def init_agni_atmos(dirs: dict, config: Config, hf_row: dict): """Initialise atmosphere struct for use by AGNI. diff --git a/src/proteus/cli.py b/src/proteus/cli.py index 825d25cf7..921ebce3e 100644 --- a/src/proteus/cli.py +++ b/src/proteus/cli.py @@ -22,7 +22,7 @@ from proteus import Proteus from proteus import __version__ as proteus_version from proteus.config import read_config_object -from proteus.utils.data import download_scattering, download_sufficient_data +from proteus.utils.data import download_sufficient_data from proteus.utils.logs import setup_logger config_option = click.option( @@ -409,6 +409,7 @@ def surfaces(): download_surface_albedos() + @click.command() def scattering(): """Get scattering radiative properties""" @@ -416,6 +417,7 @@ def scattering(): download_scattering() + @click.command() def reference(): """Get reference data (exoplanet populations, mass-radius curves, etc.)""" diff --git a/src/proteus/plot/cpl_chem_atmosphere.py b/src/proteus/plot/cpl_chem_atmosphere.py index 0617a88a9..b03c536ad 100644 --- a/src/proteus/plot/cpl_chem_atmosphere.py +++ b/src/proteus/plot/cpl_chem_atmosphere.py @@ -206,7 +206,7 @@ def plot_chem_atmosphere( # Aerosol profiles # Check for individual aerosol species for key in atm_profile.keys(): - if key.endswith('_mmr') and "cloud" not in key: + if key.endswith('_mmr') and 'cloud' not in key: aer_mmr = atm_profile[key] if np.max(aer_mmr) > 0: aer_name = key.split('_')[0] @@ -250,8 +250,17 @@ def plot_chem_atmosphere( ) # Add time annotation - fig.text(0.02, 0.98, f't = {year:.2e} yr', transform=ax1.transAxes, - ha='left', va='top',fontsize=11, weight='bold', zorder=999) + fig.text( + 0.02, + 0.98, + f't = {year:.2e} yr', + transform=ax1.transAxes, + ha='left', + va='top', + fontsize=11, + weight='bold', + zorder=999, + ) # Save file fpath = os.path.join(output_dir, 'plots', 'plot_chem_atmosphere.%s' % plot_format) diff --git a/src/proteus/utils/plot.py b/src/proteus/utils/plot.py index 359eba42a..5b6d3b0d3 100644 --- a/src/proteus/utils/plot.py +++ b/src/proteus/utils/plot.py @@ -38,7 +38,7 @@ 'O': '#88bb22', 'N': '#ffaa00', 'S': '#ff22ff', - 'P': "#33ccff", + 'P': '#33ccff', # refractory elements 'Fe': '#888888', 'Si': '#aa2277', @@ -66,13 +66,13 @@ 'int_bkg': '#fffaf2', 'cor_bkg': '#efefef', # Clouds and aerosols - 'cloud': "#027FB1", - 'freshsoot': "#876e44", - 'soot': "#63401a", - 'agedsoot': "#251607", - 'nitrate': "#ffaa00", - 'ash': "#970505", - 'sulph': "#ff22ff" + 'cloud': '#027FB1', + 'freshsoot': '#876e44', + 'soot': '#63401a', + 'agedsoot': '#251607', + 'nitrate': '#ffaa00', + 'ash': '#970505', + 'sulph': '#ff22ff', } From bb00e365d1b5dddd1345817cea79a08f8a2ed3aa Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 16:38:01 +0000 Subject: [PATCH 06/17] Add tests with copilot --- tests/atmos_clim/test_common.py | 129 ++++++++++ tests/plot/test_cpl_chem_atmosphere.py | 336 +++++++++++++++++++++++++ tests/utils/test_data.py | 15 ++ 3 files changed, 480 insertions(+) create mode 100644 tests/plot/test_cpl_chem_atmosphere.py diff --git a/tests/atmos_clim/test_common.py b/tests/atmos_clim/test_common.py index eea7ca8b0..2542ef200 100644 --- a/tests/atmos_clim/test_common.py +++ b/tests/atmos_clim/test_common.py @@ -155,6 +155,135 @@ def test_read_ncdf_profile_without_combining_edges(mock_ds, mock_isfile): assert result['converged'] == 0.0 +@pytest.mark.unit +@patch('proteus.atmos_clim.common.os.path.isfile') +@patch('netCDF4.Dataset') +def test_read_ncdf_profile_with_aerosols(mock_ds, mock_isfile): + """ + Test reading NetCDF profile data with aerosol mass mixing ratios. + + Physical scenario: AGNI can output aerosol profiles (e.g., sulfate, silicate) + in the atmosphere. These are stored as aer_mmr(nlev_c, naeros) arrays with + corresponding names in the 'aerosols' variable. + """ + mock_isfile.return_value = True + + ds_instance = MagicMock() + mock_ds.return_value = ds_instance + + # Mock basic atmospheric profile + ds_instance.variables = { + 'p': np.array([100.0]), + 'pl': np.array([110.0, 90.0]), + 'tmp': np.array([300.0]), + 'tmpl': np.array([310.0, 290.0]), + 'r': np.array([6.4e6]), + 'rl': np.array([6.3e6, 6.5e6]), + 'planet_radius': [6.0e6], + 'transparent': np.array([b'y'], dtype='S1'), + # Aerosol data + 'aerosols': np.array( + [ + [ + b'S', + b'u', + b'l', + b'f', + b'a', + b't', + b'e', + b' ', + b' ', + b' ', + b' ', + b' ', + b' ', + b' ', + b' ', + b' ', + ], + [ + b'S', + b'i', + b'l', + b'i', + b'c', + b'a', + b't', + b'e', + b' ', + b' ', + b' ', + b' ', + b' ', + b' ', + b' ', + b' ', + ], + ], + dtype='S1', + ), + 'aer_mmr': np.array([[1e-6], [2e-6]]), # 2 levels, 1 aerosol type + } + + # Read with aerosol data + result = read_ncdf_profile('dummy.nc', extra_keys=['aer_mmr']) + + # Verify basic profile data + assert 'p' in result + assert 'transparent' in result + + # Verify aerosol data was read + assert 'aer_mmr' in result + np.testing.assert_allclose(result['aer_mmr'], np.array([[1e-6], [2e-6]])) + + # Note: Named aerosol extraction happens only if names are valid (not masked) + # With mock setup above, the names should parse + # But the actual extraction logic might skip empty/invalid names + + +@pytest.mark.unit +@patch('proteus.atmos_clim.common.os.path.isfile') +@patch('netCDF4.Dataset') +def test_read_ncdf_profile_with_clouds(mock_ds, mock_isfile): + """ + Test reading NetCDF profile data with cloud properties. + + Physical scenario: AGNI outputs cloud mass mixing ratio, cloud area fraction, + and cloud particle size when cloud_enabled=True. + """ + mock_isfile.return_value = True + + ds_instance = MagicMock() + mock_ds.return_value = ds_instance + + ds_instance.variables = { + 'p': np.array([100.0, 200.0]), + 'pl': np.array([110.0, 150.0, 190.0]), + 'tmp': np.array([300.0, 280.0]), + 'tmpl': np.array([310.0, 290.0, 270.0]), + 'r': np.array([6.4e6, 6.3e6]), + 'rl': np.array([6.5e6, 6.35e6, 6.2e6]), + 'planet_radius': [6.0e6], + 'solved': np.array([b'y'], dtype='S1'), + # Cloud data + 'cloud_mmr': np.array([1e-5, 2e-5]), + 'cloud_area': np.array([0.5, 0.8]), + 'cloud_size': np.array([1e-5, 1.2e-5]), + } + + result = read_ncdf_profile('dummy.nc', extra_keys=['cloud_mmr', 'cloud_area', 'cloud_size']) + + # Verify cloud data was read + assert 'cloud_mmr' in result + assert 'cloud_area' in result + assert 'cloud_size' in result + + np.testing.assert_allclose(result['cloud_mmr'], np.array([1e-5, 2e-5])) + np.testing.assert_allclose(result['cloud_area'], np.array([0.5, 0.8])) + np.testing.assert_allclose(result['cloud_size'], np.array([1e-5, 1.2e-5])) + + @pytest.mark.unit @patch('proteus.atmos_clim.common.read_ncdf_profile') def test_read_atmosphere_data(mock_read): diff --git a/tests/plot/test_cpl_chem_atmosphere.py b/tests/plot/test_cpl_chem_atmosphere.py new file mode 100644 index 000000000..7b135269b --- /dev/null +++ b/tests/plot/test_cpl_chem_atmosphere.py @@ -0,0 +1,336 @@ +"""Unit tests for ``proteus.plot.cpl_chem_atmosphere``. + +Covers the atmospheric chemistry plotting functions including aerosol and cloud +profiles. All matplotlib rendering, file I/O, NetCDF reads, and filesystem +operations are mocked for fast unit testing (<100 ms). + +Testing standards: + - docs/test_infrastructure.md + - docs/test_categorization.md + - docs/test_building.md +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from proteus.plot.cpl_chem_atmosphere import plot_chem_atmosphere + + +@pytest.mark.unit +@patch('proteus.plot.cpl_chem_atmosphere.glob.glob') +def test_plot_chem_atmosphere_no_files(mock_glob): + """ + plot_chem_atmosphere returns early when no NetCDF files found. + + Physical scenario: simulation hasn't run yet or output was deleted. + """ + mock_glob.return_value = [] + + # Should return without error + result = plot_chem_atmosphere('output_dir', 'vulcan', plot_format='png') + assert result is None + + +@pytest.mark.unit +@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') +@patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') +@patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') +@patch('proteus.plot.cpl_chem_atmosphere.glob.glob') +@patch('proteus.plot.cpl_chem_atmosphere.os.path.join') +def test_plot_chem_atmosphere_basic( + mock_join, mock_glob, mock_read, mock_subplots, mock_savefig +): + """ + Test basic plot generation with gas species only. + + Physical scenario: plotting atmospheric composition with major volatiles + (H2O, CO2, N2, etc.) from AGNI NetCDF output. + """ + # Mock file discovery + mock_glob.return_value = ['output_dir/data/100_atm.nc'] + mock_join.return_value = 'output_dir/plots/plot_chem_atmosphere.png' + + # Mock NetCDF read with typical gas species + mock_read.return_value = { + 'pl': np.array([1e8, 1e7, 1e6, 1e5]), # Pa + 'tmpl': np.array([3000, 2000, 1000, 500]), # K + 'H2O_vmr': np.array([1e-3, 1e-4, 1e-5]), + 'CO2_vmr': np.array([1e-2, 1e-3, 1e-4]), + 'N2_vmr': np.array([0.78, 0.78, 0.78]), + } + + # Mock matplotlib + fig_mock = MagicMock() + ax1_mock = MagicMock() + ax2_mock = MagicMock() + mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) + + # Mock twiny for temperature axis + ax1_temp_mock = MagicMock() + ax2_temp_mock = MagicMock() + ax1_mock.twiny.return_value = ax1_temp_mock + ax2_mock.twiny.return_value = ax2_temp_mock + + # Run function + plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) + + # Verify NetCDF was read with correct extra keys + mock_read.assert_called_once() + call_kwargs = mock_read.call_args[1] + assert 'pl' in call_kwargs['extra_keys'] + assert 'tmpl' in call_kwargs['extra_keys'] + assert 'x_gas' in call_kwargs['extra_keys'] + assert 'cloud_mmr' in call_kwargs['extra_keys'] + assert 'aer_mmr' in call_kwargs['extra_keys'] + + # Verify two panels were created + mock_subplots.assert_called_once() + args, kwargs = mock_subplots.call_args + assert args[0] == 2 # 2 rows + assert args[1] == 1 # 1 column + + # Verify temperature axes were created + ax1_mock.twiny.assert_called_once() + ax2_mock.twiny.assert_called_once() + + # Verify plot was saved + mock_savefig.assert_called_once() + + +@pytest.mark.unit +@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') +@patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') +@patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') +@patch('proteus.plot.cpl_chem_atmosphere.glob.glob') +@patch('proteus.plot.cpl_chem_atmosphere.os.path.join') +def test_plot_chem_atmosphere_with_clouds( + mock_join, mock_glob, mock_read, mock_subplots, mock_savefig +): + """ + Test plot generation including cloud profiles. + + Physical scenario: water clouds form in the atmosphere when cloud_enabled=True. + The plot should show cloud mass mixing ratio and area fraction in panel 2. + """ + mock_glob.return_value = ['output_dir/data/500_atm.nc'] + mock_join.return_value = 'output_dir/plots/plot_chem_atmosphere.png' + + # Mock NetCDF with cloud data + mock_read.return_value = { + 'pl': np.array([1e8, 1e7, 1e6, 1e5]), + 'tmpl': np.array([3000, 2000, 1000, 500]), + 'H2O_vmr': np.array([1e-3, 1e-4, 1e-5]), + 'cloud_mmr': np.array([1e-6, 5e-6, 1e-7]), # Non-zero clouds + 'cloud_area': np.array([0.0, 0.5, 0.1]), + } + + # Mock matplotlib + fig_mock = MagicMock() + ax1_mock = MagicMock() + ax2_mock = MagicMock() + mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) + ax1_mock.twiny.return_value = MagicMock() + ax2_mock.twiny.return_value = MagicMock() + + plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) + + # Verify second panel (ax2) was used for cloud plotting + assert ax2_mock.plot.called + assert ax2_mock.set_xlabel.called + assert ax2_mock.set_ylabel.called + + +@pytest.mark.unit +@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') +@patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') +@patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') +@patch('proteus.plot.cpl_chem_atmosphere.glob.glob') +@patch('proteus.plot.cpl_chem_atmosphere.os.path.join') +def test_plot_chem_atmosphere_with_aerosols( + mock_join, mock_glob, mock_read, mock_subplots, mock_savefig +): + """ + Test plot generation including aerosol profiles. + + Physical scenario: aerosols (e.g., sulfate from volcanic outgassing) are + present when aerosols_enabled=True. Should appear in panel 2. + """ + mock_glob.return_value = ['output_dir/data/1000_atm.nc'] + mock_join.return_value = 'output_dir/plots/plot_chem_atmosphere.png' + + # Mock NetCDF with aerosol data + mock_read.return_value = { + 'pl': np.array([1e8, 1e7, 1e6, 1e5]), + 'tmpl': np.array([3000, 2000, 1000, 500]), + 'H2O_vmr': np.array([1e-3, 1e-4, 1e-5]), + 'cloud_mmr': np.array([0.0, 0.0, 0.0]), + 'cloud_area': np.array([0.0, 0.0, 0.0]), + 'aer_mmr': np.array([[1e-8], [5e-8], [1e-9]]), # Total aerosol MMR (1 species) + } + + # Mock matplotlib + fig_mock = MagicMock() + ax1_mock = MagicMock() + ax2_mock = MagicMock() + mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) + ax1_mock.twiny.return_value = MagicMock() + ax2_mock.twiny.return_value = MagicMock() + + plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) + + # Verify aerosol plotting on second panel + assert ax2_mock.plot.called + assert ax2_mock.legend.called + + +@pytest.mark.unit +@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') +@patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') +@patch('proteus.plot.cpl_chem_atmosphere.read_result') +@patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') +@patch('proteus.plot.cpl_chem_atmosphere.glob.glob') +@patch('proteus.plot.cpl_chem_atmosphere.os.path.join') +def test_plot_chem_atmosphere_with_offchem( + mock_join, mock_glob, mock_read, mock_read_result, mock_subplots, mock_savefig +): + """ + Test plot with offline chemistry results. + + Physical scenario: VULCAN or other chemistry module computes detailed + kinetics; results are overlaid as solid lines on top of equilibrium + chemistry (dashed lines). + """ + mock_glob.return_value = ['output_dir/data/100_atm.nc'] + mock_join.return_value = 'output_dir/plots/plot_chem_atmosphere.png' + + # Mock NetCDF read + mock_read.return_value = { + 'pl': np.array([1e8, 1e7, 1e6]), + 'tmpl': np.array([3000, 2000, 1000]), + 'H2O_vmr': np.array([1e-3, 1e-4]), + 'cloud_mmr': np.array([0.0, 0.0]), + 'cloud_area': np.array([0.0, 0.0]), + } + + # Mock offline chemistry data + import pandas as pd + + mock_read_result.return_value = pd.DataFrame( + { + 'H2O': [1e-2, 1e-3], + 'CO2': [1e-3, 1e-4], + 'tmp': [3000, 2000], + 'p': [1e8, 1e7], + 'z': [0, 1e5], + 'Kzz': [1e5, 1e6], + } + ) + + # Mock matplotlib + fig_mock = MagicMock() + ax1_mock = MagicMock() + ax2_mock = MagicMock() + mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) + ax1_mock.twiny.return_value = MagicMock() + ax2_mock.twiny.return_value = MagicMock() + ax1_mock.get_legend_handles_labels.return_value = ([], []) + + plot_chem_atmosphere('output_dir', 'vulcan', plot_format='png', plot_offchem=True) + + # Verify offline chem was read + mock_read_result.assert_called_once() + + # Verify plot was created + mock_savefig.assert_called_once() + + +@pytest.mark.unit +@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') +@patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') +@patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') +@patch('proteus.plot.cpl_chem_atmosphere.glob.glob') +def test_plot_chem_atmosphere_temperature_overlay( + mock_glob, mock_read, mock_subplots, mock_savefig +): + """ + Test that temperature profiles are added to both panels. + + Physical scenario: temperature provides critical context for interpreting + gas abundances (condensation fronts) and cloud/aerosol distributions. + """ + mock_glob.return_value = ['output_dir/data/100_atm.nc'] + + # Mock NetCDF with temperature gradient + mock_read.return_value = { + 'pl': np.array([1e8, 1e7, 1e6, 1e5]), + 'tmpl': np.array([3500, 2000, 800, 300]), # Strong gradient + 'H2O_vmr': np.array([1e-3, 1e-4, 1e-5]), + 'cloud_mmr': np.array([0.0, 0.0, 0.0]), + 'cloud_area': np.array([0.0, 0.0, 0.0]), + } + + # Mock matplotlib + fig_mock = MagicMock() + ax1_mock = MagicMock() + ax2_mock = MagicMock() + mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) + + ax1_temp_mock = MagicMock() + ax2_temp_mock = MagicMock() + ax1_mock.twiny.return_value = ax1_temp_mock + ax2_mock.twiny.return_value = ax2_temp_mock + + plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) + + # Verify temperature was plotted on both panels + ax1_temp_mock.plot.assert_called_once() + ax2_temp_mock.plot.assert_called_once() + + # Verify temperature axes were labeled + ax1_temp_mock.set_xlabel.assert_called_once() + ax2_temp_mock.set_xlabel.assert_called_once() + + +@pytest.mark.unit +@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') +@patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') +@patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') +@patch('proteus.plot.cpl_chem_atmosphere.glob.glob') +def test_plot_chem_atmosphere_time_annotation( + mock_glob, mock_read, mock_subplots, mock_savefig +): + """ + Test that simulation time is annotated on the plot. + + Physical scenario: tracking atmospheric evolution over time requires + timestamping each plot. + """ + # Filename encodes time: 5000_atm.nc = 5000 years + mock_glob.return_value = ['output_dir/data/5000_atm.nc'] + + mock_read.return_value = { + 'pl': np.array([1e8, 1e7]), + 'tmpl': np.array([3000, 1000]), + 'H2O_vmr': np.array([1e-3]), + 'cloud_mmr': np.array([0.0]), + 'cloud_area': np.array([0.0]), + } + + fig_mock = MagicMock() + ax1_mock = MagicMock() + ax2_mock = MagicMock() + mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) + ax1_mock.twiny.return_value = MagicMock() + ax2_mock.twiny.return_value = MagicMock() + + plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) + + # Verify time annotation was added to figure + fig_mock.text.assert_called_once() + call_args = fig_mock.text.call_args[0] + # Should contain the year value (5000) + assert '5.00e+03' in call_args[2] or '5000' in str(call_args) diff --git a/tests/utils/test_data.py b/tests/utils/test_data.py index aaadc4004..ca06554f2 100644 --- a/tests/utils/test_data.py +++ b/tests/utils/test_data.py @@ -1957,6 +1957,21 @@ def test_download_surface_albedos_no_mapping(mock_get_info): with pytest.raises(ValueError, match='No data source mapping found'): download_surface_albedos() +@pytest.mark.unit +@patch('proteus.utils.data.get_data_source_info') +def test_download_scattering_no_mapping(mock_get_info): + """ + Test scattering download raises error when no mapping found. + + Physical scenario: if the DATA_SOURCE_MAP is misconfigured, fail fast + with a clear error rather than silently skipping the download. + """ + from proteus.utils.data import download_scattering + + mock_get_info.return_value = None + + with pytest.raises(ValueError, match='No data source mapping found'): + download_scattering() @pytest.mark.unit @patch('proteus.utils.data.get_data_source_info') From 9ae54c28d37d9629183404f27f64c0749e6df169 Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 16:39:54 +0000 Subject: [PATCH 07/17] Sort aerosol order for consistency --- src/proteus/atmos_clim/agni.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/proteus/atmos_clim/agni.py b/src/proteus/atmos_clim/agni.py index e5d3286dc..1484b583f 100644 --- a/src/proteus/atmos_clim/agni.py +++ b/src/proteus/atmos_clim/agni.py @@ -141,7 +141,8 @@ def _determine_aerosols(dirs: dict) -> list: aerosols = [] for f in os.listdir(scattering_dir): if f.endswith('.mon'): - aerosols.append(f.split('.')[0]) + aerosols.append(f.replace('.mon', '')) + aerosols = sorted(aerosols) log.debug(f'Available aerosols: {aerosols}') return aerosols From ea104238b256bc656da1307837970263604a1c22 Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 17:28:04 +0000 Subject: [PATCH 08/17] Config format --- input/demos/agni.toml | 6 +++--- src/proteus/config/_atmos_clim.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/input/demos/agni.toml b/input/demos/agni.toml index 02185b9aa..491d04797 100644 --- a/input/demos/agni.toml +++ b/input/demos/agni.toml @@ -88,6 +88,9 @@ version = "2.0" albedo_pl = 0.2 # Bond albedo (scattering) module = "agni" # Which atmosphere module to use rayleigh = false + cloud_enabled = false # include water cloud radiative effects? + cloud_alpha = 0.5 # water condensate retention fraction + aerosols_enabled = false # include aerosol radiative effects? [atmos_clim.agni] verbosity = 1 # output verbosity for agni (0:none, 1:info, 2:debug) @@ -104,9 +107,6 @@ version = "2.0" overlap_method = "ee" # gas overlap method surf_roughness = 1e-3 # characteristic surface roughness [m] surf_windspeed = 2.0 # characteristic surface wind speed [m/s]. - cloud_enabled = false # include water cloud radiative effects? - cloud_alpha = 0.5 # water condensate retention fraction - aerosols_enabled = false # include aerosol radiative effects? rainout = true # include volatile condensation/evaporation aloft latent_heat = false # include latent heat release when `rainout=true`? oceans = true # form liquid oceans at planet surface? diff --git a/src/proteus/config/_atmos_clim.py b/src/proteus/config/_atmos_clim.py index afa81b60b..684203f96 100644 --- a/src/proteus/config/_atmos_clim.py +++ b/src/proteus/config/_atmos_clim.py @@ -18,7 +18,7 @@ def tmp_max_bigger_than_tmp_min(instance, attribute, value): def warn_if_dummy(instance, attribute, value): if (instance.module == 'dummy') and value: - raise ValueError('Dummy atmos_clim module is incompatible with Rayleigh scattering') + raise ValueError(f'Dummy atmos_clim module is incompatible with {attribute.name}=True') def check_overlap(instance, attribute, value): From 043738884fd27849fe8a5ba58ddb3656209c17be Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 17:28:40 +0000 Subject: [PATCH 09/17] Ruff --- tests/utils/test_data.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/utils/test_data.py b/tests/utils/test_data.py index ca06554f2..ee72c82f0 100644 --- a/tests/utils/test_data.py +++ b/tests/utils/test_data.py @@ -1957,6 +1957,7 @@ def test_download_surface_albedos_no_mapping(mock_get_info): with pytest.raises(ValueError, match='No data source mapping found'): download_surface_albedos() + @pytest.mark.unit @patch('proteus.utils.data.get_data_source_info') def test_download_scattering_no_mapping(mock_get_info): @@ -1973,6 +1974,7 @@ def test_download_scattering_no_mapping(mock_get_info): with pytest.raises(ValueError, match='No data source mapping found'): download_scattering() + @pytest.mark.unit @patch('proteus.utils.data.get_data_source_info') def test_download_massradius_data_no_mapping(mock_get_info): From e78259f218498fd93ccfaf51a8b991b47938f210 Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 17:30:20 +0000 Subject: [PATCH 10/17] Ruff --- src/proteus/atmos_clim/agni.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/proteus/atmos_clim/agni.py b/src/proteus/atmos_clim/agni.py index 1484b583f..e54af1481 100644 --- a/src/proteus/atmos_clim/agni.py +++ b/src/proteus/atmos_clim/agni.py @@ -228,7 +228,7 @@ def init_agni_atmos(dirs: dict, config: Config, hf_row: dict): p_surf = max(p_surf, p_top * 1.1) # this will happen if the atmosphere is stripped # Aerosol species dictionary (set MMR to zero initially) - aerosol_species = dict() + aerosol_species = {} if config.atmos_clim.aerosols_enabled: aerosol_species = {a: 0.0 for a in _determine_aerosols(dirs)} if len(aerosol_species) == 0: From 6597a0a7cf6559687aa2b883ff877f7a9e3d3d9a Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 18:06:04 +0000 Subject: [PATCH 11/17] Update colours. Update aerosol data on Zenodo. Read aerosol array from netcdf safely --- input/demos/agni.toml | 2 +- src/proteus/atmos_clim/common.py | 32 ++++++++++--- src/proteus/plot/cpl_chem_atmosphere.py | 63 ++++++++++++------------- src/proteus/utils/data.py | 2 +- src/proteus/utils/plot.py | 15 +++++- 5 files changed, 70 insertions(+), 44 deletions(-) diff --git a/input/demos/agni.toml b/input/demos/agni.toml index 491d04797..6c26bcfc2 100644 --- a/input/demos/agni.toml +++ b/input/demos/agni.toml @@ -90,7 +90,7 @@ version = "2.0" rayleigh = false cloud_enabled = false # include water cloud radiative effects? cloud_alpha = 0.5 # water condensate retention fraction - aerosols_enabled = false # include aerosol radiative effects? + aerosols_enabled = true # include aerosol radiative effects? [atmos_clim.agni] verbosity = 1 # output verbosity for agni (0:none, 1:info, 2:debug) diff --git a/src/proteus/atmos_clim/common.py b/src/proteus/atmos_clim/common.py index 5dcb8b264..326c314ca 100644 --- a/src/proteus/atmos_clim/common.py +++ b/src/proteus/atmos_clim/common.py @@ -141,7 +141,15 @@ def read_ncdf_profile(nc_fpath: str, extra_keys: list = [], combine_edges: bool continue # Reading composition - if key == 'x_gas': + if key == 'gases': + gas_l = ds.variables['gases'][:] # names (bytes matrix) + gases = [] + for igas, gas in enumerate(gas_l): + gas_lbl = ''.join([c.decode(encoding='utf-8') for c in gas]).strip() + gases.append(gas_lbl) + out['gases'] = gases + + elif key == 'x_gas': gas_l = ds.variables['gases'][:] # names (bytes matrix) gas_x = ds.variables['x_gas'][:] # vmrs (float matrix) @@ -150,20 +158,27 @@ def read_ncdf_profile(nc_fpath: str, extra_keys: list = [], combine_edges: bool gas_lbl = ''.join([c.decode(encoding='utf-8') for c in gas]).strip() out[gas_lbl + '_vmr'] = np.array(gas_x[:, igas]) - # Reading aerosol mass mixing ratios - elif key == 'aer_mmr': + + elif key == 'aerosols': if 'aerosols' in ds.variables.keys(): + aer_l = ds.variables['aerosols'][:] # names (bytes matrix) + aerosols = [] + for iaer, aer in enumerate(aer_l): + aer_lbl = ''.join([c.decode(encoding='utf-8') for c in aer]).strip() + aerosols.append(aer_lbl) + out['aerosols'] = aerosols + + elif key == 'aer_mmr': + if 'aer_mmr' in ds.variables.keys(): aer_l = ds.variables['aerosols'][:] # names (bytes matrix) aer_x = ds.variables['aer_mmr'][:] # mmrs (float matrix) # get data for each aerosol for iaer, aer in enumerate(aer_l): - if aer[0] is not np.bytes_: + if type(aer[0]) not in [bytes, np.bytes_]: continue aer_lbl = ''.join([c.decode(encoding='utf-8') for c in aer]).strip() out[aer_lbl + '_mmr'] = np.array(aer_x[:, iaer]) - # Also store raw array - out[key] = np.array(ds.variables[key][:]) else: out[key] = np.array(ds.variables[key][:]) @@ -173,7 +188,10 @@ def read_ncdf_profile(nc_fpath: str, extra_keys: list = [], combine_edges: bool # convert to np arrays for key in out.keys(): - out[key] = np.array(out[key], dtype=float) + try: + out[key] = np.array(out[key], dtype=float) + except (AttributeError, TypeError, ValueError): + out[key] = np.array(out[key]) return out diff --git a/src/proteus/plot/cpl_chem_atmosphere.py b/src/proteus/plot/cpl_chem_atmosphere.py index b03c536ad..20ac8cc30 100644 --- a/src/proteus/plot/cpl_chem_atmosphere.py +++ b/src/proteus/plot/cpl_chem_atmosphere.py @@ -92,7 +92,7 @@ def plot_chem_atmosphere( return nc_fpath = natural_sort(files)[-1] atm_profile = read_ncdf_profile( - nc_fpath, extra_keys=['pl', 'tmpl', 'x_gas', 'cloud_mmr', 'cloud_area', 'aer_mmr'] + nc_fpath, extra_keys=['pl', 'tmpl', 'x_gas', 'cloud_mmr', 'aer_mmr', 'aerosols'] ) parr = atm_profile['pl'] * 1e-5 # convert to bar @@ -148,7 +148,7 @@ def plot_chem_atmosphere( # create legend entry and store surface vmr if vmr > 0.0: - ax1.plot(1e30, 1e30, ls='solid', color=col, lw=_lw, alpha=al, label=lbl) + ax1.plot([], [], ls='solid', color=col, lw=_lw, alpha=al, label=lbl) vmr_surf.append(vmr) # Add temperature profile to Panel 1 (secondary x-axis at top) @@ -175,12 +175,12 @@ def plot_chem_atmosphere( [labels[idx] for idx in order], loc='center left', bbox_to_anchor=(1.00, 0.5), - ncols=1, - fontsize=9, + ncols=max(1, len(order)//18+1), + fontsize=8, borderpad=0.0, labelspacing=0.3, columnspacing=1.0, - handlelength=1.0, + handlelength=0.9, handletextpad=0.3, frameon=False, ) @@ -191,34 +191,31 @@ def plot_chem_atmosphere( # Cloud profiles if 'cloud_mmr' in atm_profile.keys(): cloud_mmr = atm_profile['cloud_mmr'] - # Extend to match parr length (edge values) - if np.max(cloud_mmr) > 0: - ax2.plot( - [cloud_mmr[0]] + list(cloud_mmr), - parr, - ls='solid', - color=get_colour('cloud'), - lw=1.5, - alpha=0.7, - label='cloud', - ) + ax2.plot( + [cloud_mmr[0]] + list(cloud_mmr), + parr, + ls='solid', + color=get_colour('cloud'), + lw=1.5, + alpha=0.7, + label='cloud', + ) # Aerosol profiles # Check for individual aerosol species - for key in atm_profile.keys(): - if key.endswith('_mmr') and 'cloud' not in key: - aer_mmr = atm_profile[key] - if np.max(aer_mmr) > 0: - aer_name = key.split('_')[0] - ax2.plot( - [aer_mmr[0]] + list(aer_mmr), - parr, - ls='solid', - lw=1.5, - alpha=0.7, - color=get_colour(aer_name), - label=aer_name, - ) + num_aerosols = 1 + for aer_name in atm_profile['aerosols']: + num_aerosols += 1 + aer_mmr = atm_profile[f'{aer_name}_mmr'] + ax2.plot( + [aer_mmr[0]] + list(aer_mmr), + parr, + ls='solid', + lw=1.5, + alpha=0.7, + color=get_colour(aer_name), + label=aer_name, + ) # Add temperature profile to Panel 2 (secondary x-axis at top) ax2_temp = ax2.twiny() @@ -239,12 +236,12 @@ def plot_chem_atmosphere( ax2.legend( loc='center left', bbox_to_anchor=(1.00, 0.5), - ncols=1, - fontsize=9, + ncols=max(1, num_aerosols//18+1), + fontsize=8, borderpad=0.0, labelspacing=0.3, columnspacing=1.0, - handlelength=1.0, + handlelength=0.9, handletextpad=0.3, frameon=False, ) diff --git a/src/proteus/utils/data.py b/src/proteus/utils/data.py index bac409bf4..e4f35bc85 100644 --- a/src/proteus/utils/data.py +++ b/src/proteus/utils/data.py @@ -479,7 +479,7 @@ def validate_zenodo_folder(zenodo_id: str, folder_dir: Path, hash_maxfilesize=10 # EOS material properties (OSF project: dpkjb) 'EOS_Seager2007': {'zenodo_id': '15727998', 'osf_id': 'dpkjb', 'osf_project': 'dpkjb'}, # Aerosol scattering data (no OSF project) - 'scattering': {'zenodo_id': '19232950'}, + 'scattering': {'zenodo_id': '19294180'}, } diff --git a/src/proteus/utils/plot.py b/src/proteus/utils/plot.py index 5b6d3b0d3..b876bc638 100644 --- a/src/proteus/utils/plot.py +++ b/src/proteus/utils/plot.py @@ -67,12 +67,23 @@ 'cor_bkg': '#efefef', # Clouds and aerosols 'cloud': '#027FB1', - 'freshsoot': '#876e44', + 'frsoot': '#876e44', 'soot': '#63401a', - 'agedsoot': '#251607', + 'agsoot': '#251607', 'nitrate': '#ffaa00', 'ash': '#970505', 'sulph': '#ff22ff', + 'biogenic': '#33ccff', + 'bioms1': '#33ccaa', + 'delta' : "#6a9064", + 'dustdiv1' : "#fcfcd7", + 'dustdiv2' : "#fbfca1", + 'dustdiv3' : "#f6f873", + 'dustdiv4' : "#fcff5d", + 'dustdiv5' : "#fcff33", + 'dustdiv6' : "#fbff21", + 'naclflm' : "#fca1ff", + 'nacljet' : "#ee29f5" } From 76757fd69b16be9e508c13caf283066d4c1f414a Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 18:06:22 +0000 Subject: [PATCH 12/17] Format --- src/proteus/atmos_clim/common.py | 1 - src/proteus/plot/cpl_chem_atmosphere.py | 4 ++-- src/proteus/utils/plot.py | 18 +++++++++--------- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/src/proteus/atmos_clim/common.py b/src/proteus/atmos_clim/common.py index 326c314ca..8cc8aa020 100644 --- a/src/proteus/atmos_clim/common.py +++ b/src/proteus/atmos_clim/common.py @@ -158,7 +158,6 @@ def read_ncdf_profile(nc_fpath: str, extra_keys: list = [], combine_edges: bool gas_lbl = ''.join([c.decode(encoding='utf-8') for c in gas]).strip() out[gas_lbl + '_vmr'] = np.array(gas_x[:, igas]) - elif key == 'aerosols': if 'aerosols' in ds.variables.keys(): aer_l = ds.variables['aerosols'][:] # names (bytes matrix) diff --git a/src/proteus/plot/cpl_chem_atmosphere.py b/src/proteus/plot/cpl_chem_atmosphere.py index 20ac8cc30..5b4997888 100644 --- a/src/proteus/plot/cpl_chem_atmosphere.py +++ b/src/proteus/plot/cpl_chem_atmosphere.py @@ -175,7 +175,7 @@ def plot_chem_atmosphere( [labels[idx] for idx in order], loc='center left', bbox_to_anchor=(1.00, 0.5), - ncols=max(1, len(order)//18+1), + ncols=max(1, len(order) // 18 + 1), fontsize=8, borderpad=0.0, labelspacing=0.3, @@ -236,7 +236,7 @@ def plot_chem_atmosphere( ax2.legend( loc='center left', bbox_to_anchor=(1.00, 0.5), - ncols=max(1, num_aerosols//18+1), + ncols=max(1, num_aerosols // 18 + 1), fontsize=8, borderpad=0.0, labelspacing=0.3, diff --git a/src/proteus/utils/plot.py b/src/proteus/utils/plot.py index b876bc638..583e7d42c 100644 --- a/src/proteus/utils/plot.py +++ b/src/proteus/utils/plot.py @@ -75,15 +75,15 @@ 'sulph': '#ff22ff', 'biogenic': '#33ccff', 'bioms1': '#33ccaa', - 'delta' : "#6a9064", - 'dustdiv1' : "#fcfcd7", - 'dustdiv2' : "#fbfca1", - 'dustdiv3' : "#f6f873", - 'dustdiv4' : "#fcff5d", - 'dustdiv5' : "#fcff33", - 'dustdiv6' : "#fbff21", - 'naclflm' : "#fca1ff", - 'nacljet' : "#ee29f5" + 'delta': '#6a9064', + 'dustdiv1': '#fcfcd7', + 'dustdiv2': '#fbfca1', + 'dustdiv3': '#f6f873', + 'dustdiv4': '#fcff5d', + 'dustdiv5': '#fcff33', + 'dustdiv6': '#fbff21', + 'naclflm': '#fca1ff', + 'nacljet': '#ee29f5', } From 46c2d63233b5ec14f30be9b01b7ab042fa0da88b Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 19:01:41 +0000 Subject: [PATCH 13/17] Add more tests for patch. --- src/proteus/cli.py | 1 - tests/atmos_clim/test_agni.py | 175 +++++++++++++++++++++++++ tests/atmos_clim/test_common.py | 154 ++++++++++++++++++++-- tests/plot/test_cpl_chem_atmosphere.py | 111 +++++++++++++++- 4 files changed, 424 insertions(+), 17 deletions(-) create mode 100644 tests/atmos_clim/test_agni.py diff --git a/src/proteus/cli.py b/src/proteus/cli.py index 921ebce3e..e1c39ccc0 100644 --- a/src/proteus/cli.py +++ b/src/proteus/cli.py @@ -101,7 +101,6 @@ def plot(plots, config_path: Path): for plot in plots: if plot not in plot_dispatch.keys(): click.echo(f'Invalid plot: {plot}') - click.echo('Available plots: ' + ', '.join(plot_dispatch.keys())) else: click.echo(f'Plotting: {plot}') plot_func = plot_dispatch[plot] diff --git a/tests/atmos_clim/test_agni.py b/tests/atmos_clim/test_agni.py new file mode 100644 index 000000000..2b9888460 --- /dev/null +++ b/tests/atmos_clim/test_agni.py @@ -0,0 +1,175 @@ +""" +Unit tests for proteus.atmos_clim.agni module. + +This module tests the AGNI atmosphere interface including: +- Aerosol discovery (_determine_aerosols) +- Condensate species determination (_determine_condensates) +- AGNI atmosphere initialization (init_agni_atmos) + +See also: +- docs/How-to/test_infrastructure.md +- docs/How-to/test_categorization.md +- docs/How-to/test_building.md +""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from proteus.atmos_clim.agni import _determine_aerosols, _determine_condensates + + +@pytest.mark.unit +@patch('proteus.atmos_clim.agni.os.listdir') +@patch('proteus.atmos_clim.agni.os.path.isdir') +def test_determine_aerosols_success(mock_isdir, mock_listdir): + """ + Test aerosol discovery when scattering data directory exists. + + Physical scenario: Scattering data for aerosols (e.g., sulfate, silicate) + is available in FWL_DATA/scattering/scattering/*.mon files. + """ + mock_isdir.return_value = True + mock_listdir.return_value = [ + 'Sulfate.mon', + 'Silicate.mon', + 'Haze.mon', + 'other_file.txt', # Should be ignored + 'readme.md', # Should be ignored + ] + + dirs = {'fwl': '/fake/fwl/path'} + aerosols = _determine_aerosols(dirs) + + # Verify correct aerosols found and sorted + assert len(aerosols) == 3 + assert aerosols == ['Haze', 'Silicate', 'Sulfate'] # alphabetically sorted + + # Verify correct directory was checked + mock_isdir.assert_called_once_with('/fake/fwl/path/scattering/scattering') + + +@pytest.mark.unit +@patch('proteus.atmos_clim.agni.os.path.isdir') +def test_determine_aerosols_missing_directory(mock_isdir): + """ + Test aerosol discovery when scattering directory doesn't exist. + + Physical scenario: FWL_DATA not properly downloaded or scattering + data not installed. Should return empty list and warn. + """ + mock_isdir.return_value = False + + dirs = {'fwl': '/nonexistent/path'} + aerosols = _determine_aerosols(dirs) + + # Should return empty list without crashing + assert aerosols == [] + mock_isdir.assert_called_once() + + +@pytest.mark.unit +@patch('proteus.atmos_clim.agni.os.listdir') +@patch('proteus.atmos_clim.agni.os.path.isdir') +def test_determine_aerosols_empty_directory(mock_isdir, mock_listdir): + """ + Test aerosol discovery when directory exists but has no .mon files. + + Physical scenario: Scattering directory present but empty or only + contains non-aerosol files. + """ + mock_isdir.return_value = True + mock_listdir.return_value = ['readme.txt', 'config.yaml'] + + dirs = {'fwl': '/path/to/fwl'} + aerosols = _determine_aerosols(dirs) + + # Should return empty list + assert aerosols == [] + + +@pytest.mark.unit +@patch('proteus.atmos_clim.agni.os.listdir') +@patch('proteus.atmos_clim.agni.os.path.isdir') +def test_determine_aerosols_single_species(mock_isdir, mock_listdir): + """ + Test aerosol discovery with only one aerosol type. + + Physical scenario: Limited scattering data with only one aerosol species + available (e.g., only sulfate aerosols). + """ + mock_isdir.return_value = True + mock_listdir.return_value = ['Sulfate.mon'] + + dirs = {'fwl': '/path/to/fwl'} + aerosols = _determine_aerosols(dirs) + + assert len(aerosols) == 1 + assert aerosols == ['Sulfate'] + + +@pytest.mark.unit +def test_determine_condensates(): + """ + Test condensate species determination from volatile list. + + Physical scenario: Given a list of volatile species, filter out + those that are always dry (H2, N2, CO) to get condensable + species like H2O, CO2, He, CH4, etc. + """ + # Test with mixed list of dry and condensable species + vol_list = ['H2O', 'CO2', 'N2', 'CH4', 'He', 'H2', 'CO'] + condensates = _determine_condensates(vol_list) + + # N2, H2, CO should be filtered out (always dry) + assert 'H2O' in condensates + assert 'CO2' in condensates + assert 'CH4' in condensates + assert 'He' in condensates # He is condensable in AGNI + assert 'N2' not in condensates + assert 'H2' not in condensates + assert 'CO' not in condensates + + +@pytest.mark.unit +def test_determine_condensates_all_dry(): + """ + Test condensate determination with only dry species. + + Physical scenario: Hydrogen-nitrogen-CO dominated atmosphere with no + condensable species. + """ + vol_list = ['H2', 'N2', 'CO'] + condensates = _determine_condensates(vol_list) + + # Should return empty list + assert condensates == [] + + +@pytest.mark.unit +def test_determine_condensates_all_condensable(): + """ + Test condensate determination with all condensable species. + + Physical scenario: Rocky planet atmosphere with water, CO2, and other + condensable volatiles but no hydrogen/helium. + """ + vol_list = ['H2O', 'CO2', 'NH3', 'CH4', 'SO2'] + condensates = _determine_condensates(vol_list) + + # All should remain + assert len(condensates) == len(vol_list) + assert set(condensates) == set(vol_list) + + +@pytest.mark.unit +def test_determine_condensates_empty_list(): + """ + Test condensate determination with empty volatile list. + + Physical scenario: Edge case where no volatiles are specified. + """ + condensates = _determine_condensates([]) + assert condensates == [] diff --git a/tests/atmos_clim/test_common.py b/tests/atmos_clim/test_common.py index 2542ef200..65f1ed98c 100644 --- a/tests/atmos_clim/test_common.py +++ b/tests/atmos_clim/test_common.py @@ -171,7 +171,7 @@ def test_read_ncdf_profile_with_aerosols(mock_ds, mock_isfile): ds_instance = MagicMock() mock_ds.return_value = ds_instance - # Mock basic atmospheric profile + # Mock basic atmospheric profile with two aerosol species ds_instance.variables = { 'p': np.array([100.0]), 'pl': np.array([110.0, 90.0]), @@ -181,7 +181,7 @@ def test_read_ncdf_profile_with_aerosols(mock_ds, mock_isfile): 'rl': np.array([6.3e6, 6.5e6]), 'planet_radius': [6.0e6], 'transparent': np.array([b'y'], dtype='S1'), - # Aerosol data + # Aerosol data: 2 species (Sulfate, Silicate) at 1 level each 'aerosols': np.array( [ [ @@ -223,23 +223,157 @@ def test_read_ncdf_profile_with_aerosols(mock_ds, mock_isfile): ], dtype='S1', ), - 'aer_mmr': np.array([[1e-6], [2e-6]]), # 2 levels, 1 aerosol type + 'aer_mmr': np.array([[1e-6, 2e-6]]), # 1 level, 2 aerosol species } # Read with aerosol data - result = read_ncdf_profile('dummy.nc', extra_keys=['aer_mmr']) + result = read_ncdf_profile('dummy.nc', extra_keys=['aer_mmr', 'aerosols']) # Verify basic profile data assert 'p' in result assert 'transparent' in result - # Verify aerosol data was read - assert 'aer_mmr' in result - np.testing.assert_allclose(result['aer_mmr'], np.array([[1e-6], [2e-6]])) + # Verify aerosol list was read (stored as numpy array) + assert 'aerosols' in result + assert len(result['aerosols']) == 2 + assert 'Silicate' in result['aerosols'] + assert 'Sulfate' in result['aerosols'] - # Note: Named aerosol extraction happens only if names are valid (not masked) - # With mock setup above, the names should parse - # But the actual extraction logic might skip empty/invalid names + # Verify individual aerosol MMRs were extracted + assert 'Sulfate_mmr' in result + assert 'Silicate_mmr' in result + np.testing.assert_allclose(result['Sulfate_mmr'], np.array([1e-6])) + np.testing.assert_allclose(result['Silicate_mmr'], np.array([2e-6])) + + +@pytest.mark.unit +@patch('proteus.atmos_clim.common.os.path.isfile') +@patch('netCDF4.Dataset') +def test_read_ncdf_profile_gases_list(mock_ds, mock_isfile): + """ + Test reading list of gas species names without VMRs. + + Physical scenario: When reading just the gas species present in the + atmosphere without their mixing ratios (useful for metadata queries). + """ + mock_isfile.return_value = True + + ds_instance = MagicMock() + mock_ds.return_value = ds_instance + + # Mock gas names only + ds_instance.variables = { + 'p': np.array([100.0]), + 'pl': np.array([110.0, 90.0]), + 'tmp': np.array([300.0]), + 'tmpl': np.array([310.0, 290.0]), + 'r': np.array([6.4e6]), + 'rl': np.array([6.3e6, 6.5e6]), + 'planet_radius': [6.0e6], + 'solved': np.array([b'y'], dtype='S1'), + # Gas species names (3 gases: H2O, CO2, N2) + 'gases': np.array( + [ + [b'H', b'2', b'O', b' ', b' ', b' '], + [b'C', b'O', b'2', b' ', b' ', b' '], + [b'N', b'2', b' ', b' ', b' ', b' '], + ], + dtype='S1', + ), + } + + # Read with gases key + result = read_ncdf_profile('dummy.nc', extra_keys=['gases']) + + # Verify gas list was read and parsed correctly + assert 'gases' in result + assert len(result['gases']) == 3 + assert 'H2O' in result['gases'] + assert 'CO2' in result['gases'] + assert 'N2' in result['gases'] + + +@pytest.mark.unit +@patch('proteus.atmos_clim.common.os.path.isfile') +@patch('netCDF4.Dataset') +def test_read_ncdf_profile_aerosols_list_only(mock_ds, mock_isfile): + """ + Test reading list of aerosol species names without MMRs. + + Physical scenario: When checking which aerosol types are available + in the simulation output without loading full profiles. + """ + mock_isfile.return_value = True + + ds_instance = MagicMock() + mock_ds.return_value = ds_instance + + ds_instance.variables = { + 'p': np.array([100.0]), + 'pl': np.array([110.0, 90.0]), + 'tmp': np.array([300.0]), + 'tmpl': np.array([310.0, 290.0]), + 'r': np.array([6.4e6]), + 'rl': np.array([6.3e6, 6.5e6]), + 'planet_radius': [6.0e6], + 'solved': np.array([b'n'], dtype='S1'), + # Aerosol species names only + 'aerosols': np.array( + [ + [b'S', b'u', b'l', b'f', b'a', b't', b'e', b' '], + [b'H', b'a', b'z', b'e', b' ', b' ', b' ', b' '], + ], + dtype='S1', + ), + } + + # Read with aerosols key + result = read_ncdf_profile('dummy.nc', extra_keys=['aerosols']) + + # Verify aerosol list was read + assert 'aerosols' in result + assert len(result['aerosols']) == 2 + assert 'Sulfate' in result['aerosols'] + assert 'Haze' in result['aerosols'] + + +@pytest.mark.unit +@patch('proteus.atmos_clim.common.os.path.isfile') +@patch('netCDF4.Dataset') +def test_read_ncdf_profile_no_aerosols_in_file(mock_ds, mock_isfile): + """ + Test reading profile when aerosols key is requested but not present. + + Physical scenario: Attempting to read aerosol data from a simulation + run without aerosols_enabled=True. Should handle gracefully. + """ + mock_isfile.return_value = True + + ds_instance = MagicMock() + mock_ds.return_value = ds_instance + + # No aerosol variables in dataset + ds_instance.variables = { + 'p': np.array([100.0]), + 'pl': np.array([110.0, 90.0]), + 'tmp': np.array([300.0]), + 'tmpl': np.array([310.0, 290.0]), + 'r': np.array([6.4e6]), + 'rl': np.array([6.3e6, 6.5e6]), + 'planet_radius': [6.0e6], + 'solved': np.array([b'y'], dtype='S1'), + } + + # Read with aerosols/aer_mmr keys (should handle missing gracefully) + result = read_ncdf_profile('dummy.nc', extra_keys=['aerosols', 'aer_mmr']) + + # Verify profile data is still read + assert 'p' in result + assert 'solved' in result + + # Missing keys should not be in result + assert 'aerosols' not in result + assert 'aer_mmr' not in result @pytest.mark.unit diff --git a/tests/plot/test_cpl_chem_atmosphere.py b/tests/plot/test_cpl_chem_atmosphere.py index 7b135269b..712994b58 100644 --- a/tests/plot/test_cpl_chem_atmosphere.py +++ b/tests/plot/test_cpl_chem_atmosphere.py @@ -75,6 +75,9 @@ def test_plot_chem_atmosphere_basic( ax1_mock.twiny.return_value = ax1_temp_mock ax2_mock.twiny.return_value = ax2_temp_mock + # Mock get_legend_handles_labels to return empty lists + ax1_mock.get_legend_handles_labels.return_value = ([], []) + # Run function plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) @@ -135,6 +138,7 @@ def test_plot_chem_atmosphere_with_clouds( mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) ax1_mock.twiny.return_value = MagicMock() ax2_mock.twiny.return_value = MagicMock() + ax1_mock.get_legend_handles_labels.return_value = ([], []) plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) @@ -154,22 +158,24 @@ def test_plot_chem_atmosphere_with_aerosols( mock_join, mock_glob, mock_read, mock_subplots, mock_savefig ): """ - Test plot generation including aerosol profiles. + Test plot generation including aerosol profiles with species names. - Physical scenario: aerosols (e.g., sulfate from volcanic outgassing) are - present when aerosols_enabled=True. Should appear in panel 2. + Physical scenario: aerosols (e.g., sulfate, silicate from volcanic outgassing) + are present when aerosols_enabled=True. Should appear in panel 2 with proper + labels for each aerosol species. """ mock_glob.return_value = ['output_dir/data/1000_atm.nc'] mock_join.return_value = 'output_dir/plots/plot_chem_atmosphere.png' - # Mock NetCDF with aerosol data + # Mock NetCDF with multiple aerosol species mock_read.return_value = { 'pl': np.array([1e8, 1e7, 1e6, 1e5]), 'tmpl': np.array([3000, 2000, 1000, 500]), 'H2O_vmr': np.array([1e-3, 1e-4, 1e-5]), 'cloud_mmr': np.array([0.0, 0.0, 0.0]), - 'cloud_area': np.array([0.0, 0.0, 0.0]), - 'aer_mmr': np.array([[1e-8], [5e-8], [1e-9]]), # Total aerosol MMR (1 species) + 'aerosols': ['Sulfate', 'Silicate'], # Two aerosol species + 'Sulfate_mmr': np.array([1e-8, 5e-8, 1e-9]), # Individual aerosol MMRs + 'Silicate_mmr': np.array([2e-9, 3e-9, 1e-10]), } # Mock matplotlib @@ -179,6 +185,7 @@ def test_plot_chem_atmosphere_with_aerosols( mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) ax1_mock.twiny.return_value = MagicMock() ax2_mock.twiny.return_value = MagicMock() + ax1_mock.get_legend_handles_labels.return_value = ([], []) plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) @@ -186,6 +193,10 @@ def test_plot_chem_atmosphere_with_aerosols( assert ax2_mock.plot.called assert ax2_mock.legend.called + # Verify both aerosol species were plotted (2 species + maybe temperature) + # Each aerosol calls plot once + assert ax2_mock.plot.call_count >= 2 + @pytest.mark.unit @patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') @@ -283,6 +294,7 @@ def test_plot_chem_atmosphere_temperature_overlay( ax2_temp_mock = MagicMock() ax1_mock.twiny.return_value = ax1_temp_mock ax2_mock.twiny.return_value = ax2_temp_mock + ax1_mock.get_legend_handles_labels.return_value = ([], []) plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) @@ -326,6 +338,7 @@ def test_plot_chem_atmosphere_time_annotation( mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) ax1_mock.twiny.return_value = MagicMock() ax2_mock.twiny.return_value = MagicMock() + ax1_mock.get_legend_handles_labels.return_value = ([], []) plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) @@ -334,3 +347,89 @@ def test_plot_chem_atmosphere_time_annotation( call_args = fig_mock.text.call_args[0] # Should contain the year value (5000) assert '5.00e+03' in call_args[2] or '5000' in str(call_args) + + +@pytest.mark.unit +@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') +@patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') +@patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') +@patch('proteus.plot.cpl_chem_atmosphere.glob.glob') +def test_plot_chem_atmosphere_clouds_and_aerosols( + mock_glob, mock_read, mock_subplots, mock_savefig +): + """ + Test plotting with both clouds and aerosols present. + + Physical scenario: Atmosphere with both water clouds and aerosol particles + (e.g., sulfate aerosols in a Venus-like atmosphere with H2SO4 clouds). + """ + mock_glob.return_value = ['output_dir/data/200_atm.nc'] + + # Mock with both cloud and aerosol data + mock_read.return_value = { + 'pl': np.array([1e8, 1e7, 1e6, 1e5]), + 'tmpl': np.array([3000, 2000, 1000, 500]), + 'H2O_vmr': np.array([1e-3, 1e-4, 1e-5]), + 'CO2_vmr': np.array([1e-2, 1e-3, 1e-4]), + 'cloud_mmr': np.array([1e-6, 5e-6, 1e-7]), # Cloud present + 'aerosols': ['Sulfate'], + 'Sulfate_mmr': np.array([1e-8, 2e-8, 5e-9]), # Aerosol present + } + + fig_mock = MagicMock() + ax1_mock = MagicMock() + ax2_mock = MagicMock() + mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) + ax1_mock.twiny.return_value = MagicMock() + ax2_mock.twiny.return_value = MagicMock() + ax1_mock.get_legend_handles_labels.return_value = ([], []) + + plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) + + # Verify both panels were plotted + assert ax1_mock.plot.called # Gas species + assert ax2_mock.plot.called # Clouds + aerosols + + # Panel 2 should have plots for cloud + aerosol (at least 2) + assert ax2_mock.plot.call_count >= 2 + + +@pytest.mark.unit +@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') +@patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') +@patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') +@patch('proteus.plot.cpl_chem_atmosphere.glob.glob') +def test_plot_chem_atmosphere_no_aerosols_key( + mock_glob, mock_read, mock_subplots, mock_savefig +): + """ + Test plotting when aerosols key is missing (backward compatibility). + + Physical scenario: Reading output from older simulation without aerosol + support. Should handle gracefully and plot other data. + """ + mock_glob.return_value = ['output_dir/data/100_atm.nc'] + + # Mock without aerosols key (older output format) + mock_read.return_value = { + 'pl': np.array([1e8, 1e7, 1e6]), + 'tmpl': np.array([3000, 2000, 1000]), + 'H2O_vmr': np.array([1e-3, 1e-4]), + 'cloud_mmr': np.array([0.0, 0.0]), + # No 'aerosols' key at all + } + + fig_mock = MagicMock() + ax1_mock = MagicMock() + ax2_mock = MagicMock() + mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) + ax1_mock.twiny.return_value = MagicMock() + ax2_mock.twiny.return_value = MagicMock() + ax1_mock.get_legend_handles_labels.return_value = ([], []) + + # Should not crash + plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) + + # Verify gas panel was still plotted + assert ax1_mock.plot.called + mock_savefig.assert_called_once() From 190f0b2fc7eba039e15856719a619188ec2ba88e Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sat, 28 Mar 2026 19:08:33 +0000 Subject: [PATCH 14/17] More tests --- tests/atmos_clim/test_common.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/atmos_clim/test_common.py b/tests/atmos_clim/test_common.py index 65f1ed98c..20a239d75 100644 --- a/tests/atmos_clim/test_common.py +++ b/tests/atmos_clim/test_common.py @@ -89,10 +89,17 @@ def test_read_ncdf_profile(mock_ds, mock_isfile): 'rl': np.array([6.3e6, 6.5e6]), 'planet_radius': [6.0e6], 'transparent': np.array([b'y'], dtype='S1'), + 'gases': np.array([[b'H', b'2', b'O'], [b'C', b'O', b'2']], dtype='S1'), + 'x_gases': np.array([0.1, 0.9]), + 'aerosols': np.array([[b's', b'o', b'o', b't'], [b's', b'u', b'l', b'f']], dtype='S1'), + 'aer_mmr': np.array([[1e-6, 2e-6]]), + 'cloud_mmr': np.array([1e-5]), } # Run function - result = read_ncdf_profile('dummy.nc') + result = read_ncdf_profile( + 'dummy.nc', extra_keys=['gases', 'x_gases', 'aerosols', 'aer_mmr', 'cloud_mmr'] + ) # Verify values are correctly extracted assert result['p'][0] == 110.0 # Should match first element of pl From 24a9aae8276540f0fc9a9468907ee23065d6b54f Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sun, 29 Mar 2026 17:46:26 +0100 Subject: [PATCH 15/17] PR comments addressed --- src/proteus/atmos_clim/common.py | 8 +++---- src/proteus/plot/cpl_chem_atmosphere.py | 29 +++++++++++++------------ src/proteus/utils/data.py | 3 ++- 3 files changed, 21 insertions(+), 19 deletions(-) diff --git a/src/proteus/atmos_clim/common.py b/src/proteus/atmos_clim/common.py index 8cc8aa020..5fb95e855 100644 --- a/src/proteus/atmos_clim/common.py +++ b/src/proteus/atmos_clim/common.py @@ -164,7 +164,8 @@ def read_ncdf_profile(nc_fpath: str, extra_keys: list = [], combine_edges: bool aerosols = [] for iaer, aer in enumerate(aer_l): aer_lbl = ''.join([c.decode(encoding='utf-8') for c in aer]).strip() - aerosols.append(aer_lbl) + if len(aer_lbl) > 0: + aerosols.append(aer_lbl) out['aerosols'] = aerosols elif key == 'aer_mmr': @@ -174,10 +175,9 @@ def read_ncdf_profile(nc_fpath: str, extra_keys: list = [], combine_edges: bool # get data for each aerosol for iaer, aer in enumerate(aer_l): - if type(aer[0]) not in [bytes, np.bytes_]: - continue aer_lbl = ''.join([c.decode(encoding='utf-8') for c in aer]).strip() - out[aer_lbl + '_mmr'] = np.array(aer_x[:, iaer]) + if len(aer_lbl) > 0: + out[aer_lbl + '_mmr'] = np.array(aer_x[:, iaer]) else: out[key] = np.array(ds.variables[key][:]) diff --git a/src/proteus/plot/cpl_chem_atmosphere.py b/src/proteus/plot/cpl_chem_atmosphere.py index 5b4997888..cbd656461 100644 --- a/src/proteus/plot/cpl_chem_atmosphere.py +++ b/src/proteus/plot/cpl_chem_atmosphere.py @@ -203,19 +203,20 @@ def plot_chem_atmosphere( # Aerosol profiles # Check for individual aerosol species - num_aerosols = 1 - for aer_name in atm_profile['aerosols']: - num_aerosols += 1 - aer_mmr = atm_profile[f'{aer_name}_mmr'] - ax2.plot( - [aer_mmr[0]] + list(aer_mmr), - parr, - ls='solid', - lw=1.5, - alpha=0.7, - color=get_colour(aer_name), - label=aer_name, - ) + num_aerosols = 0 + if 'aerosols' in atm_profile.keys(): + for aer_name in atm_profile['aerosols']: + num_aerosols += 1 + aer_mmr = atm_profile[f'{aer_name}_mmr'] + ax2.plot( + [aer_mmr[0]] + list(aer_mmr), + parr, + ls='solid', + lw=1.5, + alpha=0.7, + color=get_colour(aer_name), + label=aer_name, + ) # Add temperature profile to Panel 2 (secondary x-axis at top) ax2_temp = ax2.twiny() @@ -247,7 +248,7 @@ def plot_chem_atmosphere( ) # Add time annotation - fig.text( + ax1.text( 0.02, 0.98, f't = {year:.2e} yr', diff --git a/src/proteus/utils/data.py b/src/proteus/utils/data.py index e4f35bc85..feadd6e22 100644 --- a/src/proteus/utils/data.py +++ b/src/proteus/utils/data.py @@ -479,7 +479,7 @@ def validate_zenodo_folder(zenodo_id: str, folder_dir: Path, hash_maxfilesize=10 # EOS material properties (OSF project: dpkjb) 'EOS_Seager2007': {'zenodo_id': '15727998', 'osf_id': 'dpkjb', 'osf_project': 'dpkjb'}, # Aerosol scattering data (no OSF project) - 'scattering': {'zenodo_id': '19294180'}, + 'scattering': {'zenodo_id': '19294180', 'osf_id': 'vehxg', 'osf_project': 'vehxg'}, } @@ -1013,6 +1013,7 @@ def download_scattering(): download( folder=folder, target='scattering', + osf_id=source_info['osf_project'], zenodo_id=source_info['zenodo_id'], desc='radiative properties scattering data', ) From ae451f2f458562739dbc0570e56e58a7e71e2157 Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sun, 29 Mar 2026 19:08:15 +0100 Subject: [PATCH 16/17] fix: address PR#665 review comments - fix test mocks and validators - Fix warn_if_dummy validator tests to pass attribute.name correctly - Fix plot tests: mock legend handles matching actual gas species - Fix plot tests: assert on fig_mock.savefig instead of plt.savefig - Fix time annotation test: assert on ax1_mock.text not fig_mock.text - Remove incorrect ax2_temp xlabel assertion (label is not set in code) Resolves all 3 issues identified in PR#665 review comments. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/config/test_config.py | 8 ++- tests/plot/test_cpl_chem_atmosphere.py | 91 ++++++++++++++++---------- 2 files changed, 62 insertions(+), 37 deletions(-) diff --git a/tests/config/test_config.py b/tests/config/test_config.py index 10b0ef65f..1967eb4b0 100644 --- a/tests/config/test_config.py +++ b/tests/config/test_config.py @@ -751,7 +751,8 @@ def test_atmos_clim_warn_if_dummy_rayleigh_compatible(): # Valid: AGNI module with rayleigh enabled instance = SimpleNamespace(module='agni') - warn_if_dummy(instance, SimpleNamespace(), True) # Should not raise + attribute = SimpleNamespace(name='rayleigh') + warn_if_dummy(instance, attribute, True) # Should not raise @pytest.mark.unit @@ -761,10 +762,11 @@ def test_atmos_clim_warn_if_dummy_rayleigh_incompatible(): # Invalid: dummy module with rayleigh enabled instance = SimpleNamespace(module='dummy') + attribute = SimpleNamespace(name='rayleigh') with pytest.raises( - ValueError, match='Dummy atmos_clim module is incompatible with Rayleigh scattering' + ValueError, match='Dummy atmos_clim module is incompatible with rayleigh=True' ): - warn_if_dummy(instance, SimpleNamespace(), True) + warn_if_dummy(instance, attribute, True) @pytest.mark.unit diff --git a/tests/plot/test_cpl_chem_atmosphere.py b/tests/plot/test_cpl_chem_atmosphere.py index 712994b58..bb1261a37 100644 --- a/tests/plot/test_cpl_chem_atmosphere.py +++ b/tests/plot/test_cpl_chem_atmosphere.py @@ -36,13 +36,12 @@ def test_plot_chem_atmosphere_no_files(mock_glob): @pytest.mark.unit -@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') @patch('proteus.plot.cpl_chem_atmosphere.os.path.join') def test_plot_chem_atmosphere_basic( - mock_join, mock_glob, mock_read, mock_subplots, mock_savefig + mock_join, mock_glob, mock_read, mock_subplots ): """ Test basic plot generation with gas species only. @@ -75,8 +74,11 @@ def test_plot_chem_atmosphere_basic( ax1_mock.twiny.return_value = ax1_temp_mock ax2_mock.twiny.return_value = ax2_temp_mock - # Mock get_legend_handles_labels to return empty lists - ax1_mock.get_legend_handles_labels.return_value = ([], []) + # Mock get_legend_handles_labels to return handles for the gas species + # Need 3 handles for H2O, CO2, N2 (all have nonzero VMR) + mock_handles = [MagicMock(), MagicMock(), MagicMock()] + mock_labels = ['H2O', 'CO2', 'N2'] + ax1_mock.get_legend_handles_labels.return_value = (mock_handles, mock_labels) # Run function plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) @@ -101,17 +103,16 @@ def test_plot_chem_atmosphere_basic( ax2_mock.twiny.assert_called_once() # Verify plot was saved - mock_savefig.assert_called_once() + fig_mock.savefig.assert_called_once() @pytest.mark.unit -@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') @patch('proteus.plot.cpl_chem_atmosphere.os.path.join') def test_plot_chem_atmosphere_with_clouds( - mock_join, mock_glob, mock_read, mock_subplots, mock_savefig + mock_join, mock_glob, mock_read, mock_subplots ): """ Test plot generation including cloud profiles. @@ -138,7 +139,11 @@ def test_plot_chem_atmosphere_with_clouds( mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) ax1_mock.twiny.return_value = MagicMock() ax2_mock.twiny.return_value = MagicMock() - ax1_mock.get_legend_handles_labels.return_value = ([], []) + + # Mock get_legend_handles_labels - only H2O has nonzero VMR + mock_handles = [MagicMock()] + mock_labels = ['H2O'] + ax1_mock.get_legend_handles_labels.return_value = (mock_handles, mock_labels) plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) @@ -149,13 +154,12 @@ def test_plot_chem_atmosphere_with_clouds( @pytest.mark.unit -@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') @patch('proteus.plot.cpl_chem_atmosphere.os.path.join') def test_plot_chem_atmosphere_with_aerosols( - mock_join, mock_glob, mock_read, mock_subplots, mock_savefig + mock_join, mock_glob, mock_read, mock_subplots ): """ Test plot generation including aerosol profiles with species names. @@ -185,7 +189,11 @@ def test_plot_chem_atmosphere_with_aerosols( mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) ax1_mock.twiny.return_value = MagicMock() ax2_mock.twiny.return_value = MagicMock() - ax1_mock.get_legend_handles_labels.return_value = ([], []) + + # Mock get_legend_handles_labels - only H2O has nonzero VMR + mock_handles = [MagicMock()] + mock_labels = ['H2O'] + ax1_mock.get_legend_handles_labels.return_value = (mock_handles, mock_labels) plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) @@ -199,14 +207,13 @@ def test_plot_chem_atmosphere_with_aerosols( @pytest.mark.unit -@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_result') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') @patch('proteus.plot.cpl_chem_atmosphere.os.path.join') def test_plot_chem_atmosphere_with_offchem( - mock_join, mock_glob, mock_read, mock_read_result, mock_subplots, mock_savefig + mock_join, mock_glob, mock_read, mock_read_result, mock_subplots ): """ Test plot with offline chemistry results. @@ -248,24 +255,27 @@ def test_plot_chem_atmosphere_with_offchem( mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) ax1_mock.twiny.return_value = MagicMock() ax2_mock.twiny.return_value = MagicMock() - ax1_mock.get_legend_handles_labels.return_value = ([], []) + + # Mock get_legend_handles_labels - H2O and CO2 have nonzero VMR + mock_handles = [MagicMock(), MagicMock()] + mock_labels = ['H2O', 'CO2'] + ax1_mock.get_legend_handles_labels.return_value = (mock_handles, mock_labels) plot_chem_atmosphere('output_dir', 'vulcan', plot_format='png', plot_offchem=True) # Verify offline chem was read mock_read_result.assert_called_once() - # Verify plot was created - mock_savefig.assert_called_once() + # Verify plot was saved + fig_mock.savefig.assert_called_once() @pytest.mark.unit -@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') def test_plot_chem_atmosphere_temperature_overlay( - mock_glob, mock_read, mock_subplots, mock_savefig + mock_glob, mock_read, mock_subplots ): """ Test that temperature profiles are added to both panels. @@ -294,7 +304,11 @@ def test_plot_chem_atmosphere_temperature_overlay( ax2_temp_mock = MagicMock() ax1_mock.twiny.return_value = ax1_temp_mock ax2_mock.twiny.return_value = ax2_temp_mock - ax1_mock.get_legend_handles_labels.return_value = ([], []) + + # Mock get_legend_handles_labels - only H2O has nonzero VMR + mock_handles = [MagicMock()] + mock_labels = ['H2O'] + ax1_mock.get_legend_handles_labels.return_value = (mock_handles, mock_labels) plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) @@ -302,18 +316,17 @@ def test_plot_chem_atmosphere_temperature_overlay( ax1_temp_mock.plot.assert_called_once() ax2_temp_mock.plot.assert_called_once() - # Verify temperature axes were labeled + # Verify temperature axis was labeled on panel 1 ax1_temp_mock.set_xlabel.assert_called_once() - ax2_temp_mock.set_xlabel.assert_called_once() + # Panel 2 temperature axis has no label (xticks hidden) @pytest.mark.unit -@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') def test_plot_chem_atmosphere_time_annotation( - mock_glob, mock_read, mock_subplots, mock_savefig + mock_glob, mock_read, mock_subplots ): """ Test that simulation time is annotated on the plot. @@ -338,24 +351,27 @@ def test_plot_chem_atmosphere_time_annotation( mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) ax1_mock.twiny.return_value = MagicMock() ax2_mock.twiny.return_value = MagicMock() - ax1_mock.get_legend_handles_labels.return_value = ([], []) + + # Mock get_legend_handles_labels - only H2O has nonzero VMR + mock_handles = [MagicMock()] + mock_labels = ['H2O'] + ax1_mock.get_legend_handles_labels.return_value = (mock_handles, mock_labels) plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) - # Verify time annotation was added to figure - fig_mock.text.assert_called_once() - call_args = fig_mock.text.call_args[0] + # Verify time annotation was added to ax1 (not fig) + ax1_mock.text.assert_called_once() + call_args = ax1_mock.text.call_args[0] # Should contain the year value (5000) assert '5.00e+03' in call_args[2] or '5000' in str(call_args) @pytest.mark.unit -@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') def test_plot_chem_atmosphere_clouds_and_aerosols( - mock_glob, mock_read, mock_subplots, mock_savefig + mock_glob, mock_read, mock_subplots ): """ Test plotting with both clouds and aerosols present. @@ -382,7 +398,11 @@ def test_plot_chem_atmosphere_clouds_and_aerosols( mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) ax1_mock.twiny.return_value = MagicMock() ax2_mock.twiny.return_value = MagicMock() - ax1_mock.get_legend_handles_labels.return_value = ([], []) + + # Mock get_legend_handles_labels - H2O and CO2 have nonzero VMR + mock_handles = [MagicMock(), MagicMock()] + mock_labels = ['H2O', 'CO2'] + ax1_mock.get_legend_handles_labels.return_value = (mock_handles, mock_labels) plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) @@ -395,12 +415,11 @@ def test_plot_chem_atmosphere_clouds_and_aerosols( @pytest.mark.unit -@patch('proteus.plot.cpl_chem_atmosphere.plt.savefig') @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') def test_plot_chem_atmosphere_no_aerosols_key( - mock_glob, mock_read, mock_subplots, mock_savefig + mock_glob, mock_read, mock_subplots ): """ Test plotting when aerosols key is missing (backward compatibility). @@ -425,11 +444,15 @@ def test_plot_chem_atmosphere_no_aerosols_key( mock_subplots.return_value = (fig_mock, (ax1_mock, ax2_mock)) ax1_mock.twiny.return_value = MagicMock() ax2_mock.twiny.return_value = MagicMock() - ax1_mock.get_legend_handles_labels.return_value = ([], []) + + # Mock get_legend_handles_labels - only H2O has nonzero VMR + mock_handles = [MagicMock()] + mock_labels = ['H2O'] + ax1_mock.get_legend_handles_labels.return_value = (mock_handles, mock_labels) # Should not crash plot_chem_atmosphere('output_dir', None, plot_format='png', plot_offchem=False) # Verify gas panel was still plotted assert ax1_mock.plot.called - mock_savefig.assert_called_once() + fig_mock.savefig.assert_called_once() From 2fa6c8c1f381aecf2f1244b8b3bbf895b5c2c028 Mon Sep 17 00:00:00 2001 From: Harrison Nicholls Date: Sun, 29 Mar 2026 19:09:21 +0100 Subject: [PATCH 17/17] code style --- tests/plot/test_cpl_chem_atmosphere.py | 28 +++++++------------------- 1 file changed, 7 insertions(+), 21 deletions(-) diff --git a/tests/plot/test_cpl_chem_atmosphere.py b/tests/plot/test_cpl_chem_atmosphere.py index bb1261a37..6c8bcb772 100644 --- a/tests/plot/test_cpl_chem_atmosphere.py +++ b/tests/plot/test_cpl_chem_atmosphere.py @@ -40,9 +40,7 @@ def test_plot_chem_atmosphere_no_files(mock_glob): @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') @patch('proteus.plot.cpl_chem_atmosphere.os.path.join') -def test_plot_chem_atmosphere_basic( - mock_join, mock_glob, mock_read, mock_subplots -): +def test_plot_chem_atmosphere_basic(mock_join, mock_glob, mock_read, mock_subplots): """ Test basic plot generation with gas species only. @@ -111,9 +109,7 @@ def test_plot_chem_atmosphere_basic( @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') @patch('proteus.plot.cpl_chem_atmosphere.os.path.join') -def test_plot_chem_atmosphere_with_clouds( - mock_join, mock_glob, mock_read, mock_subplots -): +def test_plot_chem_atmosphere_with_clouds(mock_join, mock_glob, mock_read, mock_subplots): """ Test plot generation including cloud profiles. @@ -158,9 +154,7 @@ def test_plot_chem_atmosphere_with_clouds( @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') @patch('proteus.plot.cpl_chem_atmosphere.os.path.join') -def test_plot_chem_atmosphere_with_aerosols( - mock_join, mock_glob, mock_read, mock_subplots -): +def test_plot_chem_atmosphere_with_aerosols(mock_join, mock_glob, mock_read, mock_subplots): """ Test plot generation including aerosol profiles with species names. @@ -274,9 +268,7 @@ def test_plot_chem_atmosphere_with_offchem( @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') -def test_plot_chem_atmosphere_temperature_overlay( - mock_glob, mock_read, mock_subplots -): +def test_plot_chem_atmosphere_temperature_overlay(mock_glob, mock_read, mock_subplots): """ Test that temperature profiles are added to both panels. @@ -325,9 +317,7 @@ def test_plot_chem_atmosphere_temperature_overlay( @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') -def test_plot_chem_atmosphere_time_annotation( - mock_glob, mock_read, mock_subplots -): +def test_plot_chem_atmosphere_time_annotation(mock_glob, mock_read, mock_subplots): """ Test that simulation time is annotated on the plot. @@ -370,9 +360,7 @@ def test_plot_chem_atmosphere_time_annotation( @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') -def test_plot_chem_atmosphere_clouds_and_aerosols( - mock_glob, mock_read, mock_subplots -): +def test_plot_chem_atmosphere_clouds_and_aerosols(mock_glob, mock_read, mock_subplots): """ Test plotting with both clouds and aerosols present. @@ -418,9 +406,7 @@ def test_plot_chem_atmosphere_clouds_and_aerosols( @patch('proteus.plot.cpl_chem_atmosphere.plt.subplots') @patch('proteus.plot.cpl_chem_atmosphere.read_ncdf_profile') @patch('proteus.plot.cpl_chem_atmosphere.glob.glob') -def test_plot_chem_atmosphere_no_aerosols_key( - mock_glob, mock_read, mock_subplots -): +def test_plot_chem_atmosphere_no_aerosols_key(mock_glob, mock_read, mock_subplots): """ Test plotting when aerosols key is missing (backward compatibility).