Properly toggle safe functions feature from analyze-archive.py

2 files changed, 9 insertions, 4 deletions

diff --git a/bin/analyze-archive.py b/bin/analyze-archive.py
index 7779b4b..df3ecbe 100755
--- a/bin/analyze-archive.py
+++ b/bin/analyze-archive.py
@@ -60,12 +60,14 @@ if __name__ == '__main__':
     ignored_trace_indexes = None
     discard_outliers = None
     tex_output = False
+    safe_functions_enabled = False
     function_override = {}
 
     try:
         optspec = (
             'plot-unparam= plot-param= '
-            'ignored-trace-indexes= discard-outliers= function-override= tex-output'
+            'ignored-trace-indexes= discard-outliers= function-override= tex-output '
+            'with-safe-functions'
         )
         raw_opts, args = getopt.getopt(sys.argv[1:], "", optspec.split(' '))
 
@@ -89,6 +91,9 @@ if __name__ == '__main__':
             if 'tex-output' in opts:
                 tex_output = True
 
+            if 'with-safe-functions' in opts:
+                safe_functions_enabled = True
+
     except getopt.GetoptError as err:
         print(err)
         sys.exit(2)
@@ -135,7 +140,7 @@ if __name__ == '__main__':
     lut_quality = model.assess(lut_model)
 
     print('--- param model ---')
-    param_model, param_info = model.get_fitted()
+    param_model, param_info = model.get_fitted(safe_functions_enabled = safe_functions_enabled)
     if not tex_output:
         for state in model.states():
             for attribute in ['power']:
diff --git a/bin/test.py b/bin/test.py
index 60e8648..433b423 100755
--- a/bin/test.py
+++ b/bin/test.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-from dfatool import EnergyModel, RawData
+from dfatool import EnergyModel, RawData, analytic
 import unittest
 
 class TestStaticModel(unittest.TestCase):
@@ -229,5 +229,5 @@ class TestStaticModel(unittest.TestCase):
         self.assertAlmostEqual(param_info('RX', 'power')['function']._regression_args[1], 206, places=0)
 
 if __name__ == '__main__':
-    dfatool.safe_function_enabled = False
+    analytic.safe_function_enabled = False
     unittest.main()