Skip to content

Commit 9af9277

Browse files
slayooslayoo
authored
dry-radius initialisation by computing equilibrium size for a given input wet radii (incl. API module-name change: initialisation.equilibrate_wet_radii -> initialisation.hygroscopic_equilibrium) (#1676)
1 parent 3dd55a6 commit 9af9277

File tree

21 files changed

+389
-233
lines changed

21 files changed

+389
-233
lines changed

‎PySDM/environments/kinematic_1d.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from PySDM.environments.impl.moist import Moist
99

1010
from PySDM.impl import arakawa_c
11-
from PySDM.initialisation.equilibrate_wet_radii import equilibrate_wet_radii
11+
from PySDM.initialisation.hygroscopic_equilibrium import equilibrate_wet_radii
1212
from PySDM.environments.impl import register_environment
1313

1414

‎PySDM/environments/kinematic_2d.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77

88
from PySDM.environments.impl.moist import Moist
99
from PySDM.impl.mesh import Mesh
10-
from PySDM.initialisation.equilibrate_wet_radii import (
10+
from PySDM.initialisation.hygroscopic_equilibrium import (
1111
default_rtol,
1212
equilibrate_wet_radii,
1313
)

‎PySDM/environments/parcel.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99
from PySDM.environments.impl.moist import Moist
1010
from PySDM.impl.mesh import Mesh
11-
from PySDM.initialisation.equilibrate_wet_radii import (
11+
from PySDM.initialisation.hygroscopic_equilibrium import (
1212
default_rtol,
1313
equilibrate_wet_radii,
1414
)

‎PySDM/initialisation/__init__.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55

66
from . import sampling, spectra
77
from .discretise_multiplicities import discretise_multiplicities
8-
from .equilibrate_wet_radii import equilibrate_wet_radii
98
from .init_fall_momenta import init_fall_momenta
109

1110
from . import aerosol_composition # isort:skip

‎PySDM/initialisation/equilibrate_wet_radii.py‎

Lines changed: 0 additions & 129 deletions
This file was deleted.

‎PySDM/initialisation/hygroscopic_equilibrium.py‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
"""
2+
Koehler-curve equilibrium in unsaturated conditions
3+
"""
4+
5+
import numba
6+
import numpy as np
7+
8+
from ..backends.impl_numba.conf import JIT_FLAGS
9+
from ..backends.impl_numba.toms748 import toms748_solve
10+
from ..backends.impl_numba.warnings import warn
11+
12+
default_rtol = 1e-5
13+
default_max_iters = 64
14+
15+
# pylint: disable=too-many-locals,too-many-arguments
16+
17+
18+
def _solve_equilibrium_radii(
19+
*,
20+
radii_in: np.ndarray,
21+
get_bounds,
22+
get_args,
23+
minfun,
24+
environment,
25+
kappa,
26+
f_org: np.ndarray,
27+
cell_id: np.ndarray,
28+
rtol: float,
29+
max_iters: int,
30+
skip_fa_lt_zero: bool,
31+
RH_range: tuple = (0, 1),
32+
):
33+
T = environment["T"].to_ndarray()
34+
RH = environment["RH"].to_ndarray()
35+
formulae = environment.particulator.formulae
36+
within_tolerance = formulae.trivia.within_tolerance
37+
jit_flags = {**JIT_FLAGS, **{"fastmath": formulae.fastmath}}
38+
39+
if cell_id is None:
40+
cell_id = np.zeros_like(radii_in, dtype=int)
41+
if f_org is None:
42+
f_org = np.zeros_like(radii_in, dtype=float)
43+
44+
@numba.njit(**jit_flags)
45+
def impl(radii_in, iters, T, RH, cell_id, kappa, f_org):
46+
radii_out = np.empty_like(radii_in)
47+
for i in numba.prange(len(radii_in)): # pylint: disable=not-an-iterable
48+
cid = cell_id[i]
49+
a, b = get_bounds(radii_in[i], T[cid], kappa[i])
50+
51+
if not a < b:
52+
radii_out[i] = radii_in[i]
53+
iters[i] = 0
54+
continue
55+
56+
RH_i = np.maximum(RH_range[0], np.minimum(RH_range[1], RH[cid]))
57+
args = get_args(T[cid], RH_i, kappa[i], radii_in[i], f_org[i])
58+
fa, fb = minfun(a, *args), minfun(b, *args)
59+
60+
if skip_fa_lt_zero and fa < 0:
61+
radii_out[i] = radii_in[i]
62+
iters[i] = 0
63+
continue
64+
65+
radii_out[i], iters[i] = toms748_solve(
66+
minfun,
67+
args,
68+
a,
69+
b,
70+
fa,
71+
fb,
72+
rtol=rtol,
73+
max_iter=max_iters,
74+
within_tolerance=within_tolerance,
75+
)
76+
if iters[i] == -1:
77+
warn(
78+
msg="failed to find equilibrium particle size",
79+
file=__file__,
80+
context=(
81+
"i",
82+
i,
83+
"r",
84+
radii_in[i],
85+
"T",
86+
T[cid],
87+
"RH",
88+
RH_i,
89+
"f_org",
90+
f_org[i],
91+
"kappa",
92+
kappa[i],
93+
),
94+
)
95+
return radii_out
96+
97+
iters = np.empty_like(radii_in, dtype=int)
98+
radii_out = impl(radii_in, iters, T, RH, cell_id, kappa, f_org)
99+
assert (iters != max_iters).all() and (iters != -1).all()
100+
return radii_out
101+
102+
103+
def equilibrate_dry_radii(
104+
*,
105+
r_wet: np.ndarray,
106+
environment,
107+
kappa: np.ndarray,
108+
f_org: np.ndarray = None,
109+
cell_id: np.ndarray = None,
110+
rtol=default_rtol,
111+
max_iters=default_max_iters,
112+
):
113+
formulae = environment.particulator.formulae
114+
sigma = formulae.surface_tension.sigma
115+
phys_volume = formulae.trivia.volume
116+
const = environment.particulator.formulae.constants
117+
RH_eq = formulae.hygroscopicity.RH_eq
118+
jit_flags = {**JIT_FLAGS, **{"fastmath": formulae.fastmath}}
119+
120+
@numba.njit(**{**jit_flags, "parallel": False})
121+
def get_args(T_i, RH_i, kappa, r_wet, f_org):
122+
return T_i, RH_i, kappa, r_wet, f_org
123+
124+
@numba.njit(**{**jit_flags, "parallel": False})
125+
def get_bounds(r_wet, _, __):
126+
return 0.0, r_wet
127+
128+
@numba.njit(**{**jit_flags, "parallel": False})
129+
def minfun_dry(r_dry, temperature, relative_humidity, kappa, r_wet, f_org):
130+
r_dry_3 = r_dry**3
131+
sgm = sigma(
132+
temperature, phys_volume(radius=r_wet), const.PI_4_3 * r_dry_3, f_org
133+
)
134+
return relative_humidity - RH_eq(r_wet, temperature, kappa, r_dry_3, sgm)
135+
136+
return _solve_equilibrium_radii(
137+
radii_in=r_wet,
138+
get_args=get_args,
139+
get_bounds=get_bounds,
140+
minfun=minfun_dry,
141+
environment=environment,
142+
kappa=kappa,
143+
f_org=f_org,
144+
cell_id=cell_id,
145+
rtol=rtol,
146+
max_iters=max_iters,
147+
skip_fa_lt_zero=False,
148+
)
149+
150+
151+
def equilibrate_wet_radii(
152+
*,
153+
r_dry: np.ndarray,
154+
environment,
155+
kappa_times_dry_volume: np.ndarray,
156+
f_org: np.ndarray = None,
157+
cell_id: np.ndarray = None,
158+
rtol=default_rtol,
159+
max_iters=default_max_iters,
160+
):
161+
formulae = environment.particulator.formulae
162+
sigma = formulae.surface_tension.sigma
163+
phys_volume = formulae.trivia.volume
164+
const = environment.particulator.formulae.constants
165+
RH_eq = formulae.hygroscopicity.RH_eq
166+
r_cr = formulae.hygroscopicity.r_cr
167+
jit_flags = {**JIT_FLAGS, **{"fastmath": formulae.fastmath}}
168+
169+
@numba.njit(**{**jit_flags, "parallel": False})
170+
def get_args(T_i, RH_i, kappa, r_dry, f_org):
171+
return T_i, RH_i, kappa, r_dry**3, f_org
172+
173+
@numba.njit(**{**jit_flags, "parallel": False})
174+
def get_bounds(r_dry, T_i, kappa):
175+
a = r_dry
176+
b = r_cr(kappa, r_dry**3, T_i, const.sgm_w)
177+
return a, b
178+
179+
@numba.njit(**{**jit_flags, "parallel": False})
180+
def minfun_wet(r_wet, temperature, relative_humidity, kappa, r_dry_3, f_org):
181+
sgm = sigma(
182+
temperature, phys_volume(radius=r_wet), const.PI_4_3 * r_dry_3, f_org
183+
)
184+
return relative_humidity - RH_eq(r_wet, temperature, kappa, r_dry_3, sgm)
185+
186+
return _solve_equilibrium_radii(
187+
radii_in=r_dry,
188+
get_args=get_args,
189+
get_bounds=get_bounds,
190+
minfun=minfun_wet,
191+
environment=environment,
192+
kappa=kappa_times_dry_volume / phys_volume(radius=r_dry),
193+
f_org=f_org,
194+
cell_id=cell_id,
195+
rtol=rtol,
196+
max_iters=max_iters,
197+
skip_fa_lt_zero=True,
198+
)

‎docs/markdown/pysdm_landing.md‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -386,7 +386,7 @@ si = pyimport("PySDM.physics").si
386386
spectral_sampling = pyimport("PySDM.initialisation.sampling").spectral_sampling
387387
discretise_multiplicities = pyimport("PySDM.initialisation").discretise_multiplicities
388388
Lognormal = pyimport("PySDM.initialisation.spectra").Lognormal
389-
equilibrate_wet_radii = pyimport("PySDM.initialisation").equilibrate_wet_radii
389+
equilibrate_wet_radii = pyimport("PySDM.initialisation.hygroscopic_equilibrium").equilibrate_wet_radii
390390
CPU = pyimport("PySDM.backends").CPU
391391
AmbientThermodynamics = pyimport("PySDM.dynamics").AmbientThermodynamics
392392
Condensation = pyimport("PySDM.dynamics").Condensation
@@ -467,7 +467,7 @@ si = py.importlib.import_module('PySDM.physics').si;
467467
spectral_sampling = py.importlib.import_module('PySDM.initialisation.sampling').spectral_sampling;
468468
discretise_multiplicities = py.importlib.import_module('PySDM.initialisation').discretise_multiplicities;
469469
Lognormal = py.importlib.import_module('PySDM.initialisation.spectra').Lognormal;
470-
equilibrate_wet_radii = py.importlib.import_module('PySDM.initialisation').equilibrate_wet_radii;
470+
equilibrate_wet_radii = py.importlib.import_module('PySDM.initialisation.hygroscopic_equilibrium').equilibrate_wet_radii;
471471
CPU = py.importlib.import_module('PySDM.backends').CPU;
472472
AmbientThermodynamics = py.importlib.import_module('PySDM.dynamics').AmbientThermodynamics;
473473
Condensation = py.importlib.import_module('PySDM.dynamics').Condensation;
@@ -568,7 +568,8 @@ saveas(gcf, "parcel.png");
568568
```Python
569569
from matplotlib import pyplot
570570
from PySDM.physics import si
571-
from PySDM.initialisation import discretise_multiplicities, equilibrate_wet_radii
571+
from PySDM.initialisation import discretise_multiplicities
572+
from PySDM.initialisation.hygroscopic_equilibrium import equilibrate_wet_radii
572573
from PySDM.initialisation.spectra import Lognormal
573574
from PySDM.initialisation.sampling import spectral_sampling
574575
from PySDM.backends import CPU

0 commit comments

Comments
 (0)