Adaptive sstp

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):

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,23 +10,75 @@ 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"])
 
   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]
+    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]