Add benchmark of integrated probability

healpix_alchemy/__init__.py

@@ -1,4 +1,5 @@
+from ._temporary_table import temporary_table
 from .types import Point, Tile
 from . import func  # noqa: F401
 
-__all__ = ('Point', 'Tile')
+__all__ = ('temporary_table', 'Point', 'Tile')

healpix_alchemy/_temporary_table.py

@@ -0,0 +1,43 @@
+"""SQLAlchemy temporary tables.
+
+Inspired by https://github.com/sqlalchemy/sqlalchemy/wiki/Views
+"""
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.schema import DDLElement
+from sqlalchemy.sql import TableClause
+
+from uuid import uuid4
+
+__all__ = ('temporary_table',)
+
+
+class CreateTemporaryTable(DDLElement):
+
+    inherit_cache = False
+
+    def __init__(self, selectable, name):
+        self.selectable = selectable
+        self.name = name
+
+
+@compiles(CreateTemporaryTable)
+def _create_temporary_table(element, compiler, **_):
+    selectable = compiler.sql_compiler.process(
+        element.selectable, literal_binds=True)
+    return f'CREATE TEMPORARY TABLE {element.name} AS {selectable}'
+
+
+class temporary_table(TableClause):
+
+    inherit_cache = False
+
+    def __init__(self, selectable, name=None):
+        if name is None:
+            name = 'temp_' + str(uuid4()).replace('-', '_')
+        super().__init__(name)
+        self._selectable = selectable
+        self._columns._populate_separate_keys(
+            col._make_proxy(self) for col in selectable.selected_columns)
+
+    def create(self, bind):
+        return bind.execute(CreateTemporaryTable(self._selectable, self.name))

healpix_alchemy/func.py

@@ -5,4 +5,4 @@
 
 
 def union(tiles):
-    return _func.unnest(_func.range_agg(tiles), type_=_Tile)
+    return _func.multirange_unnest(_func.range_agg(tiles), type_=_Tile)

healpix_alchemy/tests/benchmarks/conftest.py

@@ -33,11 +33,17 @@ 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 random_sky_map(cursor, tables):
-    return data.get_random_sky_map(20_000, cursor)
+@pytest.fixture()
+def ztf_fields(cursor, tables):
+    return data.get_ztf_fields(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

@@ -6,12 +6,15 @@
 
 """
 import io
+import functools
 
 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
+from regions import Regions
 
 from ...constants import HPX, LEVEL, PIXEL_AREA
 from ...types import Tile
@@ -25,19 +28,18 @@
 
 
 def get_ztf_footprint_corners():
-    """Return the corner offsets of the ZTF footprint.
-
-    Notes
-    -----
-    This polygon is smaller than the spatial extent of the true ZTF field of
-    view, but has approximately the same area because the real ZTF field of
-    view has chip gaps.
+    """Return the corner offsets of the ZTF footprint."""
+    url = '/Users/lpsinger/src/skyportal/data/ZTF_Region.reg'
+    regions = Regions.read(url, cache=True)
+
+    vertices = SkyCoord(
+        [region.vertices.ra for region in regions],
+        [region.vertices.dec for region in regions]
+    ).transform_to(
+        SkyCoord(0 * u.deg, 0 * u.deg).skyoffset_frame()
+    )
 
-    For the real ZTF footprint, use the region file
-    https://github.com/skyportal/skyportal/blob/main/data/ZTF_Region.reg.
-    """
-    x = 6.86 / 2
-    return [-x, +x, +x, -x] * u.deg, [-x, -x, +x, +x] * u.deg
+    return vertices.lon, vertices.lat
 
 
 def get_footprints_grid(lon, lat, offsets):
@@ -62,7 +64,7 @@ def get_footprints_grid(lon, lat, offsets):
     """
     lon = np.repeat(lon[np.newaxis, :], len(offsets), axis=0)
     lat = np.repeat(lat[np.newaxis, :], len(offsets), axis=0)
-    result = SkyCoord(lon, lat, frame=offsets[:, np.newaxis].skyoffset_frame())
+    result = SkyCoord(lon, lat, frame=offsets.reshape((-1,) + (1,) * (lon.ndim - 1)).skyoffset_frame())
     return result.icrs
 
 
@@ -72,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)
@@ -82,10 +91,19 @@ def get_random_galaxies(n, cursor):
     return points
 
 
+def union_moc_func(a, b):
+    return a.union(b)
+
+
+def get_union_moc(footprints):
+    mocs = (MOC.from_polygon_skycoord(footprint) for footprint in footprints)
+    return functools.reduce(union_moc_func, mocs)
+
+
 def get_random_fields(n, cursor):
     centers = SkyCoord(get_random_points(n, RANDOM_FIELDS_SEED))
     footprints = get_footprints_grid(*get_ztf_footprint_corners(), centers)
-    mocs = [MOC.from_polygon_skycoord(footprint) for footprint in footprints]
+    mocs = [get_union_moc(footprint) for footprint in footprints]
 
     f = io.StringIO('\n'.join(f'{i}' for i in range(len(mocs))))
     cursor.copy_from(f, Field.__tablename__)
@@ -101,34 +119,62 @@ def get_random_fields(n, cursor):
     return mocs
 
 
-def get_random_sky_map(n, cursor):
-    rng = np.random.default_rng(RANDOM_SKY_MAP_SEED)
-    # Make a randomly subdivided sky map
-    npix = HPX.npix
-    tiles = np.arange(0, npix + 1, 4 ** LEVEL).tolist()
-    while len(tiles) < n:
-        i = rng.integers(len(tiles))
-        lo = 0 if i == 0 else tiles[i - 1]
-        hi = tiles[i]
-        diff = (hi - lo) // 4
-        if diff == 0:
-            continue
-        tiles.insert(i, hi - diff)
-        tiles.insert(i, hi - 2 * diff)
-        tiles.insert(i, hi - 3 * diff)
-
-    probdensity = rng.uniform(0, 1, size=len(tiles) - 1)
-    probdensity /= np.sum(np.diff(tiles) * probdensity) * PIXEL_AREA
-
-    f = io.StringIO('1')
-    cursor.copy_from(f, Skymap.__tablename__)
+@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__)
 
     f = io.StringIO(
         '\n'.join(
-            f'1\t[{lo},{hi})\t{p}'
-            for lo, hi, p in zip(tiles[:-1], tiles[1:], probdensity)
+            f'{i}\t{hpx}'
+            for i, moc in zip(field_data['id'], mocs) for hpx in Tile.tiles_from(moc)
         )
     )
-    cursor.copy_from(f, SkymapTile.__tablename__)
+    cursor.copy_from(f, FieldTile.__tablename__)
+
+    return mocs
+
+
+def get_random_sky_map(n, nmaps, cursor):
+    rng = np.random.default_rng(RANDOM_SKY_MAP_SEED)
+    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()
+        while len(tiles) < n:
+            i = rng.integers(len(tiles))
+            lo = 0 if i == 0 else tiles[i - 1]
+            hi = tiles[i]
+            diff = (hi - lo) // 4
+            if diff == 0:
+                continue
+            tiles.insert(i, hi - diff)
+            tiles.insert(i, hi - 2 * diff)
+            tiles.insert(i, hi - 3 * diff)
+
+        probdensity = rng.uniform(0, 1, size=len(tiles) - 1)
+        probdensity /= np.sum(np.diff(tiles) * probdensity) * PIXEL_AREA
+
+        f = io.StringIO(f'{skymap_id}')
+        cursor.copy_from(f, Skymap.__tablename__)
+
+        f = io.StringIO(
+            '\n'.join(
+                f'{skymap_id}\t[{lo},{hi})\t{p}'
+                for lo, hi, p in zip(tiles[:-1], tiles[1:], probdensity)
+            )
+        )
+        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)