रैग्ड स्लाइस वाले बड़े NumPy ऐरे को भरना तेज़ करें: Numba JIT समाधान

रैग्ड प्रीफिक्स स्लाइस के साथ बड़े NumPy ऐरे भरने में वेक्टराइज़ेशन रुकता है. वही पायथन लूप Numba JIT से कम्पाइल करें: 10× तक स्पीडअप, साफ इंडेक्सिंग और भरोसेमंद कोड.

बड़े NumPy ऐरे को कुशलता से भरना तब मुश्किल हो जाता है जब हर पंक्ति में अपडेट होने वाले तत्वों की संख्या अलग हो. ऐसे पैटर्न को वेक्टराइज़ करने की कोशिश करते ही आप जल्दी ही इंडेक्सिंग की सीमाओं या ब्रॉडकास्टिंग की उलझनों में फँस जाते हैं. सीधा Python लूप काम तो करता है, लेकिन इंडेक्स सेट बहुत बड़ा हो (दस मिलियन से कहीं अधिक इटरेशनों का सोचें) तो यह बेहद धीमा पड़ जाता है. एल्गोरिदम बदले बिना इसे तेज करने का साफ-सुथरा तरीका यहाँ दिया है.

समस्या विवरण

आपके पास m×n का एक ऐरे है और तीन 1×m ऐरे हैं, जो भराई को नियंत्रित करते हैं. पहले से सॉर्ट की गई सूची के हर इंडेक्स के लिए आप एक मान निकालते हैं और उसे कोण-बिन से चुनी गई पंक्ति के प्रीफिक्स स्लाइस में लिखते हैं. अवधारणात्मक रूप से यह ऐसा दिखता है: “किसी दिए गए कोण वाली पंक्ति में कॉलम [0:elev) तक वही मान लिखो”. लूप के बिना इसे ब्रॉडकास्ट करने की कोशिश करें तो गैर-पूर्णांक इंडेक्स या असंगत शेप जैसी समस्याएँ आती हैं. सीधा लूप सही परिणाम देता है, लेकिन करोड़ों इटरेशनों पर यही बाधा बन जाता है.

बेसलाइन कोड जो चलता है, पर धीमा है

नीचे दिया लूप वही अपडेट पैटर्न दिखाता है जो सही परिणाम देता है, लेकिन बहुत बड़े इनपुट पर समय लेता है.

import numpy as np
# grid: m x n ऐरे
# heights, radii, ang_bins: 1 x m ऐरे
# order_idx: इंडेक्स ऐरे (पहले से कहीं और सॉर्ट किया गया)
def fill_naive(grid, order_idx, heights, radii, ang_bins):
    for j in order_idx:
        stop = heights[j]
        val = 1 - (radii[j] / 10813)
        grid[int(ang_bins[j]), 0:stop] = val

लूप हटाने की कोशिश में आप शायद सीधे ragged slices पर advanced indexing जैसा कुछ आज़माएँ, और NumPy इंडेक्सिंग की जानी-पहचानी पाबंदियों से टकराएँ.

IndexError: arrays used as indices must be of integer (or boolean) type

ऐसा इसलिए होता है कि अपडेट में हर पंक्ति के लिए कॉलम का दायरा अलग चाहिए, और मानक NumPy इंडेक्सिंग एक ही अभिव्यक्ति में प्रति-पंक्ति बदलती लंबाई वाले स्लाइस स्वीकार नहीं करती.

यहाँ वेक्टराइज़ेशन कठिन क्यों है

ऑपरेशन का मूल एक ragged write है: हर चुनी पंक्ति को उसकी अपनी elevation सीमा तक ही भरना है. यानी आपके पास कोई आयताकार ब्लॉक नहीं है जिसे एक बार में असाइन किया जा सके, और साथ ही पंक्ति के इंडेक्स पूर्णांक भी होने चाहिए. ये दोनों शर्तें मिलकर सीधे-सीधे वेक्टराइज़ेशन को या तो अमान्य बनाती हैं या बोझिल, और बिना बड़े फेरबदल के लूप हटाना संभव नहीं रहता.

तेज़ और सरल: लूप को Numba से कम्पाइल करें

लूप छोड़ने की जरूरत नहीं. बस Numba के @njit से उसे कम्पाइल करें ताकि Python ओवरहेड हट जाए और तर्क वैसा ही बना रहे. यह तरीका कोड को पठनीय रखता है और अच्छा खासा स्पीडअप देता है.

from numba import njit
@njit
def paint_spans(canvas, idx_sorted, elev_vec, rad_vec, theta_vec):
    for j in idx_sorted:
        end = elev_vec[j]
        value = 1 - (rad_vec[j] / 10813)
        canvas[int(theta_vec[j]), 0:end] = value

अपनी ऐरेज़ और पहले से सॉर्ट किए गए इंडेक्स लिस्ट के साथ इस फ़ंक्शन को कॉल करें. एल्गोरिदम हूबहू वही रहता है; लूप नेटिव गति से चलता है. व्यवहार में यह पहले ही लगभग 10× सुधार दिखाता है.

यह क्यों मायने रखता है

जब इंडेक्स सूची करोड़ों में होती है, तो इंटरप्रेटर का ओवरहेड हावी हो जाता है, भले ही हर इटरेशन साधारण हो. यहाँ डेटा एक्सेस पैटर्न शुद्ध NumPy वेक्टराइज़ेशन में साफ़ नहीं बैठता, क्योंकि हर पंक्ति में भरने का स्पैन अलग है. Numba से JIT-कम्पाइल करना दोनों दिक्कतें दूर कर देता है, और कोडबेस भी सादा और संभालने लायक रहता है.

मुख्य बातें

यदि आपको प्रति पंक्ति अलग-लंबाई के स्लाइस असाइन करने हैं, तो जबरन वेक्टराइज़ेशन न करें. स्पष्ट लूप रखें, पर उसे JIT-कम्पाइल करें. पंक्ति इंडेक्स को पूर्णांक रखें, प्रति-पंक्ति स्लाइस [0:elev) ही रखें, और भरने का मान 1 - (r / 10813) के रूप में निकालें. @njit के साथ आप शुद्धता बनाए रखते हैं और बिना किसी उलझन के ज़रूरी प्रदर्शन हासिल करते हैं.