Batch search inside stack search #530

vlnistor · 2024-03-23T17:37:59Z

This tackles #356

The idea is that you'd replace the full in-memory search:

min_observations = ... # Some min observations
candidates = ... # Some trajectories
stack = ImageStack(im_list)
search = StackSearch(stack)
search.search(candidates, min_observations)

With a batched version:

 batch_results = []
 with BatchSearchManager(StackSearch(stack), candidates, min_observations) as batch_search:
    batch_results = []
    for i in range(0, width, 5):
        batch_search.set_start_bounds_x(i, i + 5)
        for j in range(0, height, 5):
            batch_search.set_start_bounds_y(j, j + 5)
            batch_results.extend(batch_search.search_batch())

The steps are:

Use the BatchSearchManager context manager to obtain a managed batch_search object
Loop through the whole grid by breaking it into smaller subsets

… setup/clear resources

removed accidental git conflict merge conflicts dirac-institute#2

…so as to be able to init gpu_search_list outside of the main search func

vlnistor · 2024-03-23T21:14:40Z

src/kbmod/batch_search.py

@@ -0,0 +1,40 @@
+class BatchSearchManager:


The main idea here was that I didn't want the user to have to call prepare_batch_search and finish_search manually

Not for this PR, but something to consider as a future extension: In theory this could be extended further by making this a generator:

for batch_search in BatchSearchManager(stack_search, search_list, min_observations, x_size, y_size): batch_results.extend(batch_search.search_batch())

where the BatchSearchManager tracks the block sizes to search, current block locations, etc.

src/kbmod/search/stack_search.cpp

jeremykubica

A bunch of points for discussion.

src/kbmod/batch_search.py

src/kbmod/search/trajectory_list.cpp

src/kbmod/search/stack_search.cpp

src/kbmod/search/stack_search.h

src/kbmod/search/trajectory_list.h

tests/test_search.py

jeremykubica · 2024-04-01T14:39:03Z

src/kbmod/batch_search.py

@@ -0,0 +1,40 @@
+class BatchSearchManager:


Not for this PR, but something to consider as a future extension: In theory this could be extended further by making this a generator:

for batch_search in BatchSearchManager(stack_search, search_list, min_observations, x_size, y_size): batch_results.extend(batch_search.search_batch())

where the BatchSearchManager tracks the block sizes to search, current block locations, etc.

…e_batch + removed move assignment operator in favour of set_trajectories + updated search to search_all across the python code

review-notebook-app · 2024-04-03T22:17:04Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

jeremykubica

This is looking really good. A few stylistic cleanups.

jeremykubica · 2024-04-04T12:06:48Z

src/kbmod/search/stack_search.cpp

+}
+
+
+int StackSearch::extract_max_results(){


Good call breaking this out into a function.

jeremykubica · 2024-04-04T12:08:14Z

tests/test_search.py

+        stack = ImageStack(imlist)
+        im_list = stack.get_images()
+        # Create a new list of LayeredImages with the added object.
+        new_im_list = []


You can remove new_im_list

Removed now

jeremykubica · 2024-04-04T12:10:35Z

src/kbmod/search/stack_search.h

@@ -29,6 +29,7 @@ using Point = indexing::Point;
 using Image = search::Image;

 class StackSearch {
+    int extract_max_results();


Maybe rename this to get_max_results or compute_max_results to avoid any confusion about this somehow extracting the results from the cached list.

Good point - renamed

src/kbmod/search/stack_search.cpp

jeremykubica · 2024-04-04T12:13:18Z

src/kbmod/search/stack_search.cpp

-
-    // Do the actual search on the GPU.
-    DebugTimer search_timer = DebugTimer("search execution", rs_logger);
+        // Do the actual search on the GPU.


Indentation of this comment appears off.

src/kbmod/search/stack_search.h

vlnistor added 5 commits March 23, 2024 17:02

batch search inside stack search

e75f93f

moved the batch search inside a context manager to remove the need to…

666d2e5

… setup/clear resources

fixed to include 2nd argument in prepare_batch

88d3289

conclude merge

12894f8

resolving search_batch merge conflicts with main

0b25c32

removed accidental git conflict merge conflicts dirac-institute#2

vlnistor force-pushed the I-356/batch-search branch from 9e322bb to 0b25c32 Compare March 23, 2024 20:48

added move assignment operators for both TrajectoryList and GPUArray …

cbcaa9c

…so as to be able to init gpu_search_list outside of the main search func

vlnistor commented Mar 23, 2024

View reviewed changes

src/kbmod/search/stack_search.cpp Outdated Show resolved Hide resolved

vlnistor requested a review from jeremykubica March 25, 2024 21:03

cleaned up comments

cc5d52d

jeremykubica reviewed Apr 1, 2024

View reviewed changes

search_batch functionality extracted from search_all and search_singl…

2106801

…e_batch + removed move assignment operator in favour of set_trajectories + updated search to search_all across the python code

vlnistor requested a review from jeremykubica April 3, 2024 22:19

jeremykubica reviewed Apr 4, 2024

View reviewed changes

vlnistor added 2 commits April 7, 2024 22:06

renamed helper to compute_max_results and minor style changes

9b5226f

kbmod_reference notebook only change search.search -> search.search_all

6180283

jeremykubica approved these changes Apr 8, 2024

View reviewed changes

jeremykubica merged commit 7412363 into dirac-institute:main Apr 8, 2024
2 checks passed

jeremykubica mentioned this pull request May 6, 2024

Refactor StackSearch to do batching at search() function #356

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Batch search inside stack search #530

Batch search inside stack search #530

vlnistor commented Mar 23, 2024 •

edited

Loading

vlnistor Mar 23, 2024 •

edited

Loading

jeremykubica Apr 1, 2024

jeremykubica left a comment

jeremykubica Apr 1, 2024

review-notebook-app bot commented Apr 3, 2024

jeremykubica left a comment

jeremykubica Apr 4, 2024

jeremykubica Apr 4, 2024

vlnistor Apr 7, 2024

jeremykubica Apr 4, 2024

vlnistor Apr 7, 2024

jeremykubica Apr 4, 2024

vlnistor Apr 7, 2024

Batch search inside stack search #530

Batch search inside stack search #530

Conversation

vlnistor commented Mar 23, 2024 • edited Loading

vlnistor Mar 23, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jeremykubica left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

review-notebook-app bot commented Apr 3, 2024

jeremykubica left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vlnistor commented Mar 23, 2024 •

edited

Loading

vlnistor Mar 23, 2024 •

edited

Loading