domino.prediction

  1import numpy as np
  2import xarray as xr
  3from sklearn.linear_model import LogisticRegression
  4from sklearn.metrics import roc_auc_score
  5
  6
  7def CV_None(predictors,target,**cv_kwargs):
  8    """No cross validation: train and test data are the same."""
  9    return [[predictors,target,predictors,target]]
 10
 11def CV_drop1year(predictors,target,**cv_kwargs):
 12    """Cross validation is performed by dropping out single years of data from the training set to use as test data. This is repeated for all years with non-nan data."""
 13
 14    X,y=xr.align(predictors.dropna('time',how='all'),target.dropna('time',how='all'))
 15    years=np.unique(X['time.year'])
 16    iters=[]
 17    for t in years:
 18        it=[]
 19        it.append(X.where(X['time.year']!=t).dropna('time'))
 20        it.append(y.where(y['time.year']!=t).dropna('time'))
 21        it.append(X.where(X['time.year']==t).dropna('time'))
 22        it.append(y.where(y['time.year']==t).dropna('time'))
 23        iters.append(it)
 24    return iters
 25
 26def CV_n_chunks(predictors,target,N=2):
 27    """The dataset is split into *N* chunks and each is successively held out as test data."""
 28    X,y=xr.align(predictors,target)
 29    L=len(y)
 30    l=L//N
 31    iters=[]
 32    for n in range(N):
 33        
 34        test_times=np.arange(n*l,(n+1)*l)
 35        train_times=np.arange(L)
 36        train_times=np.delete(train_times,test_times)
 37        
 38        X_train=X.isel(time=train_times)
 39        y_train=y.isel(time=train_times)                             
 40        X_test=X.isel(time=test_times)
 41        y_test=y.isel(time=test_times)
 42        
 43        iters.append([X_train,y_train,X_test,y_test])
 44    return iters
 45
 46def CV_manual(train_pred,train_targ,test_pred=None,test_targ=None):
 47    """Separate training and test datasets are passed manually."""
 48    if test_pred is None or test_targ is None:
 49        raise(ValueError('Both test_pred and test_targ must be specified when using CV_manual'))
 50    return [[train_pred,train_targ,test_pred,test_targ]]
 51
 52def log_regression_model(X_train,y_train,X_test,**model_kwargs):
 53    """A wrapper around sklearn's *LogisticRegression* class, which accepts the same *model_kwargs*"""
 54    LR=LogisticRegression(**model_kwargs)
 55    LR=LR.fit(X_train,y_train)
 56    
 57    det_pred=LR.predict(X_test)
 58    prob_pred=LR.predict_proba(X_test)
 59    return LR,det_pred,prob_pred
 60
 61def blank_score(det_pred,prob_pred,target,**score_kwargs):
 62    """@private"""
 63    if det_pred is not None:
 64        assert det_pred.shape[0]==target.shape[0]
 65    if prob_pred is not None:
 66        assert prob_pred.shape[0]==target.shape[0]
 67    return len(target)
 68
 69def ROC_AUC(det_pred,prob_pred,target,**score_kwargs):
 70    """Computes the ROC AUC score, as implemented in sklearn, and accepting the same *score_kwargs*"""
 71    #No ROC_AUC for a variable with only one value
 72    if len(np.unique(target))==1:
 73        return np.nan
 74    if len(np.unique(target))==2:
 75        prob_pred=prob_pred[:,1]
 76        
 77    
 78    return roc_auc_score(target,prob_pred,**score_kwargs)
 79
 80class PredictionTest(object):
 81    """
 82    Supports quick and simple testing of the predictive power of predictor variables stored in an xarray.Dataset applied to a target variable stored in an xarray.DataArray.
 83    Supports different predictive models, cross validation approaches and skill scores within a straightforward API.
 84    
 85    **Arguments**
 86    
 87    *predictors*
 88    
 89    An xarray Dataset of scalar predictors
 90    
 91    *predictand*
 92    
 93    An xarray DataArray with a shared coord to *predictors*.
 94    """
 95    def __init__(self,predictors,predictand,tolerate_empty_variables=False):
 96        """Initialise a PredictionTest object"""
 97        
 98        self.predictand=self._predictand_to_dataarray(predictand)
 99        """@private"""
100        
101        self.predictors=self._predictors_to_dataset(predictors)
102        """@private"""
103
104        
105        self.predefined_models=dict(
106            logregression=log_regression_model
107        )
108        """@private"""
109        
110        self.predefined_cv_methods=dict(
111            nchunks=CV_n_chunks,
112            drop_year=CV_drop1year,
113            manual=CV_manual
114        )
115        """@private"""
116        
117        self.predefined_scores=dict(
118            sklearn_roc_auc=ROC_AUC,
119            test=blank_score
120        )
121        """@private"""
122        
123        self.computed_scores=None
124        """@private"""
125        
126        self.tol_empty=tolerate_empty_variables
127        """@private"""
128        return
129    
130    def __repr__(self):
131        return 'A PredictionTest object'
132    def __str__(self):
133        return self.__repr__
134    
135    def _predictand_to_dataarray(self,predictand):
136        return predictand
137    
138    def _predictors_to_dataset(self,predictors):
139        return predictors
140    
141    def _handle_potentially_empty_variable_list(self,vlist):
142        if len(vlist)>0:
143            return 0
144        else:
145            if self.tol_empty:
146                return 1
147            raise(ValueError('Empty variable list passed, but tolerate_empty_variables=False was set in PredictionTest.init'))
148
149    def categorical_prediction(self,model,score='sklearn_roc_auc',cv_method=None,predictor_variables='univariate',keep_models=False,model_kwargs={},cv_kwargs={},score_kwargs={}):
150        
151        """
152        
153        **Arguments**
154        
155        *model*
156        
157        A function with the following signature:
158        
159        Input: Three numpy arrays of shape [T1,N], [T1], [T2,N] (train_predictors, train_target, and test_predictors respectively), and an arbitrary number of keyword arguments.
160        
161        Output: *model* (A scikit-learn-like fitted model object),*det_pred* (a numpy array of shape [T2] with deterministic predictions) and *prob_pred* (a numpy array of shape [T2,M] with probabilistic predictions, where M is the number of unique values in *train_target* and each element predicts the probability of a given class). Any output can instead be replaced with None, but at least one of *det_pred* and *prob_pred* must not be None.
162            
163        **Optional arguments**
164        
165        *score*
166        
167        A string specifying a predefined skill score (currently only 'sklearn_roc_auc') or a function with the signature:
168        
169        Input: det_pred (a numpy array of shape [T2]), prob_pred (a numpy array of shape [T2,M], a truth series (a numpy array of shape [T2]), and an arbitrary number of keyword arguments.
170        Output: a scalar skill score.
171            
172        *cv_method*
173        A string specifying a predefined cross-validation method, a custom cross-validation function with corresponding signature, or None, in which case no cross-validation is used (the training and test datasets are the same). Predefined cross-validation methods are:
174        
175        *nchunks* - Divide the dataset into *n* chunks, using each as the test dataset predicted by the other *n*-1 chunks, to produce *n* total skill estimates. *n* defaults to 2, and is specified in *cv_kwargs*
176        
177        *drop_year* - Split the dataset by calendar year, using each year as the test dataset predicted by the remaining years.
178        *manual* - Treat *predictors* and *predictand* as training data. Test data must be passed explicitly via *cv_kwargs* as *test_pred* and *test_targ*.
179        
180        If a custom function is passed it must have the following signature:
181        Input: predictors (a Dataset), target (a DataArray), and an arbitrary number of keyword arguments.
182        
183        Output: A train predictor Dataset, a train target DatArray, a test predictor Dataset, and a test target DataArray.
184
185        *predictor_variables*
186        
187        If 'univariate' all variables in *PredictionTest.predictors* are used individually to predict *PredictionTest.predictand*.
188        
189        If 'all' all variables in *PredictionTest.predictors* are used collectively to predict *PredictionTest.predictand*.
190        
191        If a 1D array of variable names in *PredictionTest.predictors*, each specified variable is used individually to predict *PredictionTest.predictand*.
192        
193        If a 2D array of iterables over variable names in *PredictionTest.predictors*, each specified combination of variables is used to predict *PredictionTest.predictand*.
194            
195        *keep_models*
196        If True, a dictionary of arrays of fitted models is returned, corresponding to each variable combination and cross validated model that was computed.
197        
198        *model_kwargs*
199        
200        A dictionary of keyword arguments to pass to *model*
201        *cv_kwargs*
202        
203        A dictionary of keyword arguments to pass to *cv_method*
204
205        *score_kwargs*
206        
207        A dictionary of keyword arguments to pass to *score*
208        
209        **Returns**
210        
211        If keep_models is False:
212        
213        returns *score_da*, a Dataset of skill scores for each prediction experiment, with a cross_validation coordinate.
214        
215        if keep_models is True:
216        
217        return (*score_da*,*model_da*)
218        """
219        if np.ndim(predictor_variables)==1:
220            el_type=[np.ndim(element) for element in predictor_variables]
221            
222            if len(np.unique(el_type))==1 and np.unique(el_type)==1:#a list of ragged lists
223                score_labels=[f'predictor_set_{i+1}' for i in range(len(predictor_variables))]
224                
225            else:#just a list
226                score_labels=predictor_variables
227            
228        elif np.ndim(predictor_variables)==2:
229            score_labels=[f'predictor_set_{i+1}' for i in range(len(predictor_variables))]
230
231        elif predictor_variables =='univariate':
232            predictor_variables = list(self.predictors.data_vars)
233            score_labels=list(self.predictors.data_vars)
234        elif predictor_variables=='all':
235            predictor_variables=[list(self.predictors.data_vars)]
236            score_labels=['all']
237        else:
238            raise(ValueError('predictor variables must be one of ["univariate","all",1-D array, 2-D array]'))
239        
240        if type(score)==str:
241            score=self.predefined_scores[score]
242        
243        predictors=self.predictors
244        target=self.predictand
245        predictors,target=xr.align(predictors.dropna('time',how='all'),target.dropna('time',how='all'))
246        
247        #train_pred,train_target,test_pred,test_target
248        cv_iterator=self._split_data(predictors,target,cv_method,**cv_kwargs)
249        
250        score_da=[]
251        model_da=[]
252        for split_data in cv_iterator:
253            scores={}
254            models={}
255            train_preds,train_target,test_preds,test_target=split_data
256            
257            for label,variable in zip(score_labels,predictor_variables):
258                
259                train_pred=train_preds[variable]
260                test_pred=test_preds[variable]
261                
262                #univariate
263                if type(variable)==str:
264                    
265                    train_pred=np.array(train_pred).reshape([-1,1])
266                    test_pred=np.array(test_pred).reshape([-1,1])
267
268                #multivariate
269                else:
270                    train_pred=np.array(train_pred.to_array('var').transpose(*train_pred.dims,'var'))
271                    test_pred=np.array(test_pred.to_array('var').transpose(*test_pred.dims,'var'))
272
273                fitted_model,det_pred,prob_pred=self._fit_model_and_predict(model,train_pred,train_target,test_pred,**model_kwargs)
274                test_score=score(det_pred,prob_pred,test_target,**score_kwargs)
275                scores[label]=test_score
276                if keep_models:
277                    models[label]=fitted_model
278            score_da.append(xr.Dataset(scores))
279            model_da.append(models)
280        score_da=xr.concat(score_da,'cv')
281        self.computed_scores=score_da
282
283        if keep_models:
284            return score_da,model_da
285        else:
286            return score_da
287    
288    def _fit_model_and_predict(self,model,train_pred,train_target,test_pred,**kwargs):
289        
290        if type(model)==str:
291            model=self.predefined_models[model]
292        
293        train_pred=np.array(train_pred)
294        train_target=np.array(train_target)
295        test_pred=np.array(test_pred)
296        self._verify_model_inputs(train_pred,train_target,test_pred)
297        output=model(train_pred,train_target,test_pred,**kwargs)
298        self._verify_model_output(output,test_pred)
299        return output
300    
301    def _verify_model_inputs(self, train_pred,train_target,test_pred,**model_kwargs):
302    
303        assert train_pred.ndim==2
304        T1,N=train_pred.shape
305
306        assert test_pred.ndim==2
307        T2,M=test_pred.shape
308        assert N==M
309
310        assert train_target.ndim==1
311        T=len(train_target)
312        assert T==T1
313        return 
314
315    def _verify_model_output(self,output,test_pred):
316        
317        T,N=test_pred.shape
318        assert len(output)==3
319        model,det_pred,prob_pred=output
320
321        assert (det_pred is not None)|(prob_pred is not None)
322
323        if det_pred is not None:
324            assert det_pred.shape[0]==T
325            assert np.ndim(det_pred)==1
326        if prob_pred is not None:
327            assert prob_pred.shape[0]==T
328            assert np.ndim(prob_pred)==2
329        return
330    
331    def _split_data(self,predictors,target,cv_method,**cv_kwargs):
332        if cv_method is None:
333            cv_method=CV_None
334        elif type(cv_method)==str:
335            cv_method=self.predefined_cv_methods[cv_method]
336        
337        cv_iterator=cv_method(predictors,target,**cv_kwargs)
338        return cv_iterator
339    
340    def add_score_to_index_metadata(self,indices,label='score',raise_on_missing_var=False,reduce_func=np.nanmean):
341        
342        """ Annotate a Dataset of indices with computed skill scores by adding them to the attributes of each DataArray in the Dataset.
343        
344        **Arguments**
345        
346        *indices*
347        
348        *An xarray.Dataset of indices*.
349        
350        **Optional arguments**
351        
352        *label*
353        
354        A string determining the name of the added attribute key.
355        
356        *raise_on_missing_var*
357        
358        A boolean, determining if an error is raised if not all variables present in the computed skill scores are present in the indices.
359        
360        *reduce_func*
361        
362        The function used to reduce the 'cv' vector of skill scores to a single value. Default is the mean average, ignoring nans. To add the entire vector of scores for different cross validations, pass *reduce_func*=lambda x: x
363        
364        """
365        s=self.computed_scores
366        if s is None:
367            raise(ValueError('No scores computed.'))
368        for var in s.data_vars:
369            
370            score=s[var].values
371            try:
372                indices[var].attrs[label]=reduce_func(score)
373            except:
374                if raise_on_missing_var:
375                    raise(ValueError(f'Key {var} present in self.computed_scores but not in indices.'))
376                pass
377        return
378        
379