Skip to content

Commit

Permalink
Merge pull request #151 from mrosseel/nearby-cleanup
Browse files Browse the repository at this point in the history 
Nearby cleanup
  • Loading branch information
@brickbots
brickbots authored Mar 6, 2024
2 parents a8b206b + 08bf7ad commit 8537a0f
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 59 deletions.
33 changes: 18 additions & 15 deletions python/PiFinder/catalog_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,40 +10,43 @@ class ClosestObjectsFinder:
def __init__(self):
pass

# remove from CatalogTracker into utility clas
def get_closest_objects(
self, ra, dec, n, catalogs: Catalogs
def calculate_objects_balltree(
self, ra, dec, catalogs: Catalogs
) -> Tuple[List[CompositeObject], Tuple[List[CompositeObject], BallTree]]:
"""
Takes the current catalog or a list of catalogs, gets the filtered
objects and returns the n closest objects to ra/dec
Calculates a flat list of objects and the balltree for those objects
"""
catalog_list: List[Catalog] = catalogs.get_catalogs()
catalog_list_flat = [
obj for catalog in catalog_list for obj in catalog.filtered_objects
]
if len(catalog_list_flat) < n:
n = len(catalog_list_flat)
object_radecs = [
[np.deg2rad(x.ra), np.deg2rad(x.dec)] for x in catalog_list_flat
]
objects_bt = BallTree(object_radecs, leaf_size=20, metric="haversine")
query = [[np.deg2rad(ra), np.deg2rad(dec)]]
_dist, obj_ind = objects_bt.query(query, k=n)
results = [catalog_list_flat[x] for x in obj_ind[0]]
deduplicated = self._deduplicate(results)
return deduplicated, (catalog_list_flat, objects_bt)
return (catalog_list_flat, objects_bt)

def get_closest_objects_cached(
def get_closest_objects(
self, ra, dec, n, cache: Tuple[List[CompositeObject], BallTree]
) -> List[CompositeObject]:
"""
Takes the current catalog or a list of catalogs, gets the filtered
objects and returns the n closest objects to ra/dec
"""
catalog_list_flat, balltree = cache
query = [[np.deg2rad(ra), np.deg2rad(dec)]]
_dist, obj_ind = cache[1].query(query, k=min(n, len(cache[0])))
results = [cache[0][x] for x in obj_ind[0]]
_, obj_ind = balltree.query(query, k=min(n, len(catalog_list_flat)))
results = [catalog_list_flat[x] for x in obj_ind[0]]
deduplicated = self._deduplicate(results)
return deduplicated

def _deduplicate(self, unfiltered_results):
"""
Make sure no duplicates are in the balltree results.
objects with the same object_id are considered duplicates.
If there are duplicates, the one with the higher precedence catalog_code
is kept.
"""
deduplicated_results = []
seen_ids = set()

Expand Down
75 changes: 31 additions & 44 deletions python/PiFinder/ui/nearby.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
"""
from enum import Enum
import numpy as np
from typing import List
from typing import List, Tuple, Optional

from PiFinder import config
from PiFinder.obj_types import OBJ_TYPE_MARKERS
Expand All @@ -23,26 +23,26 @@
)
from PiFinder.calc_utils import aim_degrees
from PiFinder.catalog_utils import ClosestObjectsFinder
import functools
import logging

from PiFinder import utils
from PiFinder.catalogs import CompositeObject
from PiFinder.ui.catalog import UICatalog
from PIL import Image, ImageChops
import functools
import logging
from pathlib import Path
import os
from itertools import cycle
from PiFinder import utils
from sklearn.neighbors import BallTree


class Modes(Enum):
"""
Enum for the different modes
"""

LOCATE = 0
NAME = 1
INFO = 2
LOCATE = 0 # shows distance to the target
NAME = 1 # shows common names of the target
INFO = 2 # shows magnitude, size, seen, ...


class UINearby(UIModule):
Expand Down Expand Up @@ -97,7 +97,7 @@ def __init__(self, ui_catalog: UICatalog, *args):
self.closest_objects = []
self.closest_objects_text = []
self.font_large = fonts.large
self.objects_balltree = None
self.objects_balltree: Optional[Tuple[List[CompositeObject], BallTree]]
self.closest_objects_finder = ClosestObjectsFinder()
self.current_line = -1
self.mode_cycle = cycle(Modes)
Expand Down Expand Up @@ -125,6 +125,7 @@ def __init__(self, ui_catalog: UICatalog, *args):
def update_config(self):
self.ui_catalog.update_config()
self.objects_balltree = None
return True

def update_object_info(self):
self.update()
Expand Down Expand Up @@ -193,30 +194,26 @@ def update_closest(self):
to that location
"""
if self.shared_state.solution():
closest_objects: List[CompositeObject] = None
ra, dec = (
self.shared_state.solution()["RA"],
self.shared_state.solution()["Dec"],
)
if not self.objects_balltree:
(
closest_objects,
self.objects_balltree,
) = self.closest_objects_finder.get_closest_objects(
self.shared_state.solution()["RA"],
self.shared_state.solution()["Dec"],
self.max_objects + 1,
catalogs=self.catalog_tracker.catalogs,
)
else:
closest_objects = (
self.closest_objects_finder.get_closest_objects_cached(
self.shared_state.solution()["RA"],
self.shared_state.solution()["Dec"],
self.max_objects + 1,
self.objects_balltree,
self.objects_balltree = (
self.closest_objects_finder.calculate_objects_balltree(
ra, dec, catalogs=self.catalog_tracker.catalogs
)
)
closest_objects = self.closest_objects_finder.get_closest_objects(
ra,
dec,
self.max_objects + 1,
self.objects_balltree,
)
self.current_nr_objects = len(closest_objects)
self.closest_objects = closest_objects

def create_locate_text(self) -> List[TextLayouterSimple]:
def create_locate_text(self) -> List[Tuple[str, TextLayouterSimple]]:
result = []
for obj in self.closest_objects:
az, alt = aim_degrees(
Expand All @@ -240,7 +237,7 @@ def create_locate_text(self) -> List[TextLayouterSimple]:
result.append((obj.obj_type, entry))
return result

def create_name_text(self) -> List[TextLayouterSimple]:
def create_name_text(self) -> List[Tuple[str, TextLayouterSimple]]:
result = []
for obj in self.closest_objects:
full_name = f"{','.join(obj.names)}" if obj.names else ""
Expand All @@ -258,17 +255,16 @@ def create_name_text(self) -> List[TextLayouterSimple]:
result.append((obj.obj_type, entry))
return result

def create_info_text(self) -> List[TextLayouterSimple]:
def create_info_text(self) -> List[Tuple[str, TextLayouterSimple]]:
result = []
for obj in self.closest_objects:
obj_mag, obj_color = self._obj_to_mag_color(obj)
mag = f"m{obj_mag}" if obj_mag != 99 else ""
size = f"{self.ruler}{obj.size.strip()}" if obj.size.strip() else ""
full_name = (
f"{mag} {size} {self.checkmark if obj.logged else self.checkmark_no}"
)
check = f" {self.checkmark}" if obj.logged else ""
full_name = f"{mag} {size}{check}"
if len(full_name) > 12:
full_name = mag
full_name = f"{mag}{check}"
obj_name = f"{obj.catalog_code}{obj.sequence}"
_, obj_dist = self.space_calculator.calculate_spaces(
obj_name, full_name, empty_if_exceeds=False, trunc_left=False
Expand All @@ -284,7 +280,7 @@ def create_info_text(self) -> List[TextLayouterSimple]:

def _obj_to_mag_color(self, obj: CompositeObject):
"""
Extract the magnitude safely from the object and convert it into a color
Extract the magnitude safely from the object and convert it to a color
"""
try:
obj_mag = float(obj.mag)
Expand Down Expand Up @@ -331,7 +327,7 @@ def update(self, force=True):
text_lines = self.create_info_text()

# Clear Screen
self.draw.rectangle([0, 0, 128, 128], fill=self.colors.get(0))
self.draw.rectangle((0, 0, 128, 128), fill=self.colors.get(0))
line = 17
# Draw the closest objects
for obj_type, txt in text_lines:
Expand Down Expand Up @@ -360,15 +356,6 @@ def key_long_d(self):

def key_b(self):
self.current_mode = next(self.mode_cycle)
# if self.catalog_tracker.get_current_object() is None:
# self.object_display_mode = DM_DESC
# else:
# # switch object display text
# self.object_display_mode = (
# self.object_display_mode + 1 if self.object_display_mode < 2 else 0
# )
# self.update_object_info()
# self.update()

def background_update(self):
# catalog will be filtered by the UICatalog view
Expand Down

0 comments on commit 8537a0f

Please sign in to comment.