Vectorize getFragments and dailySequences inner loops

Replace nested Python loops (stations × years × 366 days) with vectorized
numpy operations. The key change is computing conditions (missing data,
wet state, daily depth) as boolean vectors over all days at once instead
of calling np.min() per element (~2.6M numpy calls eliminated).

getFragments: 7.3s -> 1.2s (6x), dailySequences: 5.3s -> 0.5s (10x).
Overall request time: ~67s -> ~9.3s with all optimizations combined.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: joelrahman

src/pyraingen/dailysequences.py

@@ -2,21 +2,34 @@
 import numpy as np
 from datetime import date
 from numba.typed import List
-#import nvtx
 
 # Defined Functions
 from .jdtodatevec import jdToDateVec
-from .getdaysperyear import getDaysPerYear
-from .getwetstate import getWetState
 
 ## Global constants
 from .global_ import missingDay
 from .global_ import stateBad
 from .global_ import ndaysYearLeap
 from .global_ import idxfebTwentyNine
-from .global_ import yesterday, today, tomorrow
+from .global_ import stateWetWet, stateWetDry, stateDryWet, stateDryDry
+
+
+def _build_day_mapping(yearStart, yearEnd, idxYearBase):
+    """Build mapping from linear day index to (loopDay, idxYear) for a station."""
+    loopDay_list = []
+    idxYear_list = []
+    for yr_offset in range(yearEnd - yearStart + 1):
+        year = yearStart + yr_offset
+        is_leap = (year % 4 == 0 and (year % 100 != 0 or year % 400 == 0))
+        nDays = 366 if is_leap else 365
+        for d in range(ndaysYearLeap):
+            if d == idxfebTwentyNine and nDays != ndaysYearLeap:
+                continue
+            loopDay_list.append(d)
+            idxYear_list.append(idxYearBase + yr_offset)
+    return np.array(loopDay_list, dtype=np.intp), np.array(idxYear_list, dtype=np.intp)
+
 
-#@nvtx.annotate()
 def dailySequences(nSeasons,
                     nYearsPool,
                     stnDetails,
@@ -44,22 +57,23 @@ def dailySequences(nSeasons,
         dimensional with the first dimension being the seasons and the
         second the number of stations returned.  This is because there
         are possible different correlations between stations over
-        different seasons. 
+        different seasons.
     param : dict where keys are words and values are float
-        Dictionary of the run parameters. 
+        Dictionary of the run parameters.
     param_path : dict where keys are words and values are str
         Dictionary of the necessary paths.\n
 
     Returns
     ----------
     dailyDepth : list
-        A list of arrays (one for each season) containing the daily depth sequences. 
+        A list of arrays (one for each season) containing the daily depth sequences.
     dailyWetState : list
         A list of arrays (one for each season) containing the daily wetState sequences.
-    """                
+    """
     dailyDepth = List()
     dailyWetState = List()
-    workingRainDepth = np.zeros((3,1))
+    dryWetCutoff = param['dryWetCutoff']
+
     for loopSeason in range(nSeasons):
         # Allocate RAM for this daily series:
         nYears = int(nYearsPool[loopSeason].item())
@@ -71,115 +85,91 @@ def dailySequences(nSeasons,
             (ndaysYearLeap,
             nYears))
             *stateBad)
-        
+
+        depth_s = dailyDepth[loopSeason]
+        wetState_s = dailyWetState[loopSeason]
+
         #Counter through the year dimension of our arrays
         idxYear = 0
-        
+
         for loopStation in range(nearStationIdx[0,:].size):
             # Grab a conveniance variable:
             currStnIndex = int(nearStationIdx[loopSeason, loopStation])
 
             if currStnIndex == 0:
             # There are no more stations for this season
                 break
+
+            # Get the start and end years from the NetCDF then compute the
+            # number of years in the sequence.
+            stnIdx = int(stnDetails['stnIndex'][currStnIndex])
+            if station_cache is not None:
+                daySeries, rainfall_data = station_cache.get(stnIdx)
+                fnameNC = station_cache.filename(stnIdx)
             else:
-                # Get the start and end years from the NetCDF then compute the
-                # number of years in the sequence.  We also want to reshape our
-                # linear vector of rainfall data from the NetCDF into the
-                # year-on-year array herein.
-                stnIdx = int(stnDetails['stnIndex'][currStnIndex])
-                if station_cache is not None:
-                    daySeries, rainfall_data = station_cache.get(stnIdx)
-                    fnameNC = station_cache.filename(stnIdx)
-                else:
-                    import netCDF4 as nc
-                    fnameNC = ('{}/plv{:06}.nc'.format(param_path['pathSubDaily'], stnIdx))
-                    ds=nc.Dataset(fnameNC)
-                    daySeries = ds['day'][:].data
-                    rainfall_data = ds['rainfall'][:].data
-                print('Processing station:', fnameNC)
-                dayVecStart = jdToDateVec(daySeries[0])
-                dayVecEnd = jdToDateVec(daySeries[-1])
-                yearStart = int(dayVecStart[0])
-                yearEnd = int(dayVecEnd[0])
-                tmpSubDaily = rainfall_data/10
-                # The algorithm below works on the assumption that the
-                # tmpSubDaily array is populated with full years.  So pad out
-                # the data array to make full years with missingDay values.
-                nDaysKnown = (date.toordinal(date(yearEnd,12,31))
-                    - date.toordinal(date(yearStart,1,1))+1)
-                tmpDaily = np.ones((1,nDaysKnown)) * missingDay
-                # Get an index into the tmpDaily array where our known data
-                # should start.
-                dataIdxStart = (date.toordinal(date(
-                    int(dayVecStart[0]), int(dayVecStart[1]), int(dayVecStart[2])))
-                    - date.toordinal(date(yearStart,1,1)))
-                dataIdxEnd = (date.toordinal(date(
-                    int(dayVecEnd[0]),int(dayVecEnd[1]),int(dayVecEnd[2])))
-                    - date.toordinal(date(yearStart,1,1))+1)
-
-                # The sum here works because the method described in the paper
-                # is based on increments per day.  We do a bit of a fudge when
-                # dealing with daily read data aggregated to 9am.
-                tmpDaily[0,dataIdxStart:dataIdxEnd] = tmpSubDaily.sum(axis=2)
-
-                # This is the index into the linear array tmpDaily.
-                idxDayLinear = 0
-                # Loop over each year in this station
-                for loopYear in range(yearStart,yearEnd+1):
-                    # Loop over the days and compute the wetness index.  I
-                    # don't think that this can be parallelised as it works on
-                    # offset indexing.
-                    nDaysCurrYear =  getDaysPerYear(loopYear)
-                    #idxYear += 1
-
-                    #for loopDay in range(nDaysCurrYear):
-                    for loopDay in range(ndaysYearLeap):         
-                        # We are now looping through the days of a particular
-                        # year.
-                        if loopDay == 0 and loopYear == yearStart:
-                            # No previous day, assume dry as its a pretty good bet
-                            # in Australia, except in the tropics. :-[, got to
-                            # assume something.
-                            #
-                            # This is the first time through, so populate our
-                            # working array
-                            workingRainDepth[yesterday] = 0
-                            workingRainDepth[today]     = tmpDaily[0, idxDayLinear]
-                            workingRainDepth[tomorrow]  = tmpDaily[0, idxDayLinear + 1]
-                        elif (loopDay == idxfebTwentyNine and 
-                        nDaysCurrYear != ndaysYearLeap):
-                            # If we are on the 29th day of the year but this is
-                            # not a leap year increment loopDay by one as that
-                            # is the index we are looping through our 2D array
-                            # year with.  Or, just continue to the next
-                            # iteration of the day loop.  But don't index our
-                            # linear counter.
-                            continue
-                        elif loopDay == (ndaysYearLeap-1) and loopYear == yearEnd:
-                            # This is the end of our sequence, assume dry as
-                            # per the start
-                            workingRainDepth[tomorrow] = 0
-                        else:
-                            # This is a normal day out in the middle of the
-                            # sequence, so get ready for the next loop:
-                            idxDayLinear = idxDayLinear + 1
-                            workingRainDepth[yesterday] = workingRainDepth[today]
-                            workingRainDepth[today]     = workingRainDepth[tomorrow]
-                            workingRainDepth[tomorrow] = tmpDaily[0, idxDayLinear + 1]
-                            
-                    
-                        # Store the depth for this day
-                        dailyDepth[loopSeason][loopDay, idxYear] = (
-                                tmpDaily[0, idxDayLinear]
-                                )
-
-                        # Decide our wetness state for this day
-                        if np.all(workingRainDepth >= 0.0 - np.spacing(abs(workingRainDepth))):
-                            # But only if we have "good" data.  The default is
-                            # to leave it as zero
-                            dailyWetState[loopSeason][loopDay,idxYear] = (
-                                getWetState(workingRainDepth, param['dryWetCutoff'])
-                            )
-                    idxYear += 1
-    return dailyDepth, dailyWetState                
+                import netCDF4 as nc
+                fnameNC = ('{}/plv{:06}.nc'.format(param_path['pathSubDaily'], stnIdx))
+                ds=nc.Dataset(fnameNC)
+                daySeries = ds['day'][:].data
+                rainfall_data = ds['rainfall'][:].data
+            print('Processing station:', fnameNC)
+            dayVecStart = jdToDateVec(daySeries[0])
+            dayVecEnd = jdToDateVec(daySeries[-1])
+            yearStart = int(dayVecStart[0])
+            yearEnd = int(dayVecEnd[0])
+            nYearsStation = yearEnd - yearStart + 1
+
+            # Pad out the data array to make full years with missingDay values.
+            nDaysKnown = (date.toordinal(date(yearEnd,12,31))
+                - date.toordinal(date(yearStart,1,1))+1)
+            tmpDaily = np.full(nDaysKnown, missingDay)
+            dataIdxStart = (date.toordinal(date(
+                int(dayVecStart[0]), int(dayVecStart[1]), int(dayVecStart[2])))
+                - date.toordinal(date(yearStart,1,1)))
+            dataIdxEnd = (date.toordinal(date(
+                int(dayVecEnd[0]),int(dayVecEnd[1]),int(dayVecEnd[2])))
+                - date.toordinal(date(yearStart,1,1))+1)
+
+            # Sum sub-daily timesteps to get daily totals.
+            # rainfall_data has shape (1, nDays, 240); sum over subday axis
+            tmpDaily[dataIdxStart:dataIdxEnd] = (rainfall_data / 10).sum(axis=2).ravel()
+
+            # Build mapping from linear day index to (loopDay, idxYear)
+            loopDay_arr, idxYear_arr = _build_day_mapping(yearStart, yearEnd, idxYear)
+
+            # Store daily depths for all days at once
+            depth_s[loopDay_arr, idxYear_arr] = tmpDaily
+
+            # Vectorized 3-day sliding window for wet state computation.
+            # Yesterday is previous calendar day (0 for very first day).
+            # Tomorrow is next calendar day (0 for very last day).
+            yesterday_vals = np.empty(nDaysKnown)
+            yesterday_vals[0] = 0
+            yesterday_vals[1:] = tmpDaily[:-1]
+
+            tomorrow_vals = np.empty(nDaysKnown)
+            tomorrow_vals[-1] = 0
+            tomorrow_vals[:-1] = tmpDaily[1:]
+
+            # Good data: all three days non-negative (within floating point tolerance)
+            good_mask = (
+                (yesterday_vals >= -np.spacing(np.abs(yesterday_vals))) &
+                (tmpDaily >= -np.spacing(np.abs(tmpDaily))) &
+                (tomorrow_vals >= -np.spacing(np.abs(tomorrow_vals)))
+            )
+
+            # Vectorized wet state classification
+            y_wet = yesterday_vals > dryWetCutoff
+            t_wet = tomorrow_vals > dryWetCutoff
+
+            wet_states = np.full(nDaysKnown, stateBad)
+            wet_states[good_mask & y_wet & t_wet] = stateWetWet
+            wet_states[good_mask & y_wet & ~t_wet] = stateWetDry
+            wet_states[good_mask & ~y_wet & t_wet] = stateDryWet
+            wet_states[good_mask & ~y_wet & ~t_wet] = stateDryDry
+
+            wetState_s[loopDay_arr, idxYear_arr] = wet_states
+
+            idxYear += nYearsStation
+
+    return dailyDepth, dailyWetState