modify hardcoded filepaths in the quadtree to allow transplanting it to new filesystems.

Author: davemfish

src/natcap/invest/recreation/buffered_numpy_disk_map.py

@@ -109,7 +109,8 @@ def _write(self, array_id_list):
                     where array_id=? LIMIT 1""", [array_id])
             array_path = db_cursor.fetchone()
             if array_path is not None:
-                _npy_append(array_path[0], numpy.concatenate(array_deque))
+                _npy_append(os.path.join(
+                    self.manager_directory, array_path[0]), numpy.concatenate(array_deque))
                 array_deque = None
                 # cache gets wiped at end so okay to use same deque
                 # array_deque.append(numpy.load(array_path[0]))
@@ -122,16 +123,18 @@ def _write(self, array_id_list):
                 # off the last two characters in the filename
                 array_filename = uuid.uuid4().hex + '.npy'
                 # -6:-4 skips the extension and gets the last 2 characters
+                array_subdirectory = array_filename[-6:-4]
                 array_directory = os.path.join(
-                    self.manager_directory, array_filename[-6:-4])
+                    self.manager_directory, array_subdirectory)
                 if not os.path.isdir(array_directory):
                     os.mkdir(array_directory)
                 array_path = os.path.join(array_directory, array_filename)
                 # save the file
                 array_data = numpy.concatenate(array_deque)
                 array_deque = None
                 numpy.save(array_path, array_data)
-                insert_list.append((array_id, array_path))
+                # insert_list.append((array_id, array_path))
+                insert_list.append((array_id, os.path.join(array_subdirectory, array_filename)))
         db_connection.close()
         return insert_list
 
@@ -193,7 +196,7 @@ def read(self, array_id):
         db_connection.close()
 
         if array_path is not None:
-            array_data = numpy.load(array_path[0])
+            array_data = numpy.load(os.path.join(self.manager_directory, array_path[0]))
         else:
             array_data = numpy.empty(
                 0, dtype=BufferedNumpyDiskMap._ARRAY_TUPLE_TYPE)
@@ -215,10 +218,11 @@ def delete(self, array_id):
             [array_id])
         array_path = db_cursor.fetchone()
         if array_path is not None:
-            os.remove(array_path[0])
+            array_abs_path = os.path.join(self.manager_directory, array_path[0])
+            os.remove(array_abs_path)
             try:
                 # attempt to remove the directory if it's empty
-                os.rmdir(os.path.dirname(array_path[0]))
+                os.rmdir(os.path.dirname(array_abs_path))
             except OSError:
                 # it's not empty, not a big deal
                 pass

src/natcap/invest/recreation/recmodel_server.py

@@ -130,6 +130,7 @@ def __init__(
             # ooc_qt_picklefilename = os.path.join(
             #     global_cache, quadtree_pickle_filename)
             ooc_qt_picklefilename = quadtree_pickle_filename
+            global_cache = os.path.dirname(ooc_qt_picklefilename)
         else:
             raise ValueError(
                 'Both raw_csv_filename and quadtree_pickle_filename'
@@ -138,6 +139,7 @@ def __init__(
         if os.path.isfile(ooc_qt_picklefilename):
             LOGGER.info(
                 '%s quadtree already exists', ooc_qt_picklefilename)
+            ooc_qt_picklefilename = transplant_quadtree(ooc_qt_picklefilename, cache_workspace)
         else:
             LOGGER.info(
                 '%s not found, constructing quadtree', ooc_qt_picklefilename)
@@ -149,6 +151,7 @@ def __init__(
                 max_points_per_node, max_depth)
         self.qt_pickle_filename = ooc_qt_picklefilename
         self.local_cache_workspace = os.path.join(cache_workspace, 'local')
+        # self.global_cache_dir = global_cache
         self.min_year = min_year
         self.max_year = max_year
 
@@ -283,8 +286,10 @@ def _calc_aggregated_points_in_aoi(
         pud_poly_feature_queue = multiprocessing.Queue(4)
         n_polytest_processes = multiprocessing.cpu_count()
 
+        LOGGER.info(f'OPENING {self.qt_pickle_filename}')
         with open(self.qt_pickle_filename, 'rb') as qt_pickle:
             global_qt = pickle.load(qt_pickle)
+
         aoi_layer = aoi_vector.GetLayer()
         aoi_extent = aoi_layer.GetExtent()
         aoi_ref = aoi_layer.GetSpatialRef()
@@ -1010,3 +1015,32 @@ def _hash_blocks(file_buffer_queue):
     if fast_hash:
         file_hash += '_fast_hash'
     return file_hash
+
+def transplant_quadtree(qt_pickle_filepath, workspace):
+    storage_dir = os.path.dirname(qt_pickle_filepath)
+    pickle_filepath = qt_pickle_filepath
+
+    def rename_managers(qt):
+        if qt.is_leaf:
+            qt.node_data_manager.manager_filename = f'{qt_pickle_filepath}.db'
+            qt.node_data_manager.manager_directory = os.path.dirname(qt_pickle_filepath)
+            qt.quad_tree_storage_dir = storage_dir
+        else:
+            [rename_managers(qt.nodes[index]) for index in range(4)]
+        return qt
+
+    with open(qt_pickle_filepath, 'rb') as qt_pickle:
+        global_qt = pickle.load(qt_pickle)
+
+    if global_qt.quad_tree_storage_dir != storage_dir:
+        LOGGER.info(
+            f'setting quadtree node references to the local filesystem '
+            f'{storage_dir}')
+        new_qt = rename_managers(global_qt)
+        pickle_filepath = os.path.join(
+            workspace, f'transplant_{os.path.basename(qt_pickle_filepath)}')
+        LOGGER.info(
+            f'writing new quadtree index to {pickle_filepath}')
+        with open(pickle_filepath, 'wb') as qt_pickle:
+            pickle.dump(new_qt, qt_pickle)
+    return pickle_filepath