From 560ccecfd9865c3faa61be2a29418a6c311d068a Mon Sep 17 00:00:00 2001
From: Anna Price <36504362+annacprice@users.noreply.github.com>
Date: Mon, 13 Dec 2021 19:18:38 +0000
Subject: [PATCH] update warning handling for build_model

---
 covate/build_model.py | 40 +++++++++++++++++++---------------------
 1 file changed, 19 insertions(+), 21 deletions(-)

diff --git a/covate/build_model.py b/covate/build_model.py
index 48aaa85..c81205d 100644
--- a/covate/build_model.py
+++ b/covate/build_model.py
@@ -9,6 +9,8 @@
 import matplotlib.pyplot as plt
 from .utils import appendline, pairwise, getenddate
 import warnings
+warnings.simplefilter("ignore", UserWarning)
+warnings.simplefilter("ignore", RuntimeWarning)
 
 
 def buildmodel(timeseries, lineagelist, regionlist, enddate, output, maxlag,
@@ -60,22 +62,16 @@ def buildmodel(timeseries, lineagelist, regionlist, enddate, output, maxlag,
         # get basic information on timeseries
         checkdistribution(X_train, lineage, filename, errorlog)
 
-        # plot the autocorrelation, ignore warnings
-        with warnings.catch_warnings():
-
-            warnings.simplefilter("ignore", UserWarning)
-
-            warnings.simplefilter("ignore", RuntimeWarning)
-
-            if validate:
+        # plot the autocorrelation
+        if validate:
 
-                plotautocorr(X_train, lineage, maxlag, output, enddate,
-                             'additional-plots/validation')
+            plotautocorr(X_train, lineage, maxlag, output, enddate,
+                         'additional-plots/validation')
 
-            else:
+        else:
 
-                plotautocorr(X_train, lineage, maxlag, output, enddate,
-                             'additional-plots/prediction')
+            plotautocorr(X_train, lineage, maxlag, output, enddate,
+                         'additional-plots/prediction')
 
         # check for granger causality
         try:
@@ -168,15 +164,16 @@ def buildmodel(timeseries, lineagelist, regionlist, enddate, output, maxlag,
 
                 if not validate:
 
-                    vecerrcorr(X_train, lineage, VECMdeterm, lag, coint_count,
-                               regionlist, nsteps, alpha, filename, output,
-                               errorlog, enddate)
+                    vecerrcorr(X_train, lineage, VECMdeterm, lag,
+                               coint_count, regionlist, nsteps, alpha,
+                               filename, output, errorlog, enddate)
 
                 else:
 
-                    vecerrcorrvalid(X_train, X_test, lineage, VECMdeterm, lag,
-                                    coint_count, regionlist, nsteps, alpha,
-                                    filename, output, errorlog, enddate)
+                    vecerrcorrvalid(X_train, X_test, lineage, VECMdeterm,
+                                    lag, coint_count, regionlist, nsteps,
+                                    alpha, filename, output, errorlog,
+                                    enddate)
 
             # else check for stationarity and difference then run VAR
             else:
@@ -186,8 +183,9 @@ def buildmodel(timeseries, lineagelist, regionlist, enddate, output, maxlag,
 
                 if not validate:
 
-                    vecautoreg(X_train, lineage, maxlag, regionlist, nsteps,
-                               alpha, filename, output, errorlog, enddate)
+                    vecautoreg(X_train, lineage, maxlag, regionlist,
+                               nsteps, alpha, filename, output, errorlog,
+                               enddate)
 
                 else: