Jitting SMRT to improve processing time

[ECCCBen]

Below is an example on how to use Jax jit function to compile and improve computation time when running SMRT. See here for more information on Jax.

To be able to run the code below, an extra library has to be installed: NumPyro. It can be installed using conda on Linux and pip on Windows.

The code below uses parameters and snow properties used in Montpetit et al., 2024:

from smrt.core.globalconstants import PERMITTIVITY_OF_AIR
from smrt import sensor_list, make_model, make_snowpack
from smrt.substrate.reflector_backscatter import make_reflector

from jax import jit
from functools import partial

import numpy as np

##Initialize the IBA/DORT SMRT model
smrt_model = make_model("iba", "dort", rtsolver_options = {'error_handling':'nan', 'phase_normalization' : True})
#Initialize the background part of the snowpack
sub = make_reflector(temperature=253.7, specular_reflection=0)
#Initialize the sensor
sensor  = sensor_list.active(13.285e9, 35)

#Snowpack properties
thickness=[0.15,0.28]
density= np.array([335.8, 229. ])
temperature=np.array([246.0, 253.5])
p_ex=np.array([0.74,1.11])*np.array([8.485580e-05, 2.615078e-04])

sp = make_snowpack(thickness=thickness, 
                       microstructure_model='exponential',
                       density= density,
                       temperature=temperature,
                       ice_permittivity_model=None,
                       background_permittivity_model=PERMITTIVITY_OF_AIR,
                       liquid_water=0, salinity=0, 
                       corr_length = p_ex,
                       substrate = sub)

# Defining normal SMRT function
def smrtSim(sensor, sp):

    sigma_nought = smrt_model.run(sensor, sp, parallel_computation=True)

    return sigma_nought.sigmaVV_dB()

# Defining jitted SMRT function with static input variables
@partial(jit, static_argnames=['sensor','sp'])
def jitted_smrtSim(sensor, sp):

    sigma_nought = smrt_model.run(sensor, sp, parallel_computation=True)

    return sigma_nought.sigmaVV_dB()

%timeit print(smrtSim(sensor, sp))

-17.673293907762805
-17.673293907762805
-17.673293907762805
-17.673293907762805
-17.673293907762805
-17.673293907762805
-17.673293907762805
-17.673293907762805
1.7 s ± 88.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit print(jitted_smrtSim(sensor, sp))

-17.673294
-17.673294
-17.673294
-17.673294
-17.673294
-17.673294
-17.673294
-17.673294
The slowest run took 18.16 times longer than the fastest. This could mean that an intermediate result is being cached.
48.6 µs ± 78.1 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Jitted SMRT provides same output with lower precision (1e-6) but runs ~35 000 times faster than normal SMRT on average...