如何使用Numba加速输入与返回值均为字典的函数？

Great question! The short answer is that Numba's nopython mode has limited support for Python dictionaries—especially when they're used as input arguments or return values. But don't worry, there are practical workarounds to get the speed benefits while working with dictionaries. Let's break this down:

Option 1: Convert Dictionaries to Numba-Friendly Structured Arrays #

Numba shines with statically typed structures like structured arrays (which act like typed, array-based dictionaries). Here's how to adapt your function:

1. Define structured dtypes that match the keys and data types in your dictionaries:

import numba as nb
import numpy as np

# Match these to your actual dictionary keys and data types
aeolus_dtype = np.dtype([('wind_speed', np.float64), ('latitude', np.float64)])
valdata_dtype = np.dtype([('temperature', np.float64), ('pressure', np.float64)])

2. Convert input dictionaries to structured arrays before passing to the Numba-optimized function:

@nb.jit(nopython=True)
def collocation_numba(aeolus_arr, valdata_arr):
    # Your core numerical computation here
    # Access data like aeolus_arr['wind_speed'] or valdata_arr['temperature']
    processed_aeolus = np.empty_like(aeolus_arr)
    processed_valdata = np.empty_like(valdata_arr)
    
    # Example computation (replace with your logic)
    processed_aeolus['wind_speed'] = aeolus_arr['wind_speed'] * 1.1
    processed_aeolus['latitude'] = aeolus_arr['latitude']
    processed_valdata['temperature'] = valdata_arr['temperature'] + 2.0
    processed_valdata['pressure'] = valdata_arr['pressure']
    
    return processed_aeolus, processed_valdata

3. Wrap the optimized function to handle dictionary conversion:

def collocation(aeolus_data, val_data):
    # Convert dicts to structured arrays
    aeolus_arr = np.array(
        (aeolus_data['wind_speed'], aeolus_data['latitude']),
        dtype=aeolus_dtype
    )
    valdata_arr = np.array(
        (val_data['temperature'], val_data['pressure']),
        dtype=valdata_dtype
    )
    
    # Run optimized calculation
    processed_aeolus_arr, processed_valdata_arr = collocation_numba(aeolus_arr, valdata_arr)
    
    # Convert back to dictionaries for your output
    sample_aeolus = {
        'wind_speed': processed_aeolus_arr['wind_speed'],
        'latitude': processed_aeolus_arr['latitude']
    }
    sample_valdata = {
        'temperature': processed_valdata_arr['temperature'],
        'pressure': processed_valdata_arr['pressure']
    }
    
    return sample_aeolus, sample_valdata

Option 2: Unpack Dictionaries into Individual Arguments #

If your dictionaries only hold numerical arrays, you can avoid dictionaries entirely in the Numba-optimized code by unpacking values as separate arguments:

# Optimized core function with explicit array arguments
@nb.jit(nopython=True)
def collocation_core(aeolus_wind, aeolus_lat, val_temp, val_press):
    # Your core computation here using the individual arrays
    new_wind = aeolus_wind * 1.1
    new_lat = aeolus_lat
    new_temp = val_temp + 2.0
    new_press = val_press
    return new_wind, new_lat, new_temp, new_press

# Wrapper function to handle dictionaries
def collocation(aeolus_data, val_data):
    # Unpack dictionary values into separate variables
    aeolus_wind = aeolus_data['wind_speed']
    aeolus_lat = aeolus_data['latitude']
    val_temp = val_data['temperature']
    val_press = val_data['pressure']
    
    # Run optimized calculation
    new_wind, new_lat, new_temp, new_press = collocation_core(aeolus_wind, aeolus_lat, val_temp, val_press)
    
    # Repack results into dictionaries
    sample_aeolus = {'wind_speed': new_wind, 'latitude': new_lat}
    sample_valdata = {'temperature': new_temp, 'pressure': new_press}
    
    return sample_aeolus, sample_valdata

Key Things to Keep in Mind #

• Avoid dictionaries in nopython mode: While Numba supports simple dictionaries (string/integer keys, scalar/array values) inside nopython functions, using them as inputs often leads to type inference issues or falls back to unoptimized object mode.
• Stick to typed structures: Structured arrays and named tuples are designed for Numba's static typing system, so they're far more reliable for full optimization.
• Update Numba if possible: Newer Numba versions have improved support for structured types, so upgrading can help with edge cases.

内容的提问来源于stack exchange，提问作者Annerl