Issue¶

The for loop in both functions slows down the script a lot. I have tried to speed it up by numba.

Solution¶

from numba import jit

@jit(nopython=True)
def calc_field(weights, spatial_index_lat, spatial_index_lon,
               surface_pressure_external, surface_pressure, field,
               result_field, result_pressure, hya, hyb):
    for pix_idx in range(weights.shape[1]):
        if surface_pressure_external is not None:
            ps = surface_pressure_external.flat[pix_idx]
        else:
            ps = 0.0 
        for i in range(4):
            if surface_pressure_external is None:
                ps += (weights[i, pix_idx] * 
                        surface_pressure[spatial_index_lat[i, pix_idx], 
                                         spatial_index_lon[i, pix_idx]])
            result_field[:, pix_idx] += (weights[i, pix_idx] * 
                                         field[:, spatial_index_lat[i, pix_idx], 
                                         spatial_index_lon[i, pix_idx]])

        result_pressure[:, pix_idx] = hya + ps * hyb

    return result_field, result_pressure

def get_temperature(self, surface_pressure_external=None):
    ...........
    result_field, result_pressure = calc_field(self.weights, self.spatial_index_lat, self.spatial_index_lon,
                                               surface_pressure_external, surface_pressure, field,
                                               result_field, result_pressure, hyam, hybm)

def get_profiles(self, surface_pressure_external=None):
    ...........
    result_field, result_pressure = calc_field(self.weights, self.spatial_index_lat, self.spatial_index_lon,
                                               surface_pressure_external, surface_pressure, field,
                                               result_field, result_pressure, hyai, hybi)

Note that I don't use the jitclass which needs to specify all variables of class.

I just put the jit function outside the class and use it directly.

Test¶

And the execution time of the loop in get_profiles decreased from ~ 470 seconds to ~ 2.5 seconds.

For my case, the maximum differences between current slow loop and numba are 3.5527137e-15 and 9.1552734e-05 respectively for no2_vmr and temperature.