diff --git a/views.py b/views.py index d8c60ad..278f92b 100644 --- a/views.py +++ b/views.py @@ -908,12 +908,14 @@ class StackedUpperLimitView(APIView): # iterate over source vectors, making maps - # Create a single list to hold all Healpix IDs - all_map_pixels = [] - source_indices = [] + # # Create a single list to hold all Healpix IDs + # all_map_pixels = [] + # source_indices = [] + + hpid_to_source_indices_map = defaultdict(list) # Iterate over source vectors and get all unique Healpix IDs - for index, src_vec in enumerate(src_vec_list): + for source_index, src_vec in enumerate(src_vec_list): pixels = hp.query_disc( nside=4096, vec=src_vec, @@ -921,33 +923,54 @@ class StackedUpperLimitView(APIView): nest=False, radius=(map_radius * u.arcsecond).to(u.radian).value, ) - all_map_pixels.extend(pixels) - source_indices.extend([index] * len(pixels)) + for p in pixels: + hpid_to_source_indices_map[p].append(source_index) # Get only unique Healpix IDs - unique_map_pixels = list(set(all_map_pixels)) - + # unique_map_pixels = list(set(all_map_pixels)) + unique_map_pixels = list(hpid_to_source_indices_map.keys()) # Perform a single database query for all unique pixels + map_pixels_qs = ( Pixel.objects.filter(hpid__in=unique_map_pixels, survey__in=survey_numbers) .values("hpid") .annotate(counts=Sum("counts"), exposure=Sum("exposure")) - .order_by("hpid") ) - # turn the queryset to a list - map_pixels_list = list(map_pixels_qs) + db_results_map = {p["hpid"]: p for p in map_pixels_qs} - # get lists of healpix indices and count values - map_healpix_list = [d["hpid"] for d in map_pixels_list] - map_counts_list = [d["counts"] for d in map_pixels_list] - map_exposure_list = [d["exposure"] for d in map_pixels_list] - # map_contaminated_list = [d["contaminated"] for d in map_pixels_list] + map_healpix_list = [] + map_counts_list = [] + map_exposure_list = [] + map_source_indices_list = [] - # set map nside + # # turn the queryset to a list + # map_pixels_list = list(map_pixels_qs) + + # # get lists of healpix indices and count values + # map_healpix_list = [d["hpid"] for d in map_pixels_list] + # map_counts_list = [d["counts"] for d in map_pixels_list] + # map_exposure_list = [d["exposure"] for d in map_pixels_list] + # # map_contaminated_list = [d["contaminated"] for d in map_pixels_list] + + for hpid, source_indices in hpid_to_source_indices_map.items(): + # Check if this pixel was actually found in the database query + if hpid in db_results_map: + db_data = db_results_map[hpid] + counts = db_data["counts"] + exposure = db_data["exposure"] + + # For each source this pixel is associated with, create a parallel entry + for source_index in source_indices: + map_healpix_list.append(hpid) + map_counts_list.append(counts) + map_exposure_list.append(exposure) + map_source_indices_list.append(source_index) + + # set map nside explicitly map_nside = 4096 - # set map order + # set map order explicitly map_order = "ring" # assemble the result dict @@ -955,7 +978,7 @@ class StackedUpperLimitView(APIView): "healpix": map_healpix_list, "counts": map_counts_list, "exposure": map_exposure_list, - "source_index": source_indices, + "source_index": map_source_indices_list, # "contaminated": map_contaminated_list, "nside": map_nside, "order": map_order,