refactor code into separate files: IO, lgrngn, blk1m, blk1mice

.github/workflows/test_parcel.yml

@@ -61,4 +61,4 @@ jobs:
 
     - name: run unit tests debug
       working-directory: ${{github.workspace}}/parcel
-      run: apptainer exec -B${{ github.workspace }}/installed $SI python3 -m pytest -s -v unit_test_debug
+      run: apptainer exec -B${{ github.workspace }}/installed $SI python3 -m pytest -s -v unit_test_debug

io_output.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python
+from scipy.io import netcdf
+import numpy as np
+from parcel_common import _Chem_g_id, _Chem_a_id
+
+
+def _output_bins(fout, t, micro, opts, spectra):
+  for dim, dct in spectra.items():
+    for bin in range(dct["nbin"]):
+      if dct["drwt"] == 'wet':
+        micro.diag_wet_rng(
+          fout.variables[dim+"_r_wet"][bin],
+          fout.variables[dim+"_r_wet"][bin] + fout.variables[dim+"_dr_wet"][bin]
+        )
+      elif dct["drwt"] == 'dry':
+        micro.diag_dry_rng(
+          fout.variables[dim+"_r_dry"][bin],
+          fout.variables[dim+"_r_dry"][bin] + fout.variables[dim+"_dr_dry"][bin]
+        )
+      else: raise Exception("drwt should be wet or dry")
+
+      for vm in dct["moms"]:
+        if type(vm) == int:
+          # calculating moments
+          if dct["drwt"] == 'wet':
+            micro.diag_wet_mom(vm)
+          elif dct["drwt"] == 'dry':
+            micro.diag_dry_mom(vm)
+          else: raise Exception("drwt should be wet or dry")
+          fout.variables[dim+'_m'+str(vm)][int(t), int(bin)] = np.frombuffer(micro.outbuf())
+        else:
+          # calculate chemistry
+          micro.diag_chem(_Chem_a_id[vm])
+          fout.variables[dim+'_'+vm][int(t), int(bin)] = np.frombuffer(micro.outbuf())
+
+
+def _output_init(micro, opts, spectra):
+  """Initialize output file for lagrangian scheme"""
+  # file & dimensions
+  fout = netcdf.netcdf_file(opts["outfile"], 'w')
+  fout.createDimension('t', None)
+  for name, dct in spectra.items():
+    fout.createDimension(name, dct["nbin"])
+
+    tmp = name + '_r_' + dct["drwt"]
+    fout.createVariable(tmp, 'd', (name,))
+    fout.variables[tmp].unit = "m"
+    fout.variables[tmp].description = "particle wet radius (left bin edge)"
+
+    tmp = name + '_dr_' + dct["drwt"]
+    fout.createVariable(tmp, 'd', (name,))
+    fout.variables[tmp].unit = "m"
+    fout.variables[tmp].description = "bin width"
+
+    if dct["lnli"] == 'log':
+      from math import exp, log
+      dlnr = (log(dct["rght"]) - log(dct["left"])) / dct["nbin"]
+      allbins = np.exp(log(dct["left"]) + np.arange(dct["nbin"]+1) * dlnr)
+      fout.variables[name+'_r_'+dct["drwt"]][:] = allbins[0:-1]
+      fout.variables[name+'_dr_'+dct["drwt"]][:] = allbins[1:] - allbins[0:-1]
+    elif dct["lnli"] == 'lin':
+      dr = (dct["rght"] - dct["left"]) / dct["nbin"]
+      fout.variables[name+'_r_'+dct["drwt"]][:] = dct["left"] + np.arange(dct["nbin"]) * dr
+      fout.variables[name+'_dr_'+dct["drwt"]][:] = dr
+    else: raise Exception("lnli should be log or lin")
+
+    for vm in dct["moms"]:
+      if (vm in _Chem_a_id):
+        fout.createVariable(name+'_'+vm, 'd', ('t',name))
+        fout.variables[name+'_'+vm].unit = 'kg of chem species dissolved in cloud droplets (kg of dry air)^-1'
+      else:
+        assert(type(vm)==int)
+        fout.createVariable(name+'_m'+str(vm), 'd', ('t',name))
+        fout.variables[name+'_m'+str(vm)].unit = 'm^'+str(vm)+' (kg of dry air)^-1'
+
+  units = {"z"  : "m",     "t"   : "s",     "r_v"  : "kg/kg", "th_d" : "K", "rhod" : "kg/m3",
+           "p"  : "Pa",    "T"   : "K",     "RH"   : "1",    "T_blk"   : "K",     "RH_blk"   : "1"
+  }
+
+  if micro.opts_init.chem_switch:
+    for id_str in _Chem_g_id.keys():
+      units[id_str] = "gas mixing ratio [kg / kg dry air]"
+      units[id_str.replace('_g', '_a')] = "kg of chem species (both undissociated and ions) dissolved in cloud droplets (kg of dry air)^-1"
+
+  for var_name, unit in units.items():
+    fout.createVariable(var_name, 'd', ('t',))
+    fout.variables[var_name].unit = unit
+
+  return fout
+
+
+def _output_init_blk_1m(opts):
+    """Initialize output file for bulk microphysics scheme without ice processes"""
+    fout = netcdf.netcdf_file(opts["outfile"], 'w')
+    fout.createDimension('t', None)
+
+    # Basic variables with their units
+    vars_units = {
+        "z": ("m",), 
+        "t": ("s",),
+        "r_v": ("kg/kg",),
+        "rc": ("kg/kg",),
+        "rr": ("kg/kg",),
+        "th_d": ("K",),
+        "rhod": ("kg/m3",),
+        "p": ("Pa",),
+        "T": ("K",),
+        "RH": ("1",),
+        "T_blk": ("K",),
+        "RH_blk": ("1",)
+    }
+
+    for var, (unit,) in vars_units.items():
+        fout.createVariable(var, 'd', ('t',))
+        fout.variables[var].unit = unit
+
+    return fout
+
+
+def _output_init_blk_1m_ice(opts):
+    """Initialize output file for bulk ice microphysics scheme"""
+    fout = netcdf.netcdf_file(opts["outfile"], 'w')
+    fout.createDimension('t', None)
+
+    # Basic variables with their units
+    vars_units = {
+        "z": ("m",), 
+        "t": ("s",),
+        "r_v": ("kg/kg",),
+        "rc": ("kg/kg",),
+        "rr": ("kg/kg",),
+        "ria": ("kg/kg",),
+        "rib": ("kg/kg",),
+        "th_d": ("K",),
+        "rhod": ("kg/m3",),
+        "p": ("Pa",),
+        "T": ("K",),
+        "RH": ("1",),
+        "T_blk": ("K",),
+        "RH_blk": ("1",)
+    }
+
+    for var, (unit,) in vars_units.items():
+        fout.createVariable(var, 'd', ('t',))
+        fout.variables[var].unit = unit
+
+    return fout
+
+
+def _output_save(fout, state, rec):
+  for var, val in state.items():
+    fout.variables[var][int(rec)] = val
+
+
+def _save_attrs(fout, dictnr):
+  for var, val in dictnr.items():
+    setattr(fout, var, val)
+
+
+def _output(fout, opts, micro, state, rec, spectra):
+  _output_bins(fout, rec, micro, opts, spectra)
+  _output_save(fout, state, rec)

micro_blk_1m.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python
+import numpy as np
+from libcloudphxx import blk_1m
+from micro_blk_1m_common import _opts_init_blk_1m_common
+from parcel_common import _stats
+
+def _opts_init_blk_1m():
+    """Initialize options for bulk microphysics scheme"""
+    opts_init = blk_1m.opts_t()
+    opts_init = _opts_init_blk_1m_common(opts_init)
+
+    # no ice proceesses:
+    opts_init.homA1 = False
+    opts_init.homA2 = False
+    opts_init.hetA = False
+    opts_init.hetB = False 
+    opts_init.depA = False 
+    opts_init.depB = False
+    opts_init.rimA = False
+    opts_init.rimB = False 
+    opts_init.melA = False
+    opts_init.melB = False
+
+    return opts_init
+
+
+def _micro_step_blk_1m(micro_opts, state, info, opts):
+  """Microphysics step for bulk scheme without ice processes"""
+  # get state variables as numpy arrays
+  p = np.asarray(state["p"])
+  dot_th_d = np.zeros_like(state["th_d"])
+  dot_rv = np.zeros_like(state["r_v"])
+  dot_rc = np.zeros_like(state["rc"])
+  dot_rr = np.zeros_like(state["rr"])
+
+  blk_1m.adj_cellwise(
+    micro_opts,
+    state["rhod"],
+    p,
+    state["th_d"],
+    state["r_v"],
+    state["rc"],
+    state["rr"],
+    opts["dt"]
+  )
+
+  blk_1m.rhs_cellwise_revap(
+    micro_opts,
+    dot_th_d,
+    dot_rv,
+    dot_rc,
+    dot_rr,
+    state["rhod"],
+    p,
+    state["th_d"],
+    state["r_v"],
+    state["rc"],
+    state["rr"],
+    opts["dt"]
+  )
+
+  state["th_d"] += dot_th_d * opts["dt"]
+  state["r_v"] += dot_rv * opts["dt"]
+  state["rc"] += dot_rc * opts["dt"]
+  state["rr"] += dot_rr * opts["dt"]
+
+  # Update thermodynamic state
+  _stats(state, info)
+

micro_blk_1m_common.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python
+def _opts_init_blk_1m_common(opts_init):    
+    opts_init.accr = False  # no rain accretion
+    opts_init.conv = False  # no autoconversion
+    opts_init.sedi = False  # no sedimentation
+
+    opts_init.adj_nwtrph = True
+    opts_init.th_dry = True
+    opts_init.const_p = False
+
+    return opts_init

micro_blk_1m_ice.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python
+import numpy as np
+from libcloudphxx import blk_1m
+from micro_blk_1m_common import _opts_init_blk_1m_common
+from parcel_common import _stats
+
+def _opts_init_blk_1m_ice():
+    """Initialize options for bulk microphysics scheme with ice processes"""
+    opts_init = blk_1m.opts_t()
+    opts_init = _opts_init_blk_1m_common(opts_init)
+
+    # no collisions for ice:
+    opts_init.hetB = False
+    opts_init.rimA = False
+    opts_init.rimB = False
+
+    return opts_init
+
+def _micro_step_blk_1m_ice(micro_opts, state, info, opts):
+  """Microphysics step for bulk scheme with ice processes"""
+  # get state variables as numpy arrays
+  p = np.asarray(state["p"])
+  dot_th_d = np.zeros_like(state["th_d"])
+  dot_rv = np.zeros_like(state["r_v"])
+  dot_rc = np.zeros_like(state["rc"])
+  dot_rr = np.zeros_like(state["rr"])
+  dot_ria = np.zeros_like(state["ria"])
+  dot_rib = np.zeros_like(state["rib"])
+
+  blk_1m.adj_cellwise(
+    micro_opts,
+    state["rhod"],
+    p,
+    state["th_d"],
+    state["r_v"],
+    state["rc"],
+    state["rr"],
+    opts["dt"]
+    )
+
+  blk_1m.rhs_cellwise_ice(
+    micro_opts,
+    dot_th_d,
+    dot_rv,
+    dot_rc,
+    dot_rr,
+    dot_ria,
+    dot_rib,
+    state["rhod"],
+    p,
+    state["th_d"],
+    state["r_v"],
+    state["rc"],
+    state["rr"],
+    state["ria"],
+    state["rib"],
+    opts["dt"]
+  )
+
+  state["th_d"] += dot_th_d * opts["dt"]
+  state["r_v"] += dot_rv * opts["dt"]
+  state["rc"] += dot_rc * opts["dt"]
+  state["rr"] += dot_rr * opts["dt"]
+  state["ria"] += dot_ria * opts["dt"]
+  state["rib"] += dot_rib * opts["dt"]
+
+  # Update thermodynamic state
+  _stats(state, info)