# Import packages

import sys
import numba as nb
import time
from typing import Tuple

# Set constants

SHR_CONST_TKFRZ = np.float64(273.15)
lambd_a = np.float64(3.504) # Inverse of Heat Capacity
alpha = np.float64(17.67) # Constant to calculate vapor pressure
beta = np.float64(243.5) # Constant to calculate vapor pressure
epsilon = np.float64(0.6220) # Conversion between pressure/mixing ratio
es_C = np.float64(611.2) # Vapor Pressure at Freezing STD ¶
y0 = np.float64(3036) # constant
y1 = np.float64(1.78) # constant
y2 = np.float64(0.448) # constant
Cf = SHR_CONST_TKFRZ # Freezing Temp (K)
p0 = np.float64(100000) # Reference Pressure ¶
constA = np.float64(2675) # Constant used for extreme cold temperatures (K)
vkp = np.float64(0.2854) # Heat Capacity

#Define QSat_2 function

@nb.njit(fastmath=True)
def QSat_2(T_k, p_t, p0ndplam):
# Constants used to calculate es(T)
# Clausius-Clapeyron
tcfbdiff = T_k - Cf + beta
es = es_C * np.exp(alpha*(T_k - Cf)/(tcfbdiff))
dlnes_dT = alpha * beta/((tcfbdiff)*(tcfbdiff))
pminuse = p_t - es
de_dT = es * dlnes_dT

``````# Constants used to calculate rs(T)
rs = np.empty_like(T_k)
rs = epsilon * es/(p0ndplam - es + np.spacing(1)) #eps

if rs > 1 or rs < 0:
rs = np.nan

return es,rs,dlnes_dT
``````

#Define main wet-bulb-temperature function

@nb.njit(fastmath=True)
def WetBulb(TemperatureC,Pressure,Humidity,HumidityMode=0):
###
#Unless necessary, default to using specific humidity as input (simpler and tends to reduce error margins)#
###

``````"""
INPUTS:
TemperatureC	   2-m air temperature (degrees Celsius)
Pressure	       Atmospheric Pressure (Pa)
Humidity         Humidity -- meaning depends on HumidityMode
HumidityMode
0 (Default): Humidity is specific humidity (kg/kg)
1: Humidity is relative humidity (#, max = 100)
TemperatureC, Pressure, and Humidity should either be scalars or arrays of identical dimension.
OUTPUTS:
Twb	    wet bulb temperature (C)
"""
TemperatureK = TemperatureC + SHR_CONST_TKFRZ
pnd = (Pressure/p0)**(vkp)
p0ndplam = p0*pnd**lambd_a

C = SHR_CONST_TKFRZ;	# Freezing Temperature
T1 = TemperatureK;		# Use holder for T

es, rs, _ = QSat_2(TemperatureK, Pressure, p0ndplam) # first two returned values

if HumidityMode==0:
qin = Humidity                   # specific humidity
mixr = (qin / (1-qin))           # corrected by Rob Warren
vape = (Pressure * mixr) / (epsilon + mixr) #corrected by Rob Warren
relhum = 100.0 * vape/es         # corrected by Rob Warren
elif HumidityMode==1:
relhum = Humidity                # relative humidity (%)
vape = es * relhum * 0.01        # vapor pressure (Pa)
mixr = epsilon * vape / (Pressure-vape)  #corrected by Rob Warren

mixr = mixr * 1000

# Calculate Equivalent Pot. Temp (Pressure, T, mixing ratio (g/kg), pott, epott)
# Calculate Parameters for Wet Bulb Temp (epott, Pressure)
D = 1.0/(0.1859*Pressure/p0 + 0.6512)
k1 = -38.5*pnd*pnd + 137.81*pnd - 53.737
k2 = -4.392*pnd*pnd + 56.831*pnd - 0.384

# Calculate lifting condensation level
tl = (1.0/((1.0/((T1 - 55))) - (np.log(relhum/100.0)/2840.0))) + 55.0

# Theta_DL: Bolton 1980 Eqn 24.
theta_dl = T1*((p0/(Pressure-vape))**vkp) * ((T1/tl)**(mixr*0.00028))
# EPT: Bolton 1980 Eqn 39.
epott = theta_dl * np.exp(((3.036/tl)-0.00178)*mixr*(1 + 0.000448*mixr))
Teq = epott*pnd	# Equivalent Temperature at pressure
X = (C/Teq)**3.504

# Calculates the regime requirements of wet bulb equations.
invalid = Teq > 600 or Teq < 200
hot = Teq > 355.15
cold = X<1   #corrected by Rob Warren

if invalid:
return np.nan #modified 1, previously -> if invalid: return np.nan, np.nan, epott

# Calculate Wet Bulb Temperature, initial guess
# Extremely cold regimes: if X.gt.D, then need to calculate dlnesTeqdTeq
es_teq, rs_teq, dlnes_dTeq = QSat_2(Teq, Pressure, p0ndplam)
if X<=D:
wb_temp = C + (k1 - 1.21 * cold - 1.45 * hot - (k2 - 1.21 * cold) * X + (0.58 / X) * hot)
else:
wb_temp = Teq - ((constA*rs_teq)/(1 + (constA*rs_teq*dlnes_dTeq)))

# Newton-Raphson Method
maxiter = 2
iter = 0
delta = 1e6

while delta>0.01 and iter<maxiter:
foftk_wb_temp, fdwb_temp = DJ(wb_temp, Pressure, p0ndplam)
delta = (foftk_wb_temp - X)/fdwb_temp
delta = np.minimum(10,delta)
delta = np.maximum(-10,delta) #max(-10,delta)
wb_temp = wb_temp - delta
Twb = wb_temp
iter = iter+1

Tw_final=np.round(Twb-C,2)

return Tw_final
``````

# Define parallelization functions for wet-bulb (optional)

@nb.njit(fastmath=True)
def WetBulb_all(tempC, Pres, relHum, Hum_mode):
Twb = np.empty_like(tempC)
Teq = np.empty_like(tempC)
epott = np.empty_like(tempC)
for i in nb.prange(Twb.size):
Twb[i], Teq[i], epott[i] = WetBulb(tempC[i], Pres[i], relHum[i], Hum_mode)

``````# change 18-10-23
return Twb
``````

@nb.njit(fastmath=True, parallel=True)
def WetBulb_par(tempC, Pres, relHum, Hum_mode):
Twb = np.empty_like(tempC)
Teq = np.empty_like(tempC)
epott = np.empty_like(tempC)
for i in nb.prange(Twb.size):
Twb[i], Teq[i], epott[i] = WetBulb(tempC[i], Pres[i], relHum[i], Hum_mode)

``````# change 18-10-23
return Twb
``````

# Define helper functions for usage in the Davies-Jones wet-bulb algorithm

@nb.njit(fastmath=True)
def DJ(T_k, p_t, p0ndplam):
# Constants used to calculate es(T)
# Clausius-Clapeyron
tcfbdiff = T_k - Cf + beta
es = es_C * np.exp(alpha*(T_k - Cf)/(tcfbdiff))
dlnes_dT = alpha * beta/((tcfbdiff)*(tcfbdiff))
pminuse = p_t - es
de_dT = es * dlnes_dT

``````# Constants used to calculate rs(T)
rs = epsilon * es/(p0ndplam - es + np.spacing(1)) #eps)
prersdt = epsilon * p_t/((pminuse)*(pminuse))
rsdT = prersdt * de_dT

# Constants used to calculate g(T)
rsy2rs2 = rs + y2*rs*rs
oty2rs = 1 + 2.0*y2*rs
y0tky1 = y0/T_k - y1
goftk = y0tky1 * (rs + y2 * rs * rs)
gdT = - y0 * (rsy2rs2)/(T_k*T_k) + (y0tky1)*(oty2rs)*rsdT

# Calculations for calculating f(T,ndimpress)
foftk = ((Cf/T_k)**lambd_a)*(1 - es/p0ndplam)**(vkp*lambd_a)* \
np.exp(-lambd_a*goftk)
fdT = -lambd_a*(1.0/T_k + vkp*de_dT/pminuse + gdT) * foftk  #derivative corrected by Qinqin Kong

return foftk,fdT
``````

import xarray as xr

date = ‘15’
month= ‘05’
hour_str=‘0600’

# Surface temperature in Kelvin at 2m

temp = xr.open_dataset(f’/gws/nopw/j04/h2x/era5/t2m/2023/{month}/{date}/era5.t2m.2023{month}{date}_{hour_str}.nc’)
temC = temp[‘t2m’].sel(expver=1).values.flatten() - 273.15

# Dew point temperature in Kelvin at 2m

dpt = xr.open_dataset(f’/gws/nopw/j04/h2x/era5/d2m/2023/{month}/{date}/era5.d2m.2023{month}{date}_{hour_str}.nc’)
temD = dpt[‘d2m’].sel(expver=1).values.flatten() - 273.15

# Surface pressure in Pa

sfc_pres = xr.open_dataset(f’/gws/nopw/j04/h2x/era5/p0/2023/{month}/{date}/era5.p0.2023{month}{date}_{hour_str}.nc’)
pres = sfc_pres[‘p0’].sel(expver=1).values.flatten()

relhum = calc_RH_from_T_Td(temC, temD, mode=0)

# Assuming temC, temD, pres, and relhum are NumPy arrays

shape_temC = np.squeeze(temC.shape)
shape_temD = np.squeeze(temD.shape)
shape_pres = np.squeeze(pres.shape)
shape_relhum = np.squeeze(relhum.shape)

if shape_temC == shape_temD == shape_pres == shape_relhum:
print(“All arrays have the same shape:”, shape_temC)
else:
print(“Arrays have different shapes.”)

Twb = WetBulb(temC, pres, relhum, 0)
Twb

When I try to print Twb it gave me this error:
TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
Cannot unify float64 and array(float64, 1d, C) for ‘rs.1.2’, defined at /tmp/ipykernel_5712/1734631780.py (43)

When I did this, it does not result in problem, but when i applied to the real data, it gave me above error:
do_test_calc = True
if do_test_calc:
print(‘Calculate test wet-bulb temperature - result should be 21.73C’)
Twb=WetBulb(25.,100000.,0.015,0)
print(Twb)

Hi @nonanistia, Thanks for sharing the issue and your code!
It seems like Numba is not able to decide the type you are returning. In your `QSat_2`-function,
you write:

``````rs = np.empty_like(T_k)
rs = epsilon * es/(p0ndplam - es + np.spacing(1)) #eps

if rs > 1 or rs < 0:
rs = np.nan
``````

Comment out the line with `#rs = np.empty_like(T_k) `, and try to re-run without creating this array.
This way `type(rs)` stays consistent ( `numpy.float64`.)

Does that help?

Edit:
When reading your code more in detail, I think the Error comes from `WetBulb()` being given arguments of varying type. In your example, you call `WetBulb()` with arguments of type (`float, float, float, float`) , while in your real-data example you seem to call it with `(np.ndarray, np.ndarray, np.ndarray, float)`. Numba needs to decide on the type of each variable to generate efficient machine code.

• What type of arguments do you want to call `WetBulb` with?

However, if you want to call your code with `np.ndarrays` (as it seems like since you use flattened np.ndarrays in your real-data-example), there are a few changes you will have to do. The ones I can identify are:

• Let all return statements in one function, return the same type
``````if np.any(rs > 1) or np.any(rs < 0)
rs[:] = np.nan
``````
``````if invalid:
return np.full(shape=TemperatureC.shape, fill_value=np.nan)
``````
• using np.any for the condition-checks (where necessary)
``````invalid = np.any(Teq > 600) or np.any(Teq < 200)
``````
• Call `np.round` with the `out`-argument. Something like:
``````out = np.empty_like(Twb)
Tw_final=np.round(Twb-C,2, out)
``````

Also, if you are concerned about runtime with the above workarounds, then there are several ways to avoid some of those numpy-functions entirely.