Linear PSM: optional expected sign for season selection

cfr/psm.py

@@ -94,7 +94,14 @@ def __init__(self, pobj=None, climate_required=['tas']):
         self.climate_required = climate_required
 
     def calibrate(self, calib_period=None, nobs_lb=25, metric='fitR2adj',
-        season_list=[list(range(1, 13))], exog_name=None, **fit_args):
+        season_list=[list(range(1, 13))], exog_name=None, expected_sign=None, **fit_args):
+        ''' Calibrate the PSM, choosing the season from season_list with the best score.
+
+        Args:
+            expected_sign (int, optional): if given (1 or -1), only seasons whose regression slope has this sign
+                are considered, e.g., the sign of the record's interpreted relationship with the climate variable.
+                If no season has it, the PSM is not calibrated. Defaults to None (no constraint).
+        '''
         if exog_name is None:
             exog_name = f'obs.{self.climate_required[0]}'
 
@@ -127,6 +134,9 @@ def calibrate(self, calib_period=None, nobs_lb=25, metric='fitR2adj',
                 print(f'The number of overlapped data points is {nobs} < {nobs_lb}. Skipping ...')
             else:
                 mdl = smf.ols(formula=formula_spell, data=df).fit(**fit_args)
+                if expected_sign is not None and np.sign(mdl.params[exog_colname]) != np.sign(expected_sign):
+                    continue
+
                 fitR2adj =  mdl.rsquared_adj,
                 mse = np.mean(mdl.resid**2),
                 score = {

cfr/reconjob.py

@@ -443,7 +443,7 @@ def crop_clim(self, tag, lat_min=None, lat_max=None, lon_min=None, lon_max=None,
 
 
     def calib_psms(self, ptype_psm_dict=None, ptype_season_dict=None, ptype_clim_dict=None, calib_period=None,
-                   use_predefined_R=False, verbose=False, **kwargs):
+                   use_predefined_R=False, pid_sign_dict=None, verbose=False, **kwargs):
         ''' Calibrate the PSMs.
 
         Args:
@@ -452,6 +452,10 @@ def calib_psms(self, ptype_psm_dict=None, ptype_season_dict=None, ptype_clim_dic
             ptype_clim_dict (dict): the dictionary to denote the required climate variables for each proxy type; ['tas'] for all by default.
             calib_period (tuple or list): the time period for calibration.
             use_predefined_R (bool): use the predefined observation error covariance instead of by calibration.
+            pid_sign_dict (dict): the dictionary to denote the expected sign (1 or -1) of the response of each record
+                to the climate variable, e.g., from its interpretation metadata; for records with the Linear PSM,
+                only seasons whose slope has this sign are considered (see :py:meth:`cfr.psm.Linear.calibrate`).
+                Records not in the dictionary are unconstrained. Defaults to None (no constraint).
             verbose (bool, optional): print verbose information. Defaults to False.
         '''
         ptype_psm_dict_default = {ptype: 'Linear' for ptype in set(self.proxydb.type_list)}
@@ -502,6 +506,9 @@ def calib_psms(self, ptype_psm_dict=None, ptype_season_dict=None, ptype_clim_dic
                     season_list1=ptype_season_dict[pobj.ptype], season_list2=ptype_season_dict[pobj.ptype],
                     calib_period=calib_period, **kwargs,
                 )
+            elif psm_name == 'Linear' and pid_sign_dict is not None and pid in pid_sign_dict:
+                pobj.psm.calibrate(season_list=ptype_season_dict[pobj.ptype], calib_period=calib_period,
+                                   expected_sign=pid_sign_dict[pid], **kwargs)
             else:
                 pobj.psm.calibrate(season_list=ptype_season_dict[pobj.ptype], calib_period=calib_period, **kwargs)
 

tests/test_psm.py

@@ -102,6 +102,37 @@ def calibrated_linear(self):
         )
         return lpsm
 
+    def _two_season_proxy(self):
+        # the proxy responds positively to Jan-Jun and more strongly, negatively, to Jul-Dec
+        pobj = make_proxy_with_climate(seed=2, n_calib=40)
+        da = pobj.clim['obs.tas'].da
+        h1 = da.sel(time=da['time.month'] <= 6).groupby('time.year').mean().values.squeeze()
+        h2 = da.sel(time=da['time.month'] > 6).groupby('time.year').mean().values.squeeze()
+        pobj.value = 0.5*h1 - 1.0*h2
+        return pobj
+
+    def _calibrate(self, pobj, **kwargs):
+        lpsm = psm.Linear(pobj=pobj, climate_required=['tas'])
+        lpsm.calibrate(nobs_lb=10, season_list=[[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]],
+                       exog_name='obs.tas', **kwargs)
+        return lpsm
+
+    def test_season_by_score_alone(self):
+        lpsm = self._calibrate(self._two_season_proxy())
+        assert lpsm.calib_details['seasonality'] == [7, 8, 9, 10, 11, 12]
+
+    def test_expected_sign_constrains_season(self):
+        lpsm = self._calibrate(self._two_season_proxy(), expected_sign=1)
+        assert lpsm.calib_details['seasonality'] == [1, 2, 3, 4, 5, 6]
+        assert lpsm.model.params['tas'] > 0
+
+    def test_expected_sign_with_no_matching_season(self):
+        pobj = self._two_season_proxy()
+        pobj.value = -pobj.value  # now Jan-Jun negative and Jul-Dec positive
+        lpsm = psm.Linear(pobj=pobj, climate_required=['tas'])
+        lpsm.calibrate(nobs_lb=10, season_list=[[1, 2, 3, 4, 5, 6]], exog_name='obs.tas', expected_sign=1)
+        assert lpsm.calib_details is None
+
     def test_calibrate_sets_model(self, calibrated_linear):
         assert calibrated_linear.model is not None