diff --git a/healpix_alchemy/tests/benchmarks/conftest.py b/healpix_alchemy/tests/benchmarks/conftest.py index f8ceca3..bc2cd6d 100644 --- a/healpix_alchemy/tests/benchmarks/conftest.py +++ b/healpix_alchemy/tests/benchmarks/conftest.py @@ -32,11 +32,16 @@ def random_galaxies(cursor, tables): return data.get_random_galaxies(40_000, cursor) -@pytest.fixture(params=np.geomspace(1, 10_000, 10, dtype=int).tolist()) +@pytest.fixture(params=[2000]) def random_fields(cursor, tables, request): return data.get_random_fields(request.param, cursor) +@pytest.fixture() +def ztf_fields(cursor, tables): + return data.get_ztf_fields(cursor) + + @pytest.fixture def random_sky_map(cursor, tables): return data.get_random_sky_map(20_000, cursor) diff --git a/healpix_alchemy/tests/benchmarks/data.py b/healpix_alchemy/tests/benchmarks/data.py index a61fc72..9f3761e 100644 --- a/healpix_alchemy/tests/benchmarks/data.py +++ b/healpix_alchemy/tests/benchmarks/data.py @@ -10,6 +10,7 @@ from astropy.coordinates import SkyCoord, uniform_spherical_random_surface from astropy import units as u +from astropy.table import Table from mocpy import MOC import numpy as np import pytest @@ -73,6 +74,13 @@ def get_random_points(n, seed): return uniform_spherical_random_surface(n) +def get_ztf_field_data(): + url = 'https://raw.githubusercontent.com/ZwickyTransientFacility/ztf_information/master/field_grid/ZTF_Fields.txt' + data = Table.read(url, cache=True, format='ascii', comment='%') + data.rename_columns(['col1', 'col2', 'col3'], ['id', 'ra', 'dec']) + return data + + def get_random_galaxies(n, cursor): points = SkyCoord(get_random_points(n, RANDOM_GALAXIES_SEED)) hpx = HPX.skycoord_to_healpix(points) @@ -111,6 +119,26 @@ def get_random_fields(n, cursor): return mocs +def get_ztf_fields(cursor): + field_data = get_ztf_field_data() + centers = SkyCoord(field_data['ra'] * u.deg, field_data['dec'] * u.deg) + footprints = get_footprints_grid(*get_ztf_footprint_corners(), centers) + mocs = [get_union_moc(footprint) for footprint in footprints] + + f = io.StringIO('\n'.join(f'{i}' for i in field_data['id'])) + cursor.copy_from(f, Field.__tablename__) + + f = io.StringIO( + '\n'.join( + f'{i}\t{hpx}' + for i, moc in zip(field_data['id'], mocs) for hpx in Tile.tiles_from(moc) + ) + ) + cursor.copy_from(f, FieldTile.__tablename__) + + return mocs + + def get_random_sky_map(n, cursor): rng = np.random.default_rng(RANDOM_SKY_MAP_SEED) for skymap_id in range(10, 0, -1): diff --git a/healpix_alchemy/tests/benchmarks/test_benchmarks.py b/healpix_alchemy/tests/benchmarks/test_benchmarks.py index 7e2d3dd..d8674fb 100644 --- a/healpix_alchemy/tests/benchmarks/test_benchmarks.py +++ b/healpix_alchemy/tests/benchmarks/test_benchmarks.py @@ -108,10 +108,12 @@ def _func(query, expected): # bench(query) -def test_integrated_probability(bench, random_fields, random_sky_map): +def test_integrated_probability(bench, ztf_fields, random_sky_map): """Find integrated probability within union of fields.""" union = sa.select( func.union(FieldTile.hpx).label('hpx') + ).filter( + FieldTile.id.between(100, 120) ).subquery() area = (union.columns.hpx * SkymapTile.hpx).area @@ -124,4 +126,5 @@ def test_integrated_probability(bench, random_fields, random_sky_map): union.columns.hpx.overlaps(SkymapTile.hpx) ) + # Run benchmark bench(query)