diff --git a/io_output.py b/io_output.py index 22f7cd4..ffd4013 100644 --- a/io_output.py +++ b/io_output.py @@ -1,6 +1,7 @@ #!/usr/bin/env python from scipy.io import netcdf import numpy as np +import json from parcel_common import _Chem_g_id, _Chem_a_id @@ -90,6 +91,20 @@ def _output_init(micro, opts, spectra): fout.createVariable("ice_mix_ratio", 'd', ('t',)) fout.variables["ice_mix_ratio"].unit = "kg/kg" + # if micro.opts_init.exact_sstp_cond: + fout.createVariable("sstp_cond_mean", 'd', ('t',)) + fout.variables["sstp_cond_mean"].unit = "1" + + fout.createVariable("act_m0", 'd', ('t',)) + fout.variables["act_m0"].unit = "1/kg" + + fout.createVariable("sd_conc", 'd', ('t',)) + fout.variables["sd_conc"].unit = "1" + + # wall-clock time spent inside step_cond (per parcel output timestep) + fout.createVariable("step_cond_walltime_ms", 'd', ('t',)) + fout.variables["step_cond_walltime_ms"].unit = "ms" + return fout @@ -157,8 +172,46 @@ def _output_save(fout, state, rec): def _save_attrs(fout, dictnr): + """Save metadata as NetCDF global attributes. + + `scipy.io.netcdf` is picky about attribute types; it expects scalar strings/numbers + or 1-D arrays. The `opts` dict often contains Python objects (dicts, lists, None, + callables) that would crash on file close/flush. + """ + + def _coerce_attr_value(v): + # NetCDF has no null type; store None explicitly + if v is None: + return "None" + + # Basic scalar types are fine + if isinstance(v, (str, int, float, bool, np.number)): + return v + + # Numpy arrays: keep small numeric arrays, stringify object arrays + if isinstance(v, np.ndarray): + if v.dtype == object: + return repr(v.tolist()) + return v + + # Containers / complex objects: stringify deterministically + if isinstance(v, (dict, list, tuple, set)): + try: + return json.dumps(v, sort_keys=True, default=str) + except Exception: + return repr(v) + + # Callables, modules, etc. + return repr(v) + for var, val in dictnr.items(): - setattr(fout, var, val) + coerced = _coerce_attr_value(val) + # Keep the previous debug prints minimal and safe + try: + setattr(fout, var, coerced) + except Exception: + # Last resort: force string + setattr(fout, var, repr(coerced)) def _output(fout, opts, micro, state, rec, spectra): diff --git a/micro_lgrngn.py b/micro_lgrngn.py index fd5fdf9..9dfafce 100644 --- a/micro_lgrngn.py +++ b/micro_lgrngn.py @@ -1,5 +1,6 @@ #!/usr/bin/env python import numpy as np +import time from libcloudphxx import lgrngn from parcel_common import lognormal, sum_of_lognormals, _Chem_g_id, _Chem_a_id, _stats @@ -9,8 +10,24 @@ def _micro_init(aerosol, opts, state): # lagrangian scheme options opts_init = lgrngn.opts_init_t() - for opt in ["dt", "sd_conc", "chem_rho", "sstp_cond","ice_switch","time_dep_ice_nucl"]: - setattr(opts_init, opt, opts[opt]) + for opt in [ + "dt", + "sd_conc", + "chem_rho", + "sstp_cond", + "ice_switch", + "time_dep_ice_nucl", + "adaptive_sstp_cond", + "sstp_cond_adapt_drw2_eps", + "sstp_cond_adapt_drw2_max", + "sstp_cond_act", + "sstp_cond_mix", + "exact_sstp_cond", + "aerosol_independent_of_rhod" + ]: + if opt in opts and opts[opt] is not None: + setattr(opts_init, opt, opts[opt]) + opts_init.n_sd_max = opts_init.sd_conc if opts["rng_seed"] is not None: opts_init.rng_seed = int(opts["rng_seed"]) @@ -18,14 +35,50 @@ def _micro_init(aerosol, opts, state): opts_init.th_dry = True opts_init.const_p = False - # read in the initial aerosol size distribution - dry_distros = {} - for name, dct in aerosol.items(): # loop over kappas - lognormals = [] - for i in range(len(dct["mean_r"])): - lognormals.append(lognormal(dct["mean_r"][i], dct["gstdev"][i], dct["n_tot"][i])) - dry_distros[(dct["kappa"], opts["rd_insol"])] = sum_of_lognormals(lognormals) - opts_init.dry_distros = dry_distros + # --- aerosol initialization --- + # dry_distros from lognormal spec (opts['aerosol']) + if aerosol is not None and isinstance(aerosol, dict) and len(aerosol) > 0: + dry_distros = {} + for name, dct in aerosol.items(): + lognormals = [] + for i in range(len(dct["mean_r"])): + lognormals.append(lognormal(dct["mean_r"][i], dct["gstdev"][i], dct["n_tot"][i])) + dry_distros[(float(dct["kappa"]), float(opts["rd_insol"]))] = sum_of_lognormals(lognormals) + opts_init.dry_distros = dry_distros + + # dry_sizes from discrete bins (opts['dry_sizes']) + ds = opts.get("dry_sizes") + if ds is not None: + print(opts.get("dry_sizes")) + if not isinstance(ds, dict) or len(ds) == 0: + raise ValueError("dry_sizes must be a non-empty dict when provided") + + dry_sizes = {} + for name, dct in ds.items(): + print(name, dct) + if "kappa" not in dct or "bins" not in dct: + raise ValueError("Each dry_sizes mode must define 'kappa' and 'bins'") + kappa = float(dct["kappa"]) + bins = dct["bins"] + if not isinstance(bins, dict) or len(bins) == 0: + raise ValueError("dry_sizes 'bins' must be a non-empty dict of radius->[conc, n_sd]") + + print(bins) + bins_parsed = {} + for rd_key, val in bins.items(): + print(rd_key, val) + rd = float(rd_key) + if not (isinstance(val, (list, tuple)) and len(val) == 2): + raise ValueError("dry_sizes bins values must be [STP_concentration_1_per_m3, number_of_SDs]") + conc = float(val[0]) + n_sd = int(val[1]) + bins_parsed[rd] = [conc, n_sd] + + print(bins_parsed) + dry_sizes[(kappa, float(opts["rd_insol"]))] = bins_parsed + print(dry_sizes) + + opts_init.dry_sizes = dry_sizes # better resolution for the SD tail if opts["large_tail"]: @@ -43,7 +96,22 @@ def _micro_init(aerosol, opts, state): opts_init.sstp_chem = opts["sstp_chem"] # initialisation - micro = lgrngn.factory(lgrngn.backend_t.serial, opts_init) + backend_str = opts.get("backend", "serial") + if backend_str is None: + backend_str = "serial" + backend_str = str(backend_str).lower() + + backend_map = { + "serial": lgrngn.backend_t.serial, + "openmp": lgrngn.backend_t.OpenMP, + "omp": lgrngn.backend_t.OpenMP, + "cuda": lgrngn.backend_t.CUDA, + "gpu": lgrngn.backend_t.CUDA, + } + if backend_str not in backend_map: + raise ValueError(f"Unknown lgrngn backend: {backend_str!r} (expected one of: {', '.join(sorted(backend_map))})") + + micro = lgrngn.factory(backend_map[backend_str], opts_init) ambient_chem = {} if micro.opts_init.chem_switch: ambient_chem = dict((v, state[k]) for k,v in _Chem_g_id.items()) @@ -56,6 +124,7 @@ def _micro_step(micro, state, info, opts): '''Microphysics step for lagrangian scheme''' libopts = lgrngn.opts_t() libopts.cond = True + libopts.depo = True libopts.coal = False libopts.adve = False libopts.sedi = False @@ -74,7 +143,13 @@ def _micro_step(micro, state, info, opts): ambient_chem = dict((v, state[k]) for k,v in _Chem_g_id.items()) # call libcloudphxx microphysics - micro.step_sync(libopts, state["th_d"], state["r_v"], state["rhod"], ambient_chem=ambient_chem) + # micro.step_sync(libopts, state["th_d"], state["r_v"], state["rhod"], ambient_chem=ambient_chem) + micro.sync_in(state["th_d"], state["r_v"], state["rhod"], ambient_chem=ambient_chem) + + t0 = time.perf_counter() + micro.step_cond(libopts, state["th_d"], state["r_v"], ambient_chem=ambient_chem) + state["step_cond_walltime_ms"] = (time.perf_counter() - t0) * 1e3 + micro.step_async(libopts) # update state after microphysics (needed for below update for chemistry) @@ -90,4 +165,24 @@ def _micro_step(micro, state, info, opts): if micro.opts_init.ice_switch: micro.diag_ice() micro.diag_ice_mix_ratio() - state["ice_mix_ratio"] = np.frombuffer(micro.outbuf())[0] \ No newline at end of file + state["ice_mix_ratio"] = np.frombuffer(micro.outbuf())[0] + # if micro.opts_init.exact_sstp_cond: + try: # depending on options, sstp_cond_avg may not be available + micro.diag_all() + mom1 = micro.diag_sstp_cond_mom(1) + mom1 = np.frombuffer(micro.outbuf())[0] + mom0 = micro.diag_sstp_cond_mom(0) + mom0 = np.frombuffer(micro.outbuf())[0] + state["sstp_cond_mean"] = mom1/mom0 + print("sstp_cond_mean: ", state["sstp_cond_mean"]) + except Exception: + state["sstp_cond_mean"] = np.full_like(state["th_d"], np.nan) + + micro.diag_rw_ge_rc() + mom0 = micro.diag_wet_mom(0) + mom0 = np.frombuffer(micro.outbuf())[0] + state["act_m0"] = mom0 + + micro.diag_all() + micro.diag_sd_conc() + state["sd_conc"] = np.frombuffer(micro.outbuf())[0] \ No newline at end of file diff --git a/parcel.py b/parcel.py index d4e26d4..34c9bdc 100755 --- a/parcel.py +++ b/parcel.py @@ -18,7 +18,7 @@ parcel_version = subprocess.check_output(["git", "rev-parse", "HEAD"]).rstrip() # import refactored modules -from parcel_common import _Chem_g_id, _Chem_a_id, lognormal, sum_of_lognormals, _stats, _p_hydro_const_rho, _p_hydro_const_th_rv, _arguments_checking, _init_sanity_check +from parcel_common import _Chem_g_id, _Chem_a_id, lognormal, sum_of_lognormals, _stats, _p_hydro_const_rho, _p_hydro_const_th_rv, _arguments_checking, _init_sanity_check, _w_eval from micro_lgrngn import _micro_init as _micro_init_lgrngn, _micro_step as _micro_step_lgrngn from micro_blk_1m import _opts_init_blk_1m, _micro_step_blk_1m from micro_blk_1m_ice import _opts_init_blk_1m_ice, _micro_step_blk_1m_ice @@ -35,6 +35,7 @@ def parcel(dt = .1, z_max = 200., w = 1., T_0 = 300., p_0 = 101300., time_dep_ice_nucl = False, sd_conc = 64, aerosol = '{"ammonium_sulfate": {"kappa": 0.61, "mean_r": [0.02e-6], "gstdev": [1.4], "n_tot": [60.0e6]}}', + dry_sizes = None, out_bin = '{"radii": {"rght": 0.01, "moms": [0], "drwt": "wet", "nbin": 1, "lnli": "log", "left": 1e-15}}', SO2_g = 0., O3_g = 0., H2O2_g = 0., CO2_g = 0., HNO3_g = 0., NH3_g = 0., chem_dsl = False, chem_dsc = False, chem_rct = False, @@ -44,18 +45,33 @@ def parcel(dt = .1, z_max = 200., w = 1., T_0 = 300., p_0 = 101300., wait = 0, large_tail = False, rng_seed = None, - rd_insol = 0. + rd_insol = 0., + t = None, + adaptive_sstp_cond = None, + sstp_cond_adapt_drw2_eps = None, + sstp_cond_adapt_drw2_max = None, + sstp_cond_act = None, + sstp_cond_mix = None, + exact_sstp_cond = None, + aerosol_independent_of_rhod = None + ,backend = "serial" ): """ Args: dt (Optional[float]): timestep [s] z_max (Optional[float]): maximum vertical displacement [m] - w (Optional[float]): updraft velocity [m/s] + t (Optional[float|None]): simulation duration [s]. + Exactly one of `z_max` or `t` must be specified. + w (Optional[float|callable|str]): updraft velocity [m/s] + - constant: number + - time-dependent: callable w(t) or expression string in `t` (seconds) + e.g. "1 + 0.5*np.sin(2*np.pi*t/60)" T_0 (Optional[float]): initial temperature [K] p_0 (Optional[float]): initial pressure [Pa] r_0 (Optional[float]): initial water vapour mass mixing ratio [kg/kg] RH_0 (Optional[float]): initial relative humidity scheme (Optional[string]): microphysics scheme to use: 'lgrngn', 'blk_1m' + backend (Optional[str]): lgrngn backend to use: 'serial', 'openmp', 'cuda' (only used when scheme='lgrngn') ice_switch (Optional[bool]): enable ice microphysics ice_nucl (Optional[bool]): enable ice nucleation in lagrangian scheme time_dep_ice_nucl (Optional[bool]): enable time-dependent ice nucleation in lagrangian scheme @@ -68,6 +84,7 @@ def parcel(dt = .1, z_max = 200., w = 1., T_0 = 300., p_0 = 101300., (added for testing) sd_conc (Optional[int]): number of moving bins (super-droplets) + aerosol (Optional[json str]): dict of dicts defining aerosol distribution, e.g.: {"ammonium_sulfate": {"kappa": 0.61, "mean_r": [0.02e-6, 0.07e-7], "gstdev": [1.4, 1.2], "n_tot": [120.0e6, 80.0e6]} @@ -78,6 +95,20 @@ def parcel(dt = .1, z_max = 200., w = 1., T_0 = 300., p_0 = 101300., gstdev - lognormal distribution geometric standard deviation (list if multimodal distribution) n_tot - lognormal distribution total concentration under standard conditions (T=20C, p=1013.25 hPa, rv=0) [m^-3] (list if multimodal distribution) + + dry_sizes (Optional[json str|dict|None]): discrete aerosol bins used to set libcloudphxx `opts_init.dry_sizes`. + Can be used together with `aerosol`/dry_distros. + Format example: + { + "ammonium_sulfate": { + "kappa": 0.61, + "bins": { + "1e-6": [30.0, 15], + "15e-6": [10.0, 5] + } + } + } + where bins map dry_radius_m -> [STP_concentration_1_per_m3, number_of_SDs] large_tail (Optional[bool]) : use more SD to better represent the large tail of the initial aerosol distribution @@ -111,9 +142,23 @@ def parcel(dt = .1, z_max = 200., w = 1., T_0 = 300., p_0 = 101300., chem_dsl (Optional[bool]): on/off for dissolving chem species into droplets chem_dsc (Optional[bool]): on/off for dissociation of chem species in droplets chem_rct (Optional[bool]): on/off for oxidation of S_IV to S_VI + chem_rho (Optional[float]): aerosol/droplet material density for chemistry [kg/m3] + aerosol_independent_of_rhod (Optional[bool]): on/off for initial aerosol concentration independent of rhod (assumed at STP otherwise) + + # Coalescence Substepping controls + sstp_chem (Optional[int]): substeps per timestep for chemistry (>=1) -} + # condensation substepping controls + sstp_cond (Optional[int]): substeps per dynamical timestep for condensation/evaporation (>=1) + adaptive_sstp_cond (Optional[bool]): on/off for adaptive substepping for condensation/evaporation + sstp_cond_adapt_drw2_eps (Optional[float]): tolerance parameter for adaptive condensation/evaporation substepping + sstp_cond_adapt_drw2_max (Optional[float]): maximum relative change of rw2 for adaptive condensation/evaporation substepping + sstp_cond_act (Optional[int]): substeps for (de)activating droplets + sstp_cond_mix (Optional[bool]): on/off mixing of thermodynamic variables between superdroplets after each condensation substep + exact_sstp_cond (Optional[bool]): on/off for per-particle condensation substepping (per-cell if off) + # Misc + rd_insol (Optional[float]): insoluble dry radius offset/addition used by selected microphysics (if applicable) [m] """ # packing function arguments into "opts" dictionary @@ -122,11 +167,28 @@ def parcel(dt = .1, z_max = 200., w = 1., T_0 = 300., p_0 = 101300., for k in args: opts[k] = locals()[k] + # Enforce stop condition selection: either z_max or t (but not both). + # Backwards compatible default is z_max (t=None). + # if opts["t"] is not None and opts["z_max"] is not None: + # raise ValueError("Specify only one stop condition: either z_max or t (not both)") + # if opts["t"] is None and opts["z_max"] is None: + # raise ValueError("You must specify exactly one stop condition: z_max or t") + # if opts["t"] is not None and opts["t"] <= 0: + # raise ValueError("t must be > 0") + # if opts["z_max"] is not None and opts["z_max"] <= 0: + # raise ValueError("z_max must be > 0") + # parsing json specification of output spectra spectra = json.loads(opts["out_bin"]) - # parsing json specification of init aerosol spectra - aerosol = json.loads(opts["aerosol"]) + # parsing json specification of init aerosol spectra (if provided) + aerosol = json.loads(opts["aerosol"]) if isinstance(opts.get("aerosol"), str) else opts.get("aerosol") + + # allow passing dry_sizes as dict or json string + if opts.get("dry_sizes") is not None and isinstance(opts.get("dry_sizes"), str): + dry_sizes = json.loads(opts["dry_sizes"]) + else: + dry_sizes = opts.get("dry_sizes") # default water content if ((opts["r_0"] < 0) and (opts["RH_0"] < 0)): @@ -140,7 +202,22 @@ def parcel(dt = .1, z_max = 200., w = 1., T_0 = 300., p_0 = 101300., _arguments_checking(opts, spectra, aerosol, ice_switch) th_0 = T_0 * (common.p_1000 / p_0)**(common.R_d / common.c_pd) - nt = int(z_max / (w * dt)) + + # Stopping condition differs for constant vs. time-dependent w. + # For constant w: keep original behaviour (nt computed from stop condition). + # For variable w: integrate until the stop condition is met. + # w0 = _w_eval(w, 0.0) + # if w0 <= 0 and isinstance(w, (int, float, np.floating)) and opts["z_max"] is not None: + # raise ValueError("For constant w with z_max stop, expected w>0 to reach z_max") + + # if isinstance(w, (int, float, np.floating)): + # if opts["t"] is not None: + # nt = int(np.ceil(float(opts["t"]) / dt)) + # else: + # nt = int(opts["z_max"] / (float(w) * dt)) + # else: + # nt = None + state = { "t" : 0, "z" : 0, "r_v" : np.array([r_0]), "p" : p_0, @@ -198,12 +275,35 @@ def parcel(dt = .1, z_max = 200., w = 1., T_0 = 300., p_0 = 101300., _output_save(fout, state, 0) # simpler output for blk_1m # timestepping - for it in range(1, nt+1): + rec = 0 + it = 0 + max_steps_guard = 5_000_000 # safety for pathological w(t) + while True: + # if nt is not None: + # if it >= nt: + # break + # else: + # variable-w stopping conditions + if opts["t"] is not None: + if state["t"] >= opts["t"]: + break + else: + if state["z"] >= opts["z_max"]: + break + + if it >= max_steps_guard: + raise RuntimeError("Exceeded safety step limit while integrating variable w(t)") + + it += 1 + + # vertical velocity at current (start-of-step) time + w_it = _w_eval(w, state["t"]) + # diagnostics # the reasons to use analytic solution: # - independent of dt # - same as in 2D kinematic model - state["z"] += w * dt + state["z"] += w_it * dt state["t"] = it * dt # pressure @@ -218,7 +318,7 @@ def parcel(dt = .1, z_max = 200., w = 1., T_0 = 300., p_0 = 101300., elif pprof == "pprof_piecewise_const_rhod": # as in Grabowski and Wang 2009 but calculating pressure # for rho piecewise constant per each time step - state["p"] = _p_hydro_const_rho(w*dt, state["p"], state["rhod"][0]) + state["p"] = _p_hydro_const_rho(w_it*dt, state["p"], state["rhod"][0]) else: raise Exception("pprof should be pprof_const_th_rv, pprof_const_rhod, or pprof_piecewise_const_rhod") @@ -251,18 +351,25 @@ def parcel(dt = .1, z_max = 200., w = 1., T_0 = 300., p_0 = 101300., # output if (it % outfreq == 0): - print(str(round(it / (nt * 1.) * 100, 2)) + " %") + # if nt is not None and nt > 0: + # print(str(round(it / (nt * 1.) * 100, 2)) + " %") + if opts["t"] is not None: + print(str(round(state["t"] / (opts["t"] * 1.) * 100, 2)) + " %") + if opts["z_max"] is not None: + print(str(round(state["z"], 1)) + " / " + str(opts["z_max"]) + " m") + rec = it/outfreq if scheme == "lgrngn": _output(fout, opts, micro, state, rec, spectra) elif scheme == "blk_1m": _output_save(fout, state, rec) - _save_attrs(fout, info) + # _save_attrs(fout, info) _save_attrs(fout, opts) + print("post _save_attrs") if wait != 0: - for it in range (nt+1, nt+wait): + for it in range (it+1, it+wait): state["t"] = it * dt if scheme == "lgrngn": _micro_step_lgrngn(micro, state, info, opts) diff --git a/parcel_common.py b/parcel_common.py index 14f13c5..8290afc 100644 --- a/parcel_common.py +++ b/parcel_common.py @@ -24,6 +24,26 @@ "S_VI" : lgrngn.chem_species_t.S_VI } +def _w_eval(w, t): + """Evaluate vertical velocity at time t. + + Supported forms for `w`: + - float/int: constant vertical velocity + - callable: w(t) -> float + - str: Python expression in variable `t` (seconds), e.g. "1 + 0.5*np.sin(2*np.pi*t/60)" + Available names: t, np + + Note: expression strings are `eval`'d with a restricted global namespace. + """ + if callable(w): + return float(w(t)) + if isinstance(w, (int, float, np.floating)): + return float(w) + if isinstance(w, str): + # restricted eval environment + return float(eval(w, {"__builtins__": {}}, {"t": float(t), "np": np})) + raise TypeError("w must be a number, a callable w(t), or an expression string") + class lognormal(object): def __init__(self, mean_r, gstdev, n_tot): @@ -70,41 +90,53 @@ def _arguments_checking(opts, spectra, aerosol, ice_switch): raise Exception("temperature should be larger than 0C if ice_switch=False") elif ((opts["r_0"] >= 0) and (opts["RH_0"] >= 0)): raise Exception("both r_0 and RH_0 specified, please use only one") - if opts["w"] < 0: - raise Exception("vertical velocity should be larger than 0") - - for name, dct in aerosol.items(): - # TODO: check if name is valid netCDF identifier - # (http://www.unidata.ucar.edu/software/thredds/current/netcdf-java/CDM/Identifiers.html) - keys = ["kappa", "mean_r", "n_tot", "gstdev"] - for key in keys: - if key not in dct: - raise Exception(">>" + key + "<< is missing in aerosol[" + name + "]") - for key in dct: - if key not in keys: - raise Exception("invalid key >>" + key + "<< in aerosol[" + name + "]") - if dct["kappa"] <= 0: - raise Exception("kappa hygroscopicity parameter should be larger than 0 for aerosol[" + name + "]") - if type(dct["mean_r"]) != list: - raise Exception(">>mean_r<< key in aerosol["+ name +"] must be a list") - if type(dct["gstdev"]) != list: - raise Exception(">>gstdev<< key in aerosol["+ name +"] must be a list") - if type(dct["n_tot"]) != list: - raise Exception(">>n_tot<< key in aerosol["+ name +"] must be a list") - if not len(dct["mean_r"]) == len(dct["n_tot"]) == len(dct["gstdev"]): - raise Exception("mean_r, n_tot and gstdev lists should have same sizes for aerosol[" + name + "]") - for mean_r in dct["mean_r"]: - if mean_r <= 0: - raise Exception("mean radius should be > 0 for aerosol[" + name + "]") - for n_tot in dct["n_tot"]: - if n_tot <= 0: - raise Exception("concentration should be > 0 for aerosol[" + name + "]") - for gstdev in dct["gstdev"]: - if gstdev <= 0: - raise Exception("standard deviation should be > 0 for aerosol[" + name + "]") - # necessary? - if gstdev == 1.: - raise Exception("standard deviation should be != 1 to avoid monodisperse distribution for aerosol[" + name + "]") + # if opts["w"] < 0: + # raise Exception("vertical velocity should be larger than 0") + if opts["t"] is not None and opts["z_max"] is not None: + raise ValueError("Specify only one stop condition: either z_max or t (not both)") + if opts["t"] is None and opts["z_max"] is None: + raise ValueError("You must specify exactly one stop condition: z_max or t") + if opts["t"] is not None and opts["t"] <= 0: + raise ValueError("t must be > 0") + if opts["z_max"] is not None and opts["z_max"] <= 0: + raise ValueError("z_max must be > 0") + w0 = _w_eval(opts["w"], 0.0) + if w0 <= 0 and isinstance(opts["w"], (int, float, np.floating)) and opts["z_max"] is not None: + raise ValueError("For constant w with z_max stop, expected w>0 to reach z_max") + + if aerosol is not None: + for name, dct in aerosol.items(): + # TODO: check if name is valid netCDF identifier + # (http://www.unidata.ucar.edu/software/thredds/current/netcdf-java/CDM/Identifiers.html) + keys = ["kappa", "mean_r", "n_tot", "gstdev"] + for key in keys: + if key not in dct: + raise Exception(">>" + key + "<< is missing in aerosol[" + name + "]") + for key in dct: + if key not in keys: + raise Exception("invalid key >>" + key + "<< in aerosol[" + name + "]") + if dct["kappa"] <= 0: + raise Exception("kappa hygroscopicity parameter should be larger than 0 for aerosol[" + name + "]") + if type(dct["mean_r"]) != list: + raise Exception(">>mean_r<< key in aerosol["+ name +"] must be a list") + if type(dct["gstdev"]) != list: + raise Exception(">>gstdev<< key in aerosol["+ name +"] must be a list") + if type(dct["n_tot"]) != list: + raise Exception(">>n_tot<< key in aerosol["+ name +"] must be a list") + if not len(dct["mean_r"]) == len(dct["n_tot"]) == len(dct["gstdev"]): + raise Exception("mean_r, n_tot and gstdev lists should have same sizes for aerosol[" + name + "]") + for mean_r in dct["mean_r"]: + if mean_r <= 0: + raise Exception("mean radius should be > 0 for aerosol[" + name + "]") + for n_tot in dct["n_tot"]: + if n_tot <= 0: + raise Exception("concentration should be > 0 for aerosol[" + name + "]") + for gstdev in dct["gstdev"]: + if gstdev <= 0: + raise Exception("standard deviation should be > 0 for aerosol[" + name + "]") + # necessary? + if gstdev == 1.: + raise Exception("standard deviation should be != 1 to avoid monodisperse distribution for aerosol[" + name + "]") for name, dct in spectra.items(): # TODO: check if name is valid netCDF identifier diff --git a/test_adaptive_sstp_cond.py b/test_adaptive_sstp_cond.py new file mode 100644 index 0000000..cd47935 --- /dev/null +++ b/test_adaptive_sstp_cond.py @@ -0,0 +1,210 @@ +""" +Run adaptive condensation substepping and plot results as in the MSc thesis of Piotr Bartman (Sec. 3.4, Fig. 5 therein) +""" + +import sys, os +sys.path.insert(0, "../") +sys.path.insert(0, "./") + +import numpy as np +from parcel import parcel +from scipy.io import netcdf +import matplotlib.pyplot as plt +from matplotlib.collections import LineCollection +import matplotlib.colors as mcolors +from pathlib import Path +from typing import List + +sstp_cond_max = 10 +z_max = 4000.0 + +def run_scheme(w_max, adaptive, outfile, *, sstp_cond=sstp_cond_max): + args = dict( + p_0=100000, + RH_0=0.9, + T_0=260, + aerosol = None, + # aerosol = '{"pristine": {"kappa": 0.61, "mean_r": [0.011e-6, 0.06e-6], "gstdev": [1.2, 1.7], "n_tot": [125.0e6, 65.0e6]}}', # aerosol=None, + sd_conc=100,#pow(2,10),#1024,#256, + # dry_sizes={"Bartman": {"kappa": 0.2, "bins": { + # str(r_dry): [N_STP, 1] + # }}}, + dt=1, + z_max=None, + # t=300, + w=lambda t: w_max * np.pi / 2. * np.sin(np.pi*t*w_max/z_max), # z_half = z_max + # w=lambda t: w_max * np.pi / 2. * np.sin(np.pi*t/(z_max * w_max)), # z_half = z_max + # r_0=0.022, + outfile=outfile, + # outfreq=1, + scheme="lgrngn", + # out_bin='{"radius": {"rght": 1, "moms": [0,1], "drwt": "wet", "nbin": 1, "lnli": "lin", "left": 1e-15}}', + out_bin='{"cloud": {"rght": 1, "moms": [0,1], "drwt": "wet", "nbin": 1, "lnli": "lin", "left": 0.5e-6}}', + sstp_cond=sstp_cond, + adaptive_sstp_cond=adaptive, + # adaptive substepping parameters are injected below (only when adaptive=True) + sstp_cond_adapt_drw2_eps=None, + sstp_cond_adapt_drw2_max=None, + sstp_cond_act=None, + sstp_cond_mix = False, #cant be True for adaptive + exact_sstp_cond = True, # if adaptive else False, + aerosol_independent_of_rhod=True, + backend="OpenMP", + ice_switch = True, + ice_nucl = True, + time_dep_ice_nucl = True, + rd_insol = 0.1e-6 + ) + + # NOTE: we allow passing these in through function attributes set outside. + if hasattr(run_scheme, "aerosol"): + args["aerosol"] = run_scheme.aerosol + # They only make sense for adaptive_sstp_cond=True. + if adaptive: + if hasattr(run_scheme, "sstp_cond_adapt_drw2_eps"): + args["sstp_cond_adapt_drw2_eps"] = float(run_scheme.sstp_cond_adapt_drw2_eps) + if hasattr(run_scheme, "sstp_cond_adapt_drw2_max"): + args["sstp_cond_adapt_drw2_max"] = float(run_scheme.sstp_cond_adapt_drw2_max) + if hasattr(run_scheme, "sstp_cond_act"): + args["sstp_cond_act"] = int(run_scheme.sstp_cond_act) + + args["t"] = 2. * z_max / w_max # twice the time to to reach z=z_max + args["outfreq"] = 1# args["t"] // 100 # save 100 points + # args["t"] = 300 + # args["t"] = 1 + print("t: ", args["t"]) + parcel(**args) + + with netcdf.netcdf_file(outfile, 'r') as f: + # rv = np.array(f.variables['r_v'][:]) + # th_d = np.array(f.variables['th_d'][:]) + z = np.array(f.variables['z'][:]) + RH = np.array(f.variables['RH'][:]) + sstp_cond_mean = np.array(f.variables['sstp_cond_mean'][:]) if 'sstp_cond_mean' in f.variables else None + sstp_cond_mean[0] = sstp_cond_mean[1] if sstp_cond_mean is not None else None # at t=0 sstp_cond_mean=0, because its set only during the firs step (?) + act_mom0 = np.array(f.variables['act_m0'][:]).squeeze() + step_cond_walltime_ms = np.array(f.variables['step_cond_walltime_ms'][:]).squeeze() if 'step_cond_walltime_ms' in f.variables else None + return RH, z, sstp_cond_mean, act_mom0, step_cond_walltime_ms + +# --- batch scenarios --- + +# baseline - basically no adaptation, very relaxed conditions +baseline = dict( + eps=1e6, #1e-1, + max=1e6, #100, + act=1, # 1 means disabled +) + +vary_eps = [1e-1, 1e-2, 1e-3] + +def make_figure(aerosol_name, aerosol, xmax): + run_scheme.aerosol = aerosol + # rows: w_max; cols: eps + # w_max_list = [0.1, 1., 2.5, 5.0] + w_max_list = [5.0] + fig, axes = plt.subplots(len(w_max_list), len(vary_eps), figsize=(15.0, 15.0), sharex=True, sharey=True, squeeze=False) + + generated_nc_files: List[str] = [] + + # shared colormap settings for sstp_cond_dt (= sstp_cond_mean here) + cmap_dt = "gnuplot" + norm_dt = mcolors.Normalize(vmin=1, vmax=sstp_cond_max) + + for i, w_max in enumerate(w_max_list): + # --- reference run (non-adaptive) once per w_max --- + outfile_ref = f"test_adaptive_sstp_cond_{aerosol_name}_w{w_max:g}_ref_adapt0.nc" + RH_ref, z_ref, _, act_mom0_ref, step_cond_ref_ms = run_scheme(w_max, False, outfile_ref) + generated_nc_files.append(outfile_ref) + x_ref = act_mom0_ref / 1e6 + y_ref = z_ref + + ref_step_cond_mean_ms = float(np.nanmean(step_cond_ref_ms)) if step_cond_ref_ms is not None else float("nan") + + for j, eps in enumerate(vary_eps): + ax = axes[i, j] + + # overlay reference + ax.plot(x_ref, y_ref, color="0.6", linewidth=2.0, zorder=1) + + run_scheme.sstp_cond_adapt_drw2_eps = eps + run_scheme.sstp_cond_adapt_drw2_max = baseline["max"] + run_scheme.sstp_cond_act = baseline["act"] + + outfile = f"test_adaptive_sstp_cond_{aerosol_name}_w{w_max:g}_eps{eps:.0e}_adapt1.nc" + RH, z, sstp_cond_mean, act_mom0, step_cond_walltime_ms = run_scheme(w_max, True, outfile) + generated_nc_files.append(outfile) + + x = act_mom0 / 1e6 + y = z + + sstp_cond_avg = float(np.nanmean(sstp_cond_mean)) if sstp_cond_mean is not None else float("nan") + step_cond_mean_ms = float(np.nanmean(step_cond_walltime_ms)) if step_cond_walltime_ms is not None else float("nan") + speedup = (ref_step_cond_mean_ms / step_cond_mean_ms) if (np.isfinite(ref_step_cond_mean_ms) and np.isfinite(step_cond_mean_ms) and step_cond_mean_ms > 0) else float("nan") + + ax.text( + 0.98, + 0.05, + f"sstp_cond_avg={sstp_cond_avg:.2f}\nref={ref_step_cond_mean_ms:.2f} ms\nspeedup={speedup:.2f}x", + transform=ax.transAxes, + ha="right", + va="bottom", + fontsize=9, + bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="none", alpha=0.7), + zorder=5, + ) + + if sstp_cond_mean is None: + ax.plot(x, y, linewidth=2.0, color="C0", zorder=2) + else: + points = np.array([x, y]).T.reshape(-1, 1, 2) + segments = np.concatenate([points[:-1], points[1:]], axis=1) + lc = LineCollection(segments, cmap=cmap_dt, norm=norm_dt) + lc.set_array(sstp_cond_mean[1:]) + lc.set_linewidth(2.0) + lc.set_zorder=2 + ax.add_collection(lc) + + if i == 0: + ax.set_title(f"eps={eps:.0e}") + if j == 0: + ax.set_ylabel(f"w_max={w_max:g}\nHeight [m]") + + for i in range(len(w_max_list)): + for j in range(len(vary_eps)): + ax = axes[i, j] + ax.set_xlim(-5, xmax) + ax.set_ylim(0, z_max + 25) + if i == len(w_max_list) - 1: + ax.set_xlabel("activated droplets [1/mg]") + + fig.suptitle("Adaptive substepping, " + aerosol_name) + fig.tight_layout(rect=(0, 0.10, 1, 0.97)) + + # shared colorbar + sm = plt.cm.ScalarMappable(cmap=cmap_dt, norm=norm_dt) + sm.set_array([]) + cax = fig.add_axes([0.15, 0.04, 0.70, 0.025]) + cbar = fig.colorbar(sm, cax=cax, orientation="horizontal") + cbar.set_label("sstp_cond_mean [1]") + + out_png = "test_adaptive_sstp_cond_"+aerosol_name+".png" + plt.savefig(out_png, dpi=200) + + # cleanup generated NetCDF files for this figure + # for p in generated_nc_files: + # try: + # Path(p).unlink(missing_ok=True) + # except TypeError: + # try: + # if Path(p).exists(): + # Path(p).unlink() + # except OSError: + # pass + # except OSError: + # pass + + return fig + +make_figure('pristine', '{"DYCOMS": {"kappa": 0.61, "mean_r": [0.011e-6, 0.06e-6], "gstdev": [1.2, 1.7], "n_tot": [125.0e6, 65.0e6]}}', 200) +make_figure('polluted', '{"polluted": {"kappa": 0.61, "mean_r": [0.029e-6, 0.071e-6], "gstdev": [1.36, 1.57], "n_tot": [160.0e6, 380.0e6]}}', 600) +# plt.show()