WIP

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)

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):

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)