From 06780efbaf0ddf458e7cb0cf84061db497a80c8e Mon Sep 17 00:00:00 2001
From: Giacomo Nannicini <giacomo.n@gmail.com>
Date: Sun, 24 May 2020 18:37:02 -0400
Subject: [PATCH] Better module import for command-line interface

---
 bin/rbfopt_cl_interface.py          | 38 +++++++++++++++++------------
 src/rbfopt/rbfopt_test_functions.py |  6 ++---
 2 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/bin/rbfopt_cl_interface.py b/bin/rbfopt_cl_interface.py
index 0ef97a8..9623eb8 100644
--- a/bin/rbfopt_cl_interface.py
+++ b/bin/rbfopt_cl_interface.py
@@ -16,6 +16,7 @@
 
 import sys
 import os
+import importlib
 # We must set the threading options before numpy is loaded, otherwise
 # there might be issues when running several processes in parallel.
 os.environ['OMP_NUM_THREADS'] = '1'
@@ -79,7 +80,9 @@ def register_options(parser):
                         'function and the description of its ' +
                         'characteristics. This file should implement a ' +
                         'BlackBox class derived from ' +
-                        'rbfopt_black_box.BlackBox. ')
+                        'rbfopt_black_box.BlackBox, with name ' +
+                        'RbfoptBlackBox. Note: the directory containing ' +
+                        'it will be added to the Python path.')
     intset.add_argument('--load', '-l', action='store', dest='load_state',
                         help='File to read state to resume optimization')
     intset.add_argument('--log', '-o', action='store',
@@ -173,10 +176,9 @@ def rbfopt_cl_interface(args, black_box):
         except Exception as e:
             print('Exception raised reading file with initialization points',
                   file=output_stream)
-            print(type(e), file=output_stream)
             print(e, file=output_stream)
             output_stream.close()
-            exit()
+            raise
         alg = RbfoptAlgorithm(settings=settings, black_box=black_box,
                               init_node_pos=init_node_pos,
                               init_node_val=init_node_val)
@@ -212,17 +214,23 @@ def rbfopt_cl_interface(args, black_box):
         raise ValueError('The file {:s} '.format(args.black_box_file) +
                          'supposed to provide the implementation of ' +
                          'RbfoptBlackBox does not exist.')
-    if (sys.version_info[0] == 2):
-        import imp
-        bb = imp.load_source('user_bb', args.black_box_file)
-    elif (sys.version_info[0] == 3 and 3 <= sys.version_info[1] <= 4):
-        from importlib.machinery import SourceFileLoader
-        bb = SourceFileLoader('user_bb', args.black_box_file).load_module()
-    elif (sys.version_info[0] == 3 and sys.version_info[1] >= 5):
-        import importlib.util
-        spec = importlib.util.spec_from_file_location('user_bb',
-                                                      args.black_box_file)
-        bb = importlib.util.module_from_spec(spec)
-        spec.loader.exec_module(bb)
+
+    try:
+        rel_dir = os.path.dirname(args.black_box_file)
+        if (rel_dir == ''):
+            rel_dir = '.'
+        abs_dir = os.path.abspath(rel_dir)
+        # Add directory to path
+        sys.path.append(abs_dir)
+        # Import module (after removing trailing .py, if any)
+        module_name = os.path.basename(args.black_box_file)
+        if (module_name.endswith('.py')):
+            module_name = module_name[:-3]
+        bb = importlib.import_module(module_name)
+    except Exception as e:
+        print('Error while opening module with user black box',
+              file=sys.stderr)
+        print(e, file=sys.stderr)
+        raise
     # Run the interface
     rbfopt_cl_interface(vars(args), bb.RbfoptBlackBox())
diff --git a/src/rbfopt/rbfopt_test_functions.py b/src/rbfopt/rbfopt_test_functions.py
index b216b33..dbd8422 100644
--- a/src/rbfopt/rbfopt_test_functions.py
+++ b/src/rbfopt/rbfopt_test_functions.py
@@ -1051,9 +1051,9 @@ def evaluate(cls, x):
         penalty += 10*max(0, (0.0056858*x[1]*x[4] + 0.0006262*x[0]*x[3] - 
                               0.0022053*x[2]*x[4] + 85.334407) - 92)
         penalty += 10*max(0, -(0.0071317*x[1]*x[4] + 0.0029955*x[0]*x[1] +
-                               0.0021813*math.sqrt(x[2]) + 80.51249) + 90)
+                               0.0021813*np.sqrt(x[2]) + 80.51249) + 90)
         penalty += 10*max(0, (0.0071317*x[1]*x[4] + 0.0029955*x[0]*x[1] +
-                              0.0021813*math.sqrt(x[2]) + 80.51249) - 110)
+                              0.0021813*np.sqrt(x[2]) + 80.51249) - 110)
         penalty += 10*max(0, -(0.0047026*x[2]*x[4] + 0.0012547*x[0]*x[2] +
                                0.0019085*x[2]*x[3] + 9.300961) + 20) 
         penalty += 10*max(0, (0.0047026*x[2]*x[4] + 0.0012547*x[0]*x[2] +
@@ -2618,7 +2618,7 @@ def evaluate(cls, x):
         for i in range(len(x)-2):
             si = 2**i*np.sqrt((x[i]-cls.optimum_point[i])**2 +
                               (x[i+1]-cls.optimum_point[i+1])**2)
-            value += (normalizer * math.sqrt(si) * 
+            value += (normalizer * np.sqrt(si) * 
                       (fun(50*si**0.20) + 1))**2
         return value - 10