def terminal_velocity(self, values, radius, k1, k2, k3, r1, r2):
k1 = PrecisionResolver.get_floating_point(k1)
k2 = PrecisionResolver.get_floating_point(k2)
k3 = PrecisionResolver.get_floating_point(k3)
r1 = PrecisionResolver.get_floating_point(r1)
r2 = PrecisionResolver.get_floating_point(r2)
self.__terminal_velocity_body.launch_n(values.size(), [values, radius, k1, k2, k3, r1, r2])
def c_inline(fun, **args):
prae = r"([,+\-*/( ]|^)"
post = r"([ )/*\-+,]|$)"
real_t = PrecisionResolver.get_C_type()
real_fmt = ".32g"
source = ''
for lineno, line in enumerate(inspect.getsourcelines(fun)[0]):
stripped = line.strip()
if stripped.startswith('@'):
continue
if stripped.startswith('//'):
continue
if stripped.startswith('def '):
continue
source += stripped
source = source.replace("power(", "pow(")
source = re.sub("^return ", "", source)
for arg in inspect.signature(fun).parameters:
source = re.sub(f"{prae}({arg}){post}", f"\\1{real_t}({args[arg]})\\3",
source)
source = re.sub(f"{prae}const\\.([^\\d\\W]\\w*]*){post}",
"\\1" + real_t + "({const.\\2:" + real_fmt + "})\\3",
source)
source = eval(f'f"""{source}"""')
return f'{real_t}({source})'
def thrust(obj):
if isinstance(obj, list):
result = [thrust(o) for o in obj]
elif hasattr(obj, 'data'):
result = obj.data
elif isinstance(obj, float):
result = PrecisionResolver.get_floating_point(obj)
elif isinstance(obj, int):
result = trtc.DVInt64(obj)
else:
raise ValueError(f"Cannot upload {obj} to device.")
return result
def __init__(self):
phys = self.formulae
self._temperature_pressure_RH_body = trtc.For(["rhod", "thd", "qv", "T", "p", "RH"], "i", f'''
T[i] = {phys.state_variable_triplet.T.c_inline(rhod="rhod[i]", thd="thd[i]")};
p[i] = {phys.state_variable_triplet.p.c_inline(rhod="rhod[i]", T="T[i]", qv="qv[i]")};
RH[i] = {phys.state_variable_triplet.pv.c_inline(p="p[i]", qv="qv[i]")} / {phys.saturation_vapour_pressure.pvs_Celsius.c_inline(T="T[i] - const.T0")};
'''.replace("real_type", PrecisionResolver.get_C_type()))
self.__explicit_euler_body = trtc.For(("y", "dt", "dy_dt"), "i", f'''
y[i] = {phys.trivia.explicit_euler.c_inline(y="y[i]", dt="dt", dy_dt="dy_dt")};
'''.replace("real_type", PrecisionResolver.get_C_type()))
self.__critical_volume_body = trtc.For(("v_cr", "kappa", "v_dry", "v_wet", "T", "cell"), "i", f'''
auto sigma = {phys.surface_tension.sigma.c_inline(T="T[cell[i]]", v_wet="v_wet[i]", v_dry="v_dry[i]")};
auto r_cr = {phys.hygroscopicity.r_cr.c_inline(
kp="kappa",
rd3="v_dry[i] / const.pi_4_3",
T="T[cell[i]]",
sgm="sigma"
)};
v_cr[i] = {phys.trivia.volume.c_inline(radius="r_cr")};
'''.replace("real_type", PrecisionResolver.get_C_type()))
self.__terminal_velocity_body = trtc.For(["values", "radius", "k1", "k2", "k3", "r1", "r2"], "i", '''
if (radius[i] < r1) {
values[i] = k1 * radius[i] * radius[i];
}
else {
if (radius[i] < r2) {
values[i] = k2 * radius[i];
}
else {
values[i] = k3 * pow(radius[i], (real_type)(.5));
}
}
'''.replace("real_type", PrecisionResolver.get_C_type()))
def _get_empty_data(shape, dtype):
if dtype in (float, Storage.FLOAT):
elem_cls = PrecisionResolver.get_C_type()
dtype = Storage.FLOAT
elif dtype in (int, Storage.INT):
elem_cls = 'int64_t'
dtype = Storage.INT
elif dtype in (bool, Storage.BOOL):
elem_cls = 'bool'
dtype = Storage.BOOL
else:
raise NotImplementedError
data = trtc.device_vector(elem_cls, int(np.prod(shape)))
return data, shape, dtype
def __setitem__(self, key, value):
if hasattr(value, 'data') and hasattr(value, 'shape') and len(value.shape) != 0:
if isinstance(value, np.ndarray):
vector = trtc.device_vector_from_numpy(value)
trtc.Copy(vector, self.data)
else:
trtc.Copy(value.data, self.data)
else:
if isinstance(value, int):
dvalue = trtc.DVInt64(value)
elif isinstance(value, float):
dvalue = PrecisionResolver.get_floating_point(value)
else:
raise TypeError("Only Storage, int and float are supported.")
trtc.Fill(self.data, dvalue)
return self
def _to_host(self):
if isinstance(self.data, trtc.DVVector.DVRange):
if self.dtype is Storage.FLOAT:
elem_cls = PrecisionResolver.get_C_type()
elif self.dtype is Storage.INT:
elem_cls = 'int64_t'
elif self.dtype is Storage.BOOL:
elem_cls = 'bool'
else:
raise NotImplementedError()
data = trtc.device_vector(elem_cls, self.data.size())
trtc.Copy(self.data, data)
else:
data = self.data
return data.to_host()
class Storage:
FLOAT = PrecisionResolver.get_np_dtype()
INT = np.int64
BOOL = np.bool_
def __init__(self, data, shape, dtype):
self.data = data
self.shape = (shape,) if isinstance(shape, int) else shape
self.dtype = dtype
def __getitem__(self, item):
dim = len(self.shape)
if isinstance(item, slice):
start = item.start or 0
stop = item.stop or self.shape[0]
if dim == 1:
result_data = self.data.range(start, stop)
result_shape = (stop - start,)
elif dim == 2:
result_data = self.data.range(self.shape[1] * start, self.shape[1] * stop)
result_shape = (stop - start, self.shape[1])
else:
raise NotImplementedError("Only 2 or less dimensions array is supported.")
result = Storage(result_data, result_shape, self.dtype)
elif isinstance(item, tuple) and dim == 2 and isinstance(item[0], int) and isinstance(item[1], slice):
assert item[1].start is None or item[1].start == 0
assert item[1].stop is None or item[1].stop == self.shape[1]
assert item[1].step is None or item[1].step == 1
result_data = self.data.range(self.shape[1] * item[0], self.shape[1] * (item[0] + 1))
result = Storage(result_data, (*self.shape[1:],), self.dtype)
else:
result = self.to_ndarray()[item]
return result
def __setitem__(self, key, value):
if hasattr(value, 'data') and hasattr(value, 'shape') and len(value.shape) != 0:
if isinstance(value, np.ndarray):
vector = trtc.device_vector_from_numpy(value)
trtc.Copy(vector, self.data)
else:
trtc.Copy(value.data, self.data)
else:
if isinstance(value, int):
dvalue = trtc.DVInt64(value)
elif isinstance(value, float):
dvalue = PrecisionResolver.get_floating_point(value)
else:
raise TypeError("Only Storage, int and float are supported.")
trtc.Fill(self.data, dvalue)
return self
def __add__(self, other):
raise TypeError("Use +=")
def __iadd__(self, other):
impl.add(self, other)
return self
def __sub__(self, other):
raise TypeError("Use -=")
def __isub__(self, other):
impl.subtract(self, other)
return self
def __mul__(self, other):
raise TypeError("Use *=")
def __imul__(self, other):
impl.multiply(self, other)
return self
def __truediv__(self, other):
raise TypeError("Use /=")
def __itruediv__(self, other):
impl.truediv(self, other)
return self
def __mod__(self, other):
raise TypeError("Use %=")
def __imod__(self, other):
impl.row_modulo(self, other)
return self
def __pow__(self, other):
raise TypeError("Use **=")
def __ipow__(self, other):
impl.power(self, other)
return self
def __len__(self):
return self.shape[0]
def __bool__(self):
if len(self) == 1:
result = bool(self.data.to_host()[0] != 0)
else:
raise NotImplementedError("Logic value of array is ambiguous.")
return result
def _to_host(self):
if isinstance(self.data, trtc.DVVector.DVRange):
if self.dtype is Storage.FLOAT:
elem_cls = PrecisionResolver.get_C_type()
elif self.dtype is Storage.INT:
elem_cls = 'int64_t'
elif self.dtype is Storage.BOOL:
elem_cls = 'bool'
else:
raise NotImplementedError()
data = trtc.device_vector(elem_cls, self.data.size())
trtc.Copy(self.data, data)
else:
data = self.data
return data.to_host()
def amin(self):
return impl.amin(self.data)
def all(self):
assert self.dtype is Storage.BOOL
return self.amin()
def download(self, target, reshape=False):
shape = target.shape if reshape else self.shape
target[:] = np.reshape(self._to_host(), shape)
@staticmethod
def _get_empty_data(shape, dtype):
if dtype in (float, Storage.FLOAT):
elem_cls = PrecisionResolver.get_C_type()
dtype = Storage.FLOAT
elif dtype in (int, Storage.INT):
elem_cls = 'int64_t'
dtype = Storage.INT
elif dtype in (bool, Storage.BOOL):
elem_cls = 'bool'
dtype = Storage.BOOL
else:
raise NotImplementedError
data = trtc.device_vector(elem_cls, int(np.prod(shape)))
return data, shape, dtype
@staticmethod
def empty(shape, dtype):
result = Storage(*Storage._get_empty_data(shape, dtype))
return result
@staticmethod
def _get_data_from_ndarray(array):
if str(array.dtype).startswith('int'):
dtype = Storage.INT
elif str(array.dtype).startswith('float'):
dtype = Storage.FLOAT
elif str(array.dtype).startswith('bool'):
dtype = Storage.BOOL
else:
raise NotImplementedError()
data = trtc.device_vector_from_numpy(array.astype(dtype).ravel())
return data, array.shape, dtype
@staticmethod
def from_ndarray(array):
result = Storage(*Storage._get_data_from_ndarray(array))
return result
def floor(self, other=None):
if other is None:
impl.floor(self.data)
else:
impl.floor_out_of_place(self, other)
return self
def product(self, multiplicand, multiplier):
impl.multiply_out_of_place(self, multiplicand, multiplier)
return self
def ratio(self, dividend, divisor):
impl.divide_out_of_place(self, dividend, divisor)
return self
def ravel(self, other):
if isinstance(other, Storage):
trtc.Copy(other.data, self.data)
else:
self.data = trtc.device_vector_from_numpy(other.ravel())
def to_ndarray(self):
result = self._to_host()
result = np.reshape(result, self.shape)
return result
def urand(self, generator=None):
generator(self)
def upload(self, data):
trtc.Copy(
trtc.device_vector_from_numpy(data.astype(self.dtype).ravel()),
self.data
)
def explicit_euler(self, y, dt, dy_dt):
dt = PrecisionResolver.get_floating_point(dt)
dy_dt = PrecisionResolver.get_floating_point(dy_dt)
self.__explicit_euler_body.launch_n(y.shape[0], (y.data, dt, dy_dt))
def critical_volume(self, v_cr, kappa, v_dry, v_wet, T, cell):
kappa = PrecisionResolver.get_floating_point(kappa)
self.__critical_volume_body.launch_n(
v_cr.shape[0],
(v_cr.data, kappa, v_dry.data, v_wet.data, T.data, cell.data))
class PhysicsMethods:
def __init__(self):
phys = self.formulae
self._temperature_pressure_RH_body = trtc.For(
["rhod", "thd", "qv", "T", "p", "RH"], "i", f'''
T[i] = {c_inline(phys.state_variable_triplet.T, rhod="rhod[i]", thd="thd[i]")};
p[i] = {c_inline(phys.state_variable_triplet.p, rhod="rhod[i]", T="T[i]", qv="qv[i]")};
RH[i] = {c_inline(phys.state_variable_triplet.pv, p="p[i]", qv="qv[i]")} / {c_inline(phys.saturation_vapour_pressure.pvs_Celsius, T="T[i] - const.T0")};
''')
self.__explicit_euler_body = trtc.For(("y", "dt", "dy_dt"), "i", f'''
y[i] = {c_inline(phys.trivia.explicit_euler, y="y[i]", dt="dt", dy_dt="dy_dt")};
''')
self.__critical_volume_body = trtc.For(
("v_cr", "kappa", "v_dry", "v_wet", "T", "cell"), "i", f'''
auto sigma = {c_inline(phys.surface_tension.sigma, T="T[cell[i]]", v_wet="v_wet[i]", v_dry="v_dry[i]")};
auto r_cr = {c_inline(phys.hygroscopicity.r_cr,
kp="kappa",
rd3="v_dry[i] / const.pi_4_3",
T="T[cell[i]]",
sgm="sigma"
)};
v_cr[i] = {c_inline(phys.trivia.volume, radius="r_cr")};
''')
@nice_thrust(**NICE_THRUST_FLAGS)
def critical_volume(self, v_cr, kappa, v_dry, v_wet, T, cell):
kappa = PrecisionResolver.get_floating_point(kappa)
self.__critical_volume_body.launch_n(
v_cr.shape[0],
(v_cr.data, kappa, v_dry.data, v_wet.data, T.data, cell.data))
@nice_thrust(**NICE_THRUST_FLAGS)
def temperature_pressure_RH(self, rhod, thd, qv, T, p, RH):
self._temperature_pressure_RH_body.launch_n(
T.shape[0],
(rhod.data, thd.data, qv.data, T.data, p.data, RH.data))
__terminal_velocity_body = trtc.For(
["values", "radius", "k1", "k2", "k3", "r1", "r2"], "i", '''
if (radius[i] < r1) {
values[i] = k1 * radius[i] * radius[i];
}
else {
if (radius[i] < r2) {
values[i] = k2 * radius[i];
}
else {
values[i] = k3 * pow(radius[i], (real_type)(.5));
}
}
'''.replace("real_type", PrecisionResolver.get_C_type()))
@staticmethod
@nice_thrust(**NICE_THRUST_FLAGS)
def terminal_velocity(values, radius, k1, k2, k3, r1, r2):
k1 = PrecisionResolver.get_floating_point(k1)
k2 = PrecisionResolver.get_floating_point(k2)
k3 = PrecisionResolver.get_floating_point(k3)
r1 = PrecisionResolver.get_floating_point(r1)
r2 = PrecisionResolver.get_floating_point(r2)
PhysicsMethods.__terminal_velocity_body.launch_n(
values.size(), [values, radius, k1, k2, k3, r1, r2])
@nice_thrust(**NICE_THRUST_FLAGS)
def explicit_euler(self, y, dt, dy_dt):
dt = PrecisionResolver.get_floating_point(dt)
dy_dt = PrecisionResolver.get_floating_point(dy_dt)
self.__explicit_euler_body.launch_n(y.shape[0], (y.data, dt, dy_dt))
class PhysicsMethods:
@staticmethod
@nice_thrust(**NICE_THRUST_FLAGS)
def explicit_in_space(omega, c_l, c_r):
return "c_l * (1 - omega) + c_r * omega;"
@staticmethod
@nice_thrust(**NICE_THRUST_FLAGS)
def implicit_in_space(omega, c_l, c_r):
"""
see eqs 14-16 in Arabas et al. 2015 (libcloudph++)
"""
result = "(omega * (c_r - c_l) + c_l) / (1 - (c_r - c_l));"
return result
__temperature_pressure_RH_body = trtc.For(
["rhod", "thd", "qv", "T", "p", "RH"], "i", f'''
// equivalent to eqs A11 & A12 in libcloudph++ 1.0 paper
real_type exponent = {const.Rd} / {const.c_pd};
real_type pd = pow((rhod[i] * {const.Rd} * thd[i]) / pow({const.p1000}, exponent), 1 / (1 - exponent));
T[i] = thd[i] * pow((pd / {const.p1000}), exponent);
real_type R = {const.Rv} / (1 / qv[i] + 1) + {const.Rd} / (1 + qv[i]);
p[i] = rhod[i] * (1 + qv[i]) * R * T[i];
// August-Roche-Magnus formula
real_type pvs = {const.ARM_C1} * exp(({const.ARM_C2} * (T[i] - {const.T0})) / (T[i] - {const.T0} + {const.ARM_C3}));
RH[i] = (p[i] - pd) / pvs;
'''.replace("real_type", PrecisionResolver.get_C_type()))
@staticmethod
@nice_thrust(**NICE_THRUST_FLAGS)
def temperature_pressure_RH(rhod, thd, qv, T, p, RH):
PhysicsMethods.__temperature_pressure_RH_body.launch_n(
T.shape[0],
(rhod.data, thd.data, qv.data, T.data, p.data, RH.data))
__terminal_velocity_body = trtc.For(
["values", "radius", "k1", "k2", "k3", "r1", "r2"], "i", '''
if (radius[i] < r1) {
values[i] = k1 * radius[i] * radius[i];
}
else {
if (radius[i] < r2) {
values[i] = k2 * radius[i];
}
else {
values[i] = k3 * pow(radius[i], (real_type).5);
}
}
'''.replace("real_type", PrecisionResolver.get_C_type()))
@staticmethod
@nice_thrust(**NICE_THRUST_FLAGS)
def terminal_velocity(values, radius, k1, k2, k3, r1, r2):
k1 = PrecisionResolver.get_floating_point(k1)
k2 = PrecisionResolver.get_floating_point(k2)
k3 = PrecisionResolver.get_floating_point(k3)
r1 = PrecisionResolver.get_floating_point(r1)
r2 = PrecisionResolver.get_floating_point(r2)
PhysicsMethods.__terminal_velocity_body.launch_n(
values.size(), [values, radius, k1, k2, k3, r1, r2])
@staticmethod
@nice_thrust(**NICE_THRUST_FLAGS)
def radius(volume):
return ""
@staticmethod
@nice_thrust(**NICE_THRUST_FLAGS)
def dr_dt_MM(r, T, p, RH, kp, rd):
return ""
@staticmethod
@nice_thrust(**NICE_THRUST_FLAGS)
def dr_dt_FF(r, T, p, qv, kp, rd, T_i):
return ""
@staticmethod
@nice_thrust(**NICE_THRUST_FLAGS)
def dthd_dt(rhod, thd, T, dqv_dt):
return ""