| import torch |
| import numpy as np |
| import healpy as hp |
| from sklearn.neighbors import BallTree |
|
|
|
|
| def rad_to_xyz(lonlat: torch.Tensor): |
| """Convert lon-lat (in radians) to unit sphere xyz.""" |
| lon = lonlat[..., 0] |
| lat = lonlat[..., 1] |
|
|
| x = torch.cos(lat) * torch.cos(lon) |
| y = torch.cos(lat) * torch.sin(lon) |
| z = torch.sin(lat) |
|
|
| return torch.stack([x, y, z], axis=-1) |
|
|
|
|
| def get_healpix_grid(nside: int) -> torch.Tensor: |
| """Return HEALPix grid coordinates as array of shape (npix, 2).""" |
| indices = np.arange(hp.nside2npix(nside)) |
| theta, phi = hp.pix2ang(nside, indices, nest=True) |
|
|
| phi = np.rad2deg(phi) |
| theta = (90. - np.rad2deg(theta)) |
|
|
| phi = torch.from_numpy(phi) |
| theta = torch.from_numpy(theta) |
|
|
| return torch.stack((phi, theta), axis=-1).float() |
|
|
|
|
| def get_neighbors(pos_from: np.ndarray, pos_to: np.ndarray, k: int = 8) -> tuple: |
| """Build a BallTree and query k nearest neighbors with haversine metric.""" |
| pos_from_rad = pos_from[:, ::-1] |
| pos_to_rad = pos_to[:, ::-1] |
|
|
| tree = BallTree(pos_from_rad, metric='haversine') |
| _, neighbors = tree.query(pos_to_rad, k=k) |
| return neighbors |
|
|
