brick_queries.py

import healpy as hp
import numpy as np
from .brick_core import generate_bricks
from .brick_utils import angular_distance

def bricks_in_disc(nside, ra_center, dec_center, radius_deg, overlap):
    """
    Find all bricks whose centers are within a circular region.
    
    Parameters
    ----------
    nside : int
        HEALPix resolution parameter.
    ra_center : float
        Right Ascension of the circular region center (deg).
    dec_center : float
        Declination of the circular region center (deg).
    radius_deg : float
        Search radius of the circular region (deg).
    overlap : float
        Overlap factor to enlarge brick radius 
        (1.0 = no overlap, >1.0 = enlarged).
    
    Returns
    -------
    list of Brick
        A list of Brick objects whose centers fall within the given region.
    """
    theta = np.radians(90 - dec_center)
    phi = np.radians(ra_center)
    radius_rad = np.radians(radius_deg)
    vec = hp.ang2vec(theta, phi)
    pix_ids = hp.query_disc(nside, vec, radius_rad)
    return [generate_bricks(nside, pid, overlap=overlap) for pid in pix_ids]


def neighbor_bricks(nside, brick_id, overlap):
    """
    Get all neighboring bricks of a given brick.

    Parameters
    ----------
    nside : int
        HEALPix resolution parameter.
    brick_id : int
        The HEALPix pixel (brick) ID for which neighbors are to be found.
    overlap : float
        Overlap factor to enlarge brick radius 
        (1.0 = no overlap, >1.0 = enlarged).

    Returns
    -------
    list of Brick
        A list of Brick objects corresponding to the neighboring pixels.
    """
    neighbors = hp.get_all_neighbours(nside, brick_id)
    neighbors = [n for n in neighbors if n >= 0]
    return [generate_bricks(nside, n, overlap=overlap) for n in neighbors]

def nearest_brick_with_distance(nside, ra, dec, overlap):
    """
    Find the nearest brick to a given (RA, Dec) position and its angular distance.

    Parameters
    ----------
    nside : int
        HEALPix resolution parameter.
    ra : float
        Right Ascension of the target position (deg).
    dec : float
        Declination of the target position (deg).
    overlap : float
        Overlap factor to enlarge brick radius 
        (1.0 = no overlap, >1.0 = enlarged).

    Returns
    -------
    tuple
        brick : Brick
            The nearest Brick object to the given (RA, Dec) position.
        dist_deg : float
            The angular distance between the point and the brick center (deg).
    """
    theta = np.radians(90 - dec)
    phi = np.radians(ra)
    pix_id = hp.ang2pix(nside, theta, phi)
    brick = generate_bricks(nside, pix_id, overlap=overlap)
    ra1, dec1 = np.radians(ra), np.radians(dec)
    ra2, dec2 = np.radians(brick.ra), np.radians(brick.dec)
    dist_deg = np.degrees(angular_distance(ra1, dec1, ra2, dec2))
    return brick, dist_deg

def point_in_brick_with_distance(nside, ra, dec, brick_id, overlap):
    """
    Check if a given point (RA, Dec) lies within a specific brick and 
    return its angular distance to the brick center.

    Parameters
    ----------
    nside : int
        HEALPix resolution parameter.
    ra : float
        Right Ascension of the target point (deg).
    dec : float
        Declination of the target point (deg).
    brick_id : int
        The HEALPix pixel (brick) ID to check.
    overlap : float
        Overlap factor to enlarge brick radius 
        (1.0 = no overlap, >1.0 = enlarged).

    Returns
    -------
    tuple
        inside : bool
            True if the point lies within the brick (distance <= brick.radius).
        dist : float
            Angular distance between the point and the brick center (deg).
    """
    brick = generate_bricks(nside, brick_id, overlap=overlap)
    ra1, dec1 = np.radians(ra), np.radians(dec)
    ra2, dec2 = np.radians(brick.ra), np.radians(brick.dec)
    dist = np.degrees(angular_distance(ra1, dec1, ra2, dec2))
    return dist <= brick.radius, dist

def bricks_containing_point(nside, ra, dec, overlap, search_margin=5):
    """
    Find all bricks that contain a given point (RA, Dec).
    
    Parameters
    ----------
    nside : int
        HEALPix resolution parameter.
    ra : float
        Right Ascension of the target point (deg).
    dec : float
        Declination of the target point (deg).
    overlap : float
        Overlap factor to enlarge brick radius 
        (1.0 = no overlap, >1.0 = enlarged).
    search_margin : float, optional
        Search margin to expand the maximum search radius. Default is 5.

    Returns
    -------
    list of Brick
        A list of Brick objects whose radius contains the given point.
    """
    max_radius = (180 / (np.pi * np.sqrt(3) * nside)) * search_margin * overlap
    theta = np.radians(90 - dec)
    phi = np.radians(ra)
    vec = hp.ang2vec(theta, phi)
    pix_ids = hp.query_disc(nside, vec, np.radians(max_radius), inclusive=True)
    result = []
    for pid in pix_ids:
        brick = generate_bricks(nside, pid, overlap=overlap)
        ra1, dec1 = np.radians(ra), np.radians(dec)
        ra2, dec2 = np.radians(brick.ra), np.radians(brick.dec)
        dist = np.degrees(angular_distance(ra1, dec1, ra2, dec2))
        if dist <= brick.radius:
            result.append(brick)
    return result


class BrickQueries:
    def __init__(self, nside, overlap=1.0):
        self.nside = nside
        self.overlap = overlap
        self.methods = {
            "disc_bricks": self.get_bricks_in_disc,
            "neighbors": self.get_neighbor_bricks,
            "nearest": self.get_nearest_brick_with_distance,
            "contains_point": self.point_in_brick,
            "find_bricks": self.get_bricks_containing_point
        }

    def __getattr__(self, name):
        if name in self.methods:
            return self.methods[name]
        else:
            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")

    def get_bricks_in_disc(self, ra_center, dec_center, radius_deg):
        return bricks_in_disc(self.nside, ra_center, dec_center, radius_deg, self.overlap)

    def get_neighbor_bricks(self, brick_id):
        return neighbor_bricks(self.nside, brick_id, self.overlap)

    def get_nearest_brick_with_distance(self, ra, dec):
        return nearest_brick_with_distance(self.nside, ra, dec, self.overlap)

    def point_in_brick(self, ra, dec, brick_id):
        return point_in_brick_with_distance(self.nside, ra, dec, brick_id, self.overlap)

    def get_bricks_containing_point(self, ra, dec, search_margin=5):
        return bricks_containing_point(self.nside, ra, dec, self.overlap, search_margin)