WIP

healpix_alchemy/tests/benchmarks/conftest.py

@@ -42,6 +42,7 @@ 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)
+@pytest.fixture(params=np.geomspace(1, 1000, 10, dtype=int).tolist())
+# @pytest.fixture(params=[50])
+def random_sky_map(cursor, tables, request):
+    return data.get_random_sky_map(20_000, request.param, cursor)

healpix_alchemy/tests/benchmarks/data.py

@@ -119,11 +119,17 @@ def get_random_fields(n, cursor):
     return mocs
 
 
-def get_ztf_fields(cursor):
+@functools.cache
+def get_ztf_mocs():
     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]
+    return field_data, mocs
+
+
+def get_ztf_fields(cursor):
+    field_data, mocs = get_ztf_mocs()
 
     f = io.StringIO('\n'.join(f'{i}' for i in field_data['id']))
     cursor.copy_from(f, Field.__tablename__)
@@ -139,9 +145,9 @@ def get_ztf_fields(cursor):
     return mocs
 
 
-def get_random_sky_map(n, cursor):
+def get_random_sky_map(n, nmaps, cursor):
     rng = np.random.default_rng(RANDOM_SKY_MAP_SEED)
-    for skymap_id in range(10, 0, -1):
+    for skymap_id in range(nmaps, 0, -1):
         # Make a randomly subdivided sky map
         npix = HPX.npix
         tiles = np.arange(0, npix + 1, 4 ** LEVEL).tolist()
@@ -170,4 +176,5 @@ def get_random_sky_map(n, cursor):
         )
         cursor.copy_from(f, SkymapTile.__tablename__)
 
-    return tiles, probdensity
+    # return tiles, probdensity
+    return nmaps

healpix_alchemy/tests/benchmarks/models.py

@@ -23,7 +23,7 @@ class Field(Base):
 
 
 class FieldTile(Base):
-    id = sa.Column(sa.ForeignKey(Field.id), primary_key=True, index=True)
+    id = sa.Column(sa.ForeignKey(Field.id), primary_key=True)
     hpx = sa.Column(Tile, primary_key=True, index=True)
 
 

healpix_alchemy/tests/benchmarks/test_benchmarks.py

@@ -1,5 +1,5 @@
 from functools import reduce
-import time
+from .explain import explain
 
 import numpy as np
 import sqlalchemy as sa
@@ -108,7 +108,7 @@ def _func(query, expected):
 #     bench(query)
 
 
-def test_integrated_probability(bench, ztf_fields, random_sky_map):
+def test_integrated_probability(bench, ztf_fields, random_sky_map, session):
     """Find integrated probability within union of fields."""
     union = sa.select(
         func.union(FieldTile.hpx).label('hpx')
@@ -126,5 +126,8 @@ def test_integrated_probability(bench, ztf_fields, random_sky_map):
         union.columns.hpx.overlaps(SkymapTile.hpx)
     )
 
-    # Run benchmark
-    bench(query)
+    session.execute('ANALYZE')
+    session.execute("LOAD 'auto_explain'")
+    session.execute("SET auto_explain.log_min_duration=0")
+    session.execute(query).all()
+    session.execute("SET auto_explain.log_min_duration='1000s'")

healpix_alchemy/types.py

@@ -93,4 +93,4 @@ def _create_indices(index, parent):
         isinstance(index.expressions[0], sa.Column) and
         isinstance(index.expressions[0].type, Tile)
     ):
-        index.dialect_options['postgresql']['using'] = 'spgist'
+        index.dialect_options['postgresql']['using'] = 'gist'