Merge pull request #1134 from xylar/fix-el-nino-main-vs-control

Fix control run support for El Nino 3.4 Index

Author: xylar

mpas_analysis/ocean/index_nino34.py

@@ -221,6 +221,9 @@ def run_task(self):
             ninoIndexNumber))
         varName = self.variableList[0]
         regionSST = ds[varName]
+        self.logger.debug('Main run SST dims=%s shape=%s',
+                          getattr(regionSST, 'dims', None),
+                          getattr(regionSST, 'shape', None))
         nino34Main = self._compute_nino34_index(regionSST, calendar)
 
         # Compute the observational index over the entire time range
@@ -270,7 +273,14 @@ def run_task(self):
             dsRef = add_standard_regions_and_subset(
                 dsRef, self.controlConfig, regionShortNames=[regionToPlot])
 
+            # we want to collapse the nOceanRegions dimension (same as main)
+            if 'nOceanRegions' in dsRef.dims:
+                dsRef = dsRef.isel(nOceanRegions=0)
+
             regionSSTRef = dsRef[varName]
+            self.logger.debug('Control run SST dims=%s shape=%s',
+                              getattr(regionSSTRef, 'dims', None),
+                              getattr(regionSSTRef, 'shape', None))
             nino34Ref = self._compute_nino34_index(regionSSTRef, calendar)
 
             nino34s = [nino34Subset, nino34Main[2:-3], nino34Ref[2:-3]]
@@ -499,7 +509,19 @@ def _running_mean(self, inputData, wgts):
         sp = (len(wgts) - 1) // 2
         runningMean = inputData.copy()
         for k in range(sp, nt - (sp + 1)):
-            runningMean[k] = sum(wgts * inputData[k - sp:k + sp + 1].values)
+            windowValues = np.asarray(inputData[k - sp:k + sp + 1].values)
+            if windowValues.shape[0] != len(wgts):
+                raise ValueError(
+                    'Unexpected running-mean window shape. '
+                    f'Expected first dimension {len(wgts)} but got {windowValues.shape}. '
+                    f'inputData dims={getattr(inputData, "dims", None)} '
+                    f'shape={getattr(inputData, "shape", None)}')
+
+            if windowValues.ndim == 1:
+                runningMean[k] = np.sum(wgts * windowValues)
+            else:
+                # weighted sum over the first axis (the running-mean window)
+                runningMean[k] = np.tensordot(wgts, windowValues, axes=(0, 0))
 
         return runningMean