From e01941f06b98914c6a4909f0801fbc3788ab14ce Mon Sep 17 00:00:00 2001
From: tyrin <nytyrin@stud.kpfu.ru>
Date: Tue, 27 May 2025 10:40:36 +0300
Subject: [PATCH] added mask_radius_deg field to nearby source dict

---
 management/commands/set_contaminated.py |  2 +-
 views.py                                | 89 ++++++++++++++++++-------
 2 files changed, 65 insertions(+), 26 deletions(-)

diff --git a/management/commands/set_contaminated.py b/management/commands/set_contaminated.py
index 539b12f..ee5a309 100644
--- a/management/commands/set_contaminated.py
+++ b/management/commands/set_contaminated.py
@@ -156,7 +156,7 @@ class Command(BaseCommand):
             self.stdout.write("All catalog rows already exist in the database.")
 
         # hard coded nside and flux-radius mapping
-        # maybe change that
+        # maybe change
 
         nside = 4096
         npix = hp.nside2npix(nside)
diff --git a/views.py b/views.py
index 6cf1b62..763ba1d 100644
--- a/views.py
+++ b/views.py
@@ -12,7 +12,18 @@ from astropy.stats import poisson_conf_interval
 
 from collections import defaultdict
 
-from django.db.models import Sum, Max
+from django.db.models import (
+    Max,
+    Sum,
+    F,
+    IntegerField,
+    BooleanField,
+    ExpressionWrapper,
+    Case,
+    When,
+    Value,
+)
+from django.db.models.functions import Cast
 from django.shortcuts import get_object_or_404
 from rest_framework.views import APIView
 from rest_framework.response import Response
@@ -366,25 +377,37 @@ class UpperLimitView(APIView):
 
         nearby_sources = []
 
+        radius_map = {0: 0.06, 125: 0.15, 250: 0.5, 2000: 0.9, 20000: 2.5}
+
+        sorted_bounds = sorted(radius_map.keys())
+
         # refine belt to circular region using astropy separation
         for catsrc in belt_sources:
             catsrc_coord = SkyCoord(catsrc.ra_deg, catsrc.dec_deg, unit="deg")
-            if center_coord.separation(catsrc_coord).deg <= radius_deg:
-                nearby_sources.append(
-                    {
-                        "srcid": catsrc.srcid,
-                        "name": catsrc.name,
-                        "ra_deg": catsrc.ra_deg,
-                        "dec_deg": catsrc.dec_deg,
-                        "pos_error": catsrc.pos_error,
-                        "significance": catsrc.significance,
-                        "flux": catsrc.flux,
-                        "flux_error": catsrc.flux_error,
-                        "catalog_name": catsrc.catalog_name,
-                        "new_xray": catsrc.new_xray,
-                        "source_type": catsrc.source_type,
-                    }
-                )
+            if center_coord.separation(catsrc_coord).deg > radius_deg:
+                continue
+            f = catsrc.flux or 0.0
+            for lb in reversed(sorted_bounds):
+                if f >= lb:
+                    mask_radius = radius_map[lb]
+                    break
+
+            nearby_sources.append(
+                {
+                    "srcid": catsrc.srcid,
+                    "name": catsrc.name,
+                    "ra_deg": catsrc.ra_deg,
+                    "dec_deg": catsrc.dec_deg,
+                    "pos_error": catsrc.pos_error,
+                    "significance": catsrc.significance,
+                    "flux": catsrc.flux,
+                    "flux_error": catsrc.flux_error,
+                    "catalog_name": catsrc.catalog_name,
+                    "new_xray": catsrc.new_xray,
+                    "source_type": catsrc.source_type,
+                    "mask_radius_deg": mask_radius,
+                }
+            )
 
         # REGION IMAGE SERVING
         # ****************************************************************
@@ -407,11 +430,26 @@ class UpperLimitView(APIView):
         # fetch those pixels for the requested surveys
         # summing counts and sorting by hpid
 
+        # map_pixels_qs = (
+        #     Pixel.objects.filter(hpid__in=map_pixel_list, survey__in=survey_numbers)
+        #     .values("hpid")
+        #     .annotate(counts=Sum("counts"))
+        #     .order_by("hpid")
+        # )
+
         map_pixels_qs = (
             Pixel.objects.filter(hpid__in=map_pixel_list, survey__in=survey_numbers)
             .values("hpid")
             .annotate(
-                counts=Sum("counts"),
+                total_counts=Sum("counts"),
+                max_contaminated_int=Max(Cast("contaminated", IntegerField())),
+            )
+            .annotate(
+                contaminated=Case(
+                    When(max_contaminated_int=1, then=Value(True)),
+                    default=Value(False),
+                    output_field=BooleanField(),
+                )
             )
             .order_by("hpid")
         )
@@ -421,15 +459,16 @@ class UpperLimitView(APIView):
 
         # 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_counts_list = [d["total_counts"] for d in map_pixels_list]
+        map_contaminated_list = [d["contaminated"] for d in map_pixels_list]
 
-        cont_dict = dict(
-            Pixel.objects.filter(hpid__in=map_healpix_list, survey__in=survey_numbers)
-            .values_list("hpid", "contaminated")
-            .distinct()
-        )
+        # cont_dict = dict(
+        #     Pixel.objects.filter(hpid__in=map_healpix_list, survey__in=survey_numbers)
+        #     .values_list("hpid", "contaminated")
+        #     .distinct()
+        # )
 
-        map_contaminated_list = [cont_dict[h] for h in map_healpix_list]
+        # map_contaminated_list = [cont_dict[h] for h in map_healpix_list]
 
         # set map nside
         map_nside = 4096