|
| 1 | +""" |
| 2 | +Make sure to run the machine-specific commands below before |
| 3 | +running this script: |
| 4 | +
|
| 5 | +Compy: |
| 6 | + srun --pty --nodes=1 --time=01:00:00 /bin/bash |
| 7 | + source /share/apps/E3SM/conda_envs/load_latest_e3sm_unified_compy.sh |
| 8 | +
|
| 9 | +LCRC: |
| 10 | + srun --pty --nodes=1 --time=01:00:00 /bin/bash |
| 11 | + source /lcrc/soft/climate/e3sm-unified/load_latest_e3sm_unified_chrysalis.sh |
| 12 | + Or: source /lcrc/soft/climate/e3sm-unified/load_latest_e3sm_unified_anvil.sh |
| 13 | +
|
| 14 | +NERSC perlmutter cpu: |
| 15 | + salloc --nodes 1 --qos interactive --time 01:00:00 --constraint cpu --account=e3sm |
| 16 | + source /global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh |
| 17 | +""" |
| 18 | +# flake8: noqa E501 |
| 19 | + |
| 20 | +import os |
| 21 | +import timeit |
| 22 | +from typing import Tuple, TypedDict |
| 23 | + |
| 24 | +from mache import MachineInfo |
| 25 | + |
| 26 | +from e3sm_diags.parameter.annual_cycle_zonal_mean_parameter import ACzonalmeanParameter |
| 27 | +from e3sm_diags.parameter.area_mean_time_series_parameter import ( |
| 28 | + AreaMeanTimeSeriesParameter, |
| 29 | +) |
| 30 | +from e3sm_diags.parameter.arm_diags_parameter import ARMDiagsParameter |
| 31 | +from e3sm_diags.parameter.core_parameter import CoreParameter |
| 32 | +from e3sm_diags.parameter.diurnal_cycle_parameter import DiurnalCycleParameter |
| 33 | +from e3sm_diags.parameter.enso_diags_parameter import EnsoDiagsParameter |
| 34 | +from e3sm_diags.parameter.mp_partition_parameter import MPpartitionParameter |
| 35 | +from e3sm_diags.parameter.qbo_parameter import QboParameter |
| 36 | +from e3sm_diags.parameter.streamflow_parameter import StreamflowParameter |
| 37 | +from e3sm_diags.parameter.tc_analysis_parameter import TCAnalysisParameter |
| 38 | +from e3sm_diags.parameter.zonal_mean_2d_stratosphere_parameter import ( |
| 39 | + ZonalMean2dStratosphereParameter, |
| 40 | +) |
| 41 | +from e3sm_diags.run import runner |
| 42 | + |
| 43 | + |
| 44 | +class MachinePaths(TypedDict): |
| 45 | + html_path: str |
| 46 | + obs_climo: str |
| 47 | + test_climo: str |
| 48 | + obs_ts: str |
| 49 | + test_ts: str |
| 50 | + dc_obs_climo: str |
| 51 | + dc_test_climo: str |
| 52 | + arm_obs: str |
| 53 | + arm_test: str |
| 54 | + tc_obs: str |
| 55 | + tc_test: str |
| 56 | + |
| 57 | + |
| 58 | +def run_all_sets(): |
| 59 | + machine_paths: MachinePaths = _get_machine_paths() |
| 60 | + |
| 61 | + param = CoreParameter() |
| 62 | + |
| 63 | + param.reference_data_path = machine_paths["obs_climo"] |
| 64 | + param.test_data_path = machine_paths["test_climo"] |
| 65 | + param.test_name = "20210528.v2rc3e.piControl.ne30pg2_EC30to60E2r2.chrysalis" |
| 66 | + param.seasons = [ |
| 67 | + "ANN", |
| 68 | + "JJA", |
| 69 | + ] # Default setting: seasons = ["ANN", "DJF", "MAM", "JJA", "SON"] |
| 70 | + param.results_dir = "/global/cfs/cdirs/e3sm/www/cdat-migration-fy24/25-01-14-branch-907-v2-data-streamflow" |
| 71 | + param.multiprocessing = True |
| 72 | + param.num_workers = 24 |
| 73 | + |
| 74 | + # Set specific parameters for new sets |
| 75 | + enso_param = EnsoDiagsParameter() |
| 76 | + enso_param.reference_data_path = machine_paths["obs_ts"] |
| 77 | + enso_param.test_data_path = machine_paths["test_ts"] |
| 78 | + enso_param.test_name = "e3sm_v2" |
| 79 | + enso_param.test_start_yr = "0051" |
| 80 | + enso_param.test_end_yr = "0060" |
| 81 | + # Enso obs data range from year 1979 to 2016 |
| 82 | + enso_param.ref_start_yr = "2001" |
| 83 | + enso_param.ref_end_yr = "2010" |
| 84 | + |
| 85 | + qbo_param = QboParameter() |
| 86 | + qbo_param.reference_data_path = machine_paths["obs_ts"] |
| 87 | + qbo_param.test_data_path = machine_paths["test_ts"] |
| 88 | + qbo_param.test_name = "e3sm_v2" |
| 89 | + qbo_param.start_yr = "0051" |
| 90 | + qbo_param.end_yr = "0060" |
| 91 | + # Qbo obs data range from year 1979 to 2019 |
| 92 | + # Number of years of test and ref should match |
| 93 | + qbo_param.ref_start_yr = "2001" |
| 94 | + qbo_param.ref_end_yr = "2010" |
| 95 | + |
| 96 | + ts_param = AreaMeanTimeSeriesParameter() |
| 97 | + ts_param.reference_data_path = machine_paths["obs_ts"] |
| 98 | + ts_param.test_data_path = machine_paths["test_ts"] |
| 99 | + ts_param.test_name = "e3sm_v2" |
| 100 | + ts_param.start_yr = "0051" |
| 101 | + ts_param.end_yr = "0060" |
| 102 | + |
| 103 | + dc_param = DiurnalCycleParameter() |
| 104 | + dc_param.reference_data_path = machine_paths["dc_obs_climo"] |
| 105 | + dc_param.test_data_path = machine_paths["dc_test_climo"] |
| 106 | + dc_param.short_test_name = "e3sm_v2" |
| 107 | + # Plotting diurnal cycle amplitude on different scales. Default is True |
| 108 | + dc_param.normalize_test_amp = False |
| 109 | + |
| 110 | + streamflow_param = StreamflowParameter() |
| 111 | + streamflow_param.reference_data_path = machine_paths["obs_ts"] |
| 112 | + streamflow_param.test_data_path = machine_paths["test_ts"] |
| 113 | + streamflow_param.short_test_name = "e3sm_v2" |
| 114 | + streamflow_param.test_start_yr = "0051" |
| 115 | + streamflow_param.test_end_yr = "0060" |
| 116 | + # Streamflow gauge station data range from year 1986 to 1995 |
| 117 | + streamflow_param.ref_start_yr = "1986" |
| 118 | + streamflow_param.ref_end_yr = "1995" |
| 119 | + |
| 120 | + arm_param = ARMDiagsParameter() |
| 121 | + arm_param.reference_data_path = machine_paths["arm_obs"] |
| 122 | + arm_param.ref_name = "armdiags" |
| 123 | + arm_param.test_data_path = machine_paths["arm_test"] |
| 124 | + arm_param.test_name = "e3sm_v2" |
| 125 | + # arm_param.test_start_yr = "1996" |
| 126 | + # arm_param.test_end_yr = "2010" |
| 127 | + arm_param.test_start_yr = "1985" |
| 128 | + arm_param.test_end_yr = "2014" |
| 129 | + # For model vs obs, the ref start and end year can be any four digit strings. |
| 130 | + # For now, will use all available years form obs |
| 131 | + arm_param.ref_start_yr = "0001" |
| 132 | + arm_param.ref_end_yr = "0001" |
| 133 | + |
| 134 | + tc_param = TCAnalysisParameter() |
| 135 | + tc_param.reference_data_path = machine_paths["tc_obs"] |
| 136 | + tc_param.test_data_path = machine_paths["tc_test"] |
| 137 | + tc_param.short_test_name = "e3sm_v2" |
| 138 | + tc_param.test_start_yr = "0051" |
| 139 | + tc_param.test_end_yr = "0060" |
| 140 | + # For model vs obs, the ref start and end year can be any four digit strings. |
| 141 | + # For now, use all available years form obs by default. |
| 142 | + tc_param.ref_start_yr = "1979" |
| 143 | + tc_param.ref_end_yr = "2018" |
| 144 | + |
| 145 | + ac_param = ACzonalmeanParameter() |
| 146 | + |
| 147 | + zm_param = ZonalMean2dStratosphereParameter() |
| 148 | + |
| 149 | + mp_param = MPpartitionParameter() |
| 150 | + # mp_param.reference_data_path = machine_paths["obs_ts"] |
| 151 | + mp_param.test_data_path = machine_paths["test_ts"] |
| 152 | + mp_param.short_test_name = "e3sm_v2" |
| 153 | + mp_param.test_start_yr = "0051" |
| 154 | + mp_param.test_end_yr = "0060" |
| 155 | + |
| 156 | + param.save_netcdf = True |
| 157 | + runner.sets_to_run = [ |
| 158 | + "lat_lon", |
| 159 | + "zonal_mean_xy", |
| 160 | + "zonal_mean_2d", |
| 161 | + "zonal_mean_2d_stratosphere", |
| 162 | + "polar", |
| 163 | + "cosp_histogram", |
| 164 | + "meridional_mean_2d", |
| 165 | + "annual_cycle_zonal_mean", |
| 166 | + "enso_diags", |
| 167 | + "qbo", |
| 168 | + "area_mean_time_series", |
| 169 | + "diurnal_cycle", |
| 170 | + "streamflow", |
| 171 | + "arm_diags", |
| 172 | + "tc_analysis", |
| 173 | + "aerosol_aeronet", |
| 174 | + "aerosol_budget", |
| 175 | + "mp_partition", |
| 176 | + ] |
| 177 | + |
| 178 | + runner.run_diags( |
| 179 | + [ |
| 180 | + param, |
| 181 | + zm_param, |
| 182 | + ac_param, |
| 183 | + enso_param, |
| 184 | + qbo_param, |
| 185 | + ts_param, |
| 186 | + dc_param, |
| 187 | + streamflow_param, |
| 188 | + arm_param, |
| 189 | + tc_param, |
| 190 | + mp_param, |
| 191 | + ] |
| 192 | + ) |
| 193 | + |
| 194 | + return param.results_dir |
| 195 | + |
| 196 | + |
| 197 | +def _get_machine_paths() -> MachinePaths: |
| 198 | + """Returns the paths on the machine that are required to run e3sm_diags. |
| 199 | +
|
| 200 | + Returns |
| 201 | + ------- |
| 202 | + MachinePaths |
| 203 | + A dictionary of paths on the machine, with the key being the path type |
| 204 | + and the value being the absolute path string. |
| 205 | + """ |
| 206 | + # Get the current machine's configuration info. |
| 207 | + machine_info = MachineInfo() |
| 208 | + machine = machine_info.machine |
| 209 | + |
| 210 | + if machine not in [ |
| 211 | + "anvil", |
| 212 | + "chrysalis", |
| 213 | + "compy", |
| 214 | + "pm-cpu", |
| 215 | + "cori-haswell", |
| 216 | + "cori-knl", |
| 217 | + ]: |
| 218 | + raise ValueError(f"e3sm_diags is not supported on this machine ({machine}).") |
| 219 | + |
| 220 | + # Path to the HTML outputs for the current user. |
| 221 | + web_portal_base_path = machine_info.config.get("web_portal", "base_path") |
| 222 | + html_path = f"{web_portal_base_path}/{machine_info.username}/" |
| 223 | + |
| 224 | + # Path to the reference data directory. |
| 225 | + diags_base_path = machine_info.diagnostics_base |
| 226 | + ref_data_dir = f"{diags_base_path}/observations/Atm" |
| 227 | + |
| 228 | + # Paths to the test data directories. |
| 229 | + test_data_dir, test_data_dir2 = _get_test_data_dirs(machine) |
| 230 | + |
| 231 | + # Construct the paths required by e3sm_diags using the base paths above. |
| 232 | + machine_paths: MachinePaths = { |
| 233 | + "html_path": html_path, |
| 234 | + "obs_climo": f"{ref_data_dir}/climatology", |
| 235 | + "test_climo": f"{test_data_dir}/climatology/rgr/", |
| 236 | + "obs_ts": f"{ref_data_dir}/time-series/", |
| 237 | + "test_ts": f"{test_data_dir}/time-series/rgr/", |
| 238 | + "dc_obs_climo": f"{ref_data_dir}/climatology", |
| 239 | + "dc_test_climo": f"{test_data_dir}/diurnal_climatology/rgr", |
| 240 | + "arm_obs": f"{ref_data_dir}/arm-diags-data/", |
| 241 | + "arm_test": f"{test_data_dir2}/arm-diags-data/", |
| 242 | + "tc_obs": f"{ref_data_dir}/tc-analysis/", |
| 243 | + "tc_test": f"{test_data_dir}/tc-analysis/", |
| 244 | + } |
| 245 | + |
| 246 | + return machine_paths |
| 247 | + |
| 248 | + |
| 249 | +def _get_test_data_dirs(machine: str) -> Tuple[str, str]: |
| 250 | + """Get the directories for test data based on the machine. |
| 251 | +
|
| 252 | + The second path is for using the high frequency grid box output at ARM sites |
| 253 | + from another simulation when the output is available. |
| 254 | +
|
| 255 | + Parameters |
| 256 | + ---------- |
| 257 | + machine : str |
| 258 | + The name of the machine. |
| 259 | +
|
| 260 | + Returns |
| 261 | + ------- |
| 262 | + Tuple[str, str] |
| 263 | + A tuple of two strings, each representing a test data directory path. |
| 264 | + """ |
| 265 | + test_data_dirs = None |
| 266 | + |
| 267 | + # TODO: Update this function to use `mache` after the directories are updated. |
| 268 | + if machine in ["chrysalis", "anvil"]: |
| 269 | + base = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/postprocessed_e3sm_v2_data_for_e3sm_diags" |
| 270 | + elif machine in ["compy"]: |
| 271 | + base = "/compyfs/e3sm_diags_data/postprocessed_e3sm_v2_data_for_e3sm_diags" |
| 272 | + elif machine in ["cori-haswell", "cori-knl", "pm-cpu"]: |
| 273 | + base = "/global/cfs/cdirs/e3sm/e3sm_diags/postprocessed_e3sm_v2_data_for_e3sm_diags" |
| 274 | + |
| 275 | + test_data_dirs = ( |
| 276 | + f"{base}/20210528.v2rc3e.piControl.ne30pg2_EC30to60E2r2.chrysalis", |
| 277 | + # f"{base}/20210719.PhaseII.F20TR-P3.NGD.ne30pg2.compy", |
| 278 | + f"{base}/20221103.v2.LR.amip.NGD_v3atm.chrysalis", |
| 279 | + ) |
| 280 | + |
| 281 | + return test_data_dirs # type: ignore |
| 282 | + |
| 283 | + |
| 284 | +if __name__ == "__main__": |
| 285 | + start_time = timeit.default_timer() |
| 286 | + run_all_sets() |
| 287 | + end_time = timeit.default_timer() |
| 288 | + elapsed_time = end_time - start_time |
| 289 | + print(f"Elapsed time: {elapsed_time} seconds") |
0 commit comments