import numpy as np
import matplotlib.pyplot as plt
from astropy.coordinates import SkyCoord
import astropy.units as u
from mocpy import MOC
from tqdm import tqdm

import pandas as pd
from astropy.io import fits
from astropy.wcs import WCS

import os

Speeding up Butter

Run times are shown below each code cell.

skycell_csv = "./data/SkyCells/skycell_coordinates.csv"
skycell_df = pd.read_csv(skycell_csv)

def angular_distance(ra1, dec1, ra2, dec2):
    """
    Angular distance between two points on the sky
    """

    d1 = np.sin(np.radians(dec1)) * np.sin(np.radians(dec2))
    d2 = np.cos(np.radians(dec1)) * np.cos(np.radians(dec2)) * np.cos(np.radians(ra1 - ra2))
    return np.degrees(np.arccos(d1 + d2))

def max_radius(ra1, dec1, ra2, dec2):
    """
    Maximum angular distance between two points on the sky
    """

    return np.nanmax(angular_distance(ra1, dec1, ra2, dec2))

# def _image1(self, im1_file):
# """
# Opens the TESS image and extracts the necessary information
# """

im1_file = "../data/tess/20_3_3/tess2020019135923-s0020-3-3-0165-s_ffic.fits"
hdul = fits.open(im1_file)
super_data = hdul[1].data
data_shape = np.shape(super_data)
super_wcs = WCS(hdul[1].header)
im1_poly = super_wcs.calc_footprint()
hdul.close()

ra_centre, dec_centre = super_wcs.all_pix2world(data_shape[1]/2, data_shape[0]/2, 0)

WARNING: FITSFixedWarning: 'datfix' made the change 'Set MJD-OBS to 58867.584712 from DATE-OBS.
Set MJD-END to 58867.605545 from DATE-END'. [astropy.wcs.wcs]

# def _ravelling(self):
# """
# Unravels the TESS image to get the pixel coordinates
# """

t_y, t_x = data_shape
_ty, _tx = np.mgrid[:t_y, :t_x]

ty_input = _ty.ravel()
tx_input = _tx.ravel()

tpix_coord_input = np.column_stack([ty_input, tx_input])

_x_im1 = tpix_coord_input[:,1]
_y_im1 = tpix_coord_input[:,0]
_ra_im1, _dec_im1 = super_wcs.all_pix2world(_x_im1, _y_im1, 0)

ravelled_indices = np.array([(i, j) for i in range(data_shape[0]) for j in range(data_shape[1])])
ravelled_index = np.arange(data_shape[0] * data_shape[1])

1.6s

ra_corners = im1_poly[:, 0]
dec_corners = im1_poly[:, 1]

radius_var = max_radius(ra_centre, dec_centre, ra_corners, dec_corners) + np.sqrt(2) * 0.4

radius = angular_distance(ra_centre, dec_centre, skycell_df['RA'].to_numpy(), skycell_df['DEC'].to_numpy())

mask = radius < radius_var

skycells_list = skycell_df[mask].Name.to_numpy()
skycells_id = np.arange(len(skycells_list))

0.0s

skycell_wcs_csv = "./data/SkyCells/skycell_wcs.csv"
skycell_wcs_df = pd.read_csv(skycell_wcs_csv)

# i = 22
i = 915

skycell = skycells_list[i]
skycell_index = skycells_id[i]
buffer = 120

skycell, skycell_index

('skycell.2589.099', 915)

def ps1_wcs_function(skycell):
    """
    Grabs and stores the WCS information for the PanSTARRS1 images
    """

    other_temp = skycell_wcs_df[skycell_wcs_df['NAME'] == skycell].reset_index(drop=True)

    records = other_temp.to_dict(orient='records')
    header_dict = {k: v for d in records for k, v in d.items()}

    temp_head = fits.Header(header_dict)
    ps1_data_shape = (other_temp['NAXIS2'].iloc[0], other_temp['NAXIS1'].iloc[0])

    temp_wcs = WCS(temp_head)

    # header_dicts.append(temp_head)
    # ps1_wcs_master.append(temp_wcs)
    return temp_wcs, ps1_data_shape

# def _ps1_image(skycell):
# """
# Grabs the PanSTARRS1 skycells information for extraction
# """

# other_temp = skycell_wcs_df[skycell_wcs_df['NAME'] == skycell].reset_index(drop=True)
# ps1_data_shape = (other_temp['NAXIS2'].iloc[0], other_temp['NAXIS1'].iloc[0])

# temp_ind = complete_skycells[complete_skycells['Name'] == skycell].index[0]
# ps1_wcs = ps1_wcs_master[temp_ind]
ps1_wcs, ps1_data_shape = ps1_wcs_function(skycell)

corns = np.array([[buffer, buffer], [ps1_data_shape[1] - buffer, buffer], 
                [ps1_data_shape[1] - buffer, ps1_data_shape[0] - buffer], 
                [buffer, ps1_data_shape[0] - buffer]])

ra_corners_ps1, dec_corners_ps1 = ps1_wcs.all_pix2world(corns[:, 0], corns[:, 1], 0)

min_ps1_ra = np.min(ra_corners_ps1)
max_ps1_ra = np.max(ra_corners_ps1)

min_ps1_dec = np.min(dec_corners_ps1)
max_ps1_dec = np.max(dec_corners_ps1)

ps1_poly = np.column_stack((ra_corners_ps1, dec_corners_ps1))

0.0s

# def _pixel_vertices(self):
# """
# Grab the pixel vertices for the TESS image
# """

pixel_vertices = []

_ra, _dec = super_wcs.all_pix2world(tpix_coord_input[:, 1], tpix_coord_input[:, 0], 0)
mask = (_ra >= min_ps1_ra) & (_ra < max_ps1_ra) & (_dec >= min_ps1_dec) & (_dec < max_ps1_dec)
filtered_indices = np.where(mask)[0]

pix_centers = np.column_stack((_ra[filtered_indices], _dec[filtered_indices]))
pix_center_ra = _ra[filtered_indices]
pix_center_dec = _dec[filtered_indices]
im1_indices = ravelled_indices[filtered_indices]
im1_index = ravelled_index[filtered_indices]

for c in tpix_coord_input[filtered_indices]:
    x = c[1]
    y = c[0]

    upper_left = (x - 0.5, y - 0.5)
    upper_right = (x + 0.5, y - 0.5)
    lower_right = (x + 0.5, y + 0.5)
    lower_left = (x - 0.5, y + 0.5)

    t_poly = super_wcs.all_pix2world([upper_left, upper_right, lower_right, lower_left], 0)
    pixel_vertices.append(t_poly)

0.6s

coords = tpix_coord_input[filtered_indices]
x_coords = coords[:, 1]
y_coords = coords[:, 0]

corners = np.array([
    np.column_stack([x_coords - 0.5, y_coords - 0.5]),  # upper_left
    np.column_stack([x_coords + 0.5, y_coords - 0.5]),  # upper_right
    np.column_stack([x_coords + 0.5, y_coords + 0.5]),  # lower_right
    np.column_stack([x_coords - 0.5, y_coords + 0.5])   # lower_left
])

corners_reshaped = corners.transpose(1, 0, 2).reshape(-1, 2)

world_coords = super_wcs.all_pix2world(corners_reshaped, 0)

pixel_vertices_vectorized = world_coords.reshape(len(filtered_indices), 4, 2)

np.sum(np.array(pixel_vertices) != pixel_vertices_vectorized)

0

# def _ps1_ravelling(self):
    # """
    # PanSTARRS1 image unravelling
    # """

p_y, p_x = ps1_data_shape

py, px = np.mgrid[:p_y, :p_x]

py_input = py.ravel()
px_input = px.ravel()

ppix_coord_input = np.column_stack((py_input, px_input))

x2 = ppix_coord_input[:,1]
y2 = ppix_coord_input[:,0]
_ra2, _dec2 = ps1_wcs.all_pix2world(x2, y2, 0)

from itertools import compress

# def moccy(self):
    # """
    # Grabbing enclosed PanSTARRS1 pixels
    # """

ps1_skycoord = SkyCoord(ps1_poly, unit="deg", frame="icrs")
ps1_moc = MOC.from_polygon_skycoord(ps1_skycoord, complement=False, max_depth=21)

ps1_mask = ps1_moc.contains_lonlat(pix_center_ra * u.degree, pix_center_dec * u.degree)

enc_pix_vertices = list(compress(pixel_vertices, ps1_mask))
enc_pix_center_ra = np.asarray(pix_center_ra)[ps1_mask]
enc_pix_center_dec = np.asarray(pix_center_dec)[ps1_mask]
enc_pix_indices = np.asarray(im1_indices)[ps1_mask]
enc_pix_index = np.asarray(im1_index)[ps1_mask]

if len(enc_pix_vertices) > 0:
    enc_pix = []
    for t in enc_pix_vertices:
        pix_skycoord = SkyCoord(t, unit="deg", frame="icrs")
        tess_pix_moc = MOC.from_polygon_skycoord(pix_skycoord, complement=False, max_depth=21)
        enc_pix.append(tess_pix_moc)
else:
    skip = True
    print('Skipping:', skycell, skycell_index)

4.9s

enc_pix_vertices_small = enc_pix_vertices
# enc_pix_vertices_small = enc_pix_vertices[:500]

original but simplified code

## OLD but simplified
pix_ras = _ra2
pix_decs = _dec2

# output_vec = []
fll = np.full((ps1_data_shape[0], ps1_data_shape[1], 2), np.nan)

for i, tp in enumerate(enc_pix_vertices_small):
    pix_skycoord = SkyCoord(tp, unit="deg", frame="icrs")
    tess_pix_moc = MOC.from_polygon_skycoord(pix_skycoord, complement=False, max_depth=21)

    tess_ind = enc_pix_indices[i]

    _min_ra, _max_ra = np.min(tp[:, 0]) - 0.05, np.max(tp[:, 0]) + 0.05
    _min_dec, _max_dec = np.min(tp[:, 1]) - 0.05, np.max(tp[:, 1]) + 0.05

    search_indices = np.where((pix_ras >= _min_ra) &
                            (pix_ras <= _max_ra) &
                            (pix_decs >= _min_dec) &
                            (pix_decs <= _max_dec))[0]

    enc_ps1_pix_mask = tess_pix_moc.contains_lonlat(pix_ras[search_indices] * u.degree, pix_decs[search_indices] * u.degree)
    ps1_ind = np.arange(len(pix_ras))
    ps1_ind = ps1_ind[search_indices][enc_ps1_pix_mask]
    temp_ra = np.asarray(pix_ras[search_indices])[enc_ps1_pix_mask]
    temp_dec = np.asarray(pix_decs[search_indices])[enc_ps1_pix_mask]

    # output_vec.append((tess_ind, temp_ra, temp_dec, ps1_ind))

    unravel_index = np.unravel_index(ps1_ind, ps1_data_shape)
    length = len(unravel_index[0])
    zeros = np.zeros(length)
    tup = (unravel_index[0], unravel_index[1], zeros)

    fll[tup[0], tup[1], np.int16(tup[2])] = tess_ind[1]
    fll[tup[0], tup[1], np.int16(tup[2]) + 1] = tess_ind[0]

15m52.3s

New code

from numba import jit

pix_ras = _ra2
pix_decs = _dec2

polygons = np.array(enc_pix_vertices_small)
rust_result = MOC.filter_points_in_polygons(polygons=polygons, pix_ras=pix_ras, pix_decs=pix_decs, buffer=0.05, max_depth=21)

fll_1d_vec_rust = np.full((ps1_data_shape[0] * ps1_data_shape[1]), -1, dtype=np.int32)


@jit(nopython=True)
def populate_array_numba(fll_1d_vec_rust, rust_result, enc_pix_indices):
    for i in range(len(rust_result)):
        ps1_pix_in_tess = rust_result[i]
        tess_ind = enc_pix_index[i]
        for j in range(len(ps1_pix_in_tess)):
            fll_1d_vec_rust[ps1_pix_in_tess[j]] = tess_ind


populate_array_numba(fll_1d_vec_rust, rust_result, enc_pix_indices)

6m12.1s

61% reduction in run time

np.array_equiv(rust_result[-1], ps1_ind)

True

Original code

pix_obj_ind = np.arange(len(enc_pix))
pix_obj_etp = enc_pix
pix_obj_pix = enc_pix_indices[pix_obj_ind]

pix_indexs, tess_pix_mocs, tess_inds = pix_obj_ind, pix_obj_etp, pix_obj_pix

# def initialize_moc_pixel(self, pix_obj):
# """
# Initializes the PanSTARRS1 pixels to TESS pixels matching
# """

pix_ras = _ra2
pix_decs = _dec2
poly = enc_pix_vertices

output = []

# pix_indexs, tess_pix_mocs, tess_inds = pix_obj

if len(pix_indexs) != 0:

    for i in range(len(pix_indexs)):
        pix_index = pix_indexs[i]
        tess_ind = tess_inds[i]
        tess_pix_moc = tess_pix_mocs[i]

        tp = poly[pix_index]

        _min_ra, _max_ra = np.min(tp[:, 0]) - 0.05, np.max(tp[:, 0]) + 0.05
        _min_dec, _max_dec = np.min(tp[:, 1]) - 0.05, np.max(tp[:, 1]) + 0.05

        search_indices = np.where((pix_ras >= _min_ra) &
                                (pix_ras <= _max_ra) &
                                (pix_decs >= _min_dec) &
                                (pix_decs <= _max_dec))[0]

        enc_ps1_pix_mask = tess_pix_moc.contains_lonlat(pix_ras[search_indices] * u.degree, pix_decs[search_indices] * u.degree)
        ps1_ind = np.arange(len(pix_ras))
        ps1_ind = ps1_ind[search_indices][enc_ps1_pix_mask]
        temp_ra = np.asarray(pix_ras[search_indices])[enc_ps1_pix_mask]
        temp_dec = np.asarray(pix_decs[search_indices])[enc_ps1_pix_mask]

        output.append((tess_ind, temp_ra, temp_dec, ps1_ind))

    # fll = _arrayifying(output)
    # _fitsify(fll)
else:
    pass

pix_output = output

14m54.9s

# def _arrayifying(self, pix_output):
    # """
    # Converts the PanSTARRS1 pixels to TESS pixels output to an array
    # """

fll = np.full((ps1_data_shape[0], ps1_data_shape[1], 2), np.nan)

for i in range(len(pix_output)):
    unravel_index = np.unravel_index(pix_output[i][-1], ps1_data_shape)
    length = len(unravel_index[0])
    zeros = np.zeros(length)
    tup = (unravel_index[0], unravel_index[1], zeros)

    fll[tup[0], tup[1], np.int16(tup[2])] = pix_output[i][0][1]
    fll[tup[0], tup[1], np.int16(tup[2]) + 1] = pix_output[i][0][0]

# return fll

0.7s

# checking if the output matches the Rust implementation
fll_1d = (fll[:, :, 1] * data_shape[1] + fll[:, :, 0]).flatten()
fll_1d[np.isnan(fll_1d)] = -1
np.array_equal(fll_1d_vec_rust, fll_1d, equal_nan=True)

True

plt.imshow(fll_1d.reshape(ps1_data_shape[0], ps1_data_shape[1]), origin="lower", cmap="gray")
plt.colorbar()
plt.show()