TypingError: Failed in nopython mode pipeline (step: nopython frontend) Failed in nopython mode pipeline (step: nopython frontend)

Numba Support: How do I do ...?
1

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

# avoid bad numbers
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

Define the date

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)

2

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

# avoid bad numbers
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.