From df6013dd3999831361298874866ef77984149b80 Mon Sep 17 00:00:00 2001 From: Leo Singer Date: Tue, 3 Sep 2024 07:56:34 -0400 Subject: [PATCH] Run pre-commit hooks for the first time --- .github/workflows/runtime_plots.py | 25 +++++---- conftest.py | 20 +++++--- healpix_alchemy/__init__.py | 4 +- healpix_alchemy/constants.py | 8 +-- healpix_alchemy/func.py | 1 + healpix_alchemy/types.py | 32 ++++++------ tests/benchmarks/conftest.py | 2 +- tests/benchmarks/data.py | 48 +++++++++--------- tests/benchmarks/models.py | 2 +- tests/benchmarks/test_benchmarks.py | 78 ++++++++++++----------------- 10 files changed, 107 insertions(+), 113 deletions(-) diff --git a/.github/workflows/runtime_plots.py b/.github/workflows/runtime_plots.py index ff69751..3d3c3ef 100644 --- a/.github/workflows/runtime_plots.py +++ b/.github/workflows/runtime_plots.py @@ -3,16 +3,23 @@ import matplotlib.pyplot as plt import pandas as pd -with open('benchmark_results.json') as f: +with open("benchmark_results.json") as f: data = json.load(f) -df = pd.json_normalize(data['benchmarks']) -df['func'] = df['name'].str.split('[', expand=True)[0] -df['fullfunc'] = df['fullname'].str.split('[', expand=True)[0] -df['param'] = pd.to_numeric(df['param']) -for func, group in df.groupby('func'): +df = pd.json_normalize(data["benchmarks"]) +df["func"] = df["name"].str.split("[", expand=True)[0] +df["fullfunc"] = df["fullname"].str.split("[", expand=True)[0] +df["param"] = pd.to_numeric(df["param"]) +for func, group in df.groupby("func"): fig, ax = plt.subplots() - group.plot.scatter(x='param', y='stats.mean', yerr='stats.stddev', - loglog=True, xlabel='N', ylabel='Runtime (s)', ax=ax) + group.plot.scatter( + x="param", + y="stats.mean", + yerr="stats.stddev", + loglog=True, + xlabel="N", + ylabel="Runtime (s)", + ax=ax, + ) fig.tight_layout() - fig.savefig(f'plots/{func}') + fig.savefig(f"plots/{func}") diff --git a/conftest.py b/conftest.py index 61d3c2e..d15388b 100644 --- a/conftest.py +++ b/conftest.py @@ -1,21 +1,25 @@ """Pytest configuration for running the doctests in README.md.""" + from unittest.mock import Mock -import sqlalchemy as sa + import pytest +import sqlalchemy as sa @pytest.fixture def engine(postgresql): """Create an SQLAlchemy engine with a disposable PostgreSQL database.""" - return sa.create_engine('postgresql+psycopg://', - poolclass=sa.pool.StaticPool, - pool_reset_on_return=None, - creator=lambda: postgresql) + return sa.create_engine( + "postgresql+psycopg://", + poolclass=sa.pool.StaticPool, + pool_reset_on_return=None, + creator=lambda: postgresql, + ) @pytest.fixture(autouse=True) def add_mock_create_engine(monkeypatch, request): """Monkey patch sqlalchemy.create_engine for doctests in README.md.""" - if request.node.name == 'README.md': - engine = request.getfixturevalue('engine') - monkeypatch.setattr(sa, 'create_engine', Mock(return_value=engine)) + if request.node.name == "README.md": + engine = request.getfixturevalue("engine") + monkeypatch.setattr(sa, "create_engine", Mock(return_value=engine)) diff --git a/healpix_alchemy/__init__.py b/healpix_alchemy/__init__.py index 5a627e0..6b4119c 100644 --- a/healpix_alchemy/__init__.py +++ b/healpix_alchemy/__init__.py @@ -1,4 +1,4 @@ -from .types import Point, Tile from . import func # noqa: F401 +from .types import Point, Tile -__all__ = ('Point', 'Tile') +__all__ = ("Point", "Tile") diff --git a/healpix_alchemy/constants.py b/healpix_alchemy/constants.py index 3daeff0..1c771e4 100644 --- a/healpix_alchemy/constants.py +++ b/healpix_alchemy/constants.py @@ -1,14 +1,14 @@ -from astropy.coordinates import ICRS +import sqlalchemy as sa from astropy import units as u -from astropy_healpix import level_to_nside, HEALPix +from astropy.coordinates import ICRS +from astropy_healpix import HEALPix, level_to_nside from mocpy import MOC -import sqlalchemy as sa LEVEL = MOC.MAX_ORDER """Base HEALPix resolution. This is the maximum HEALPix level that can be stored in a signed 8-byte integer data type.""" -HPX = HEALPix(nside=level_to_nside(LEVEL), order='nested', frame=ICRS()) +HPX = HEALPix(nside=level_to_nside(LEVEL), order="nested", frame=ICRS()) """HEALPix projection object.""" PIXEL_AREA = HPX.pixel_area.to_value(u.sr) diff --git a/healpix_alchemy/func.py b/healpix_alchemy/func.py index 35b2e9f..18929d3 100644 --- a/healpix_alchemy/func.py +++ b/healpix_alchemy/func.py @@ -1,4 +1,5 @@ """SQLAlchemy functions.""" + from sqlalchemy import func as _func from .types import Tile as _Tile diff --git a/healpix_alchemy/types.py b/healpix_alchemy/types.py index d01bbf0..648abe9 100644 --- a/healpix_alchemy/types.py +++ b/healpix_alchemy/types.py @@ -1,21 +1,21 @@ """SQLAlchemy types for multiresolution HEALPix data.""" + from collections.abc import Sequence from numbers import Integral +import numpy as np +import sqlalchemy as sa from astropy.coordinates import SkyCoord from astropy_healpix import uniq_to_level_ipix from mocpy import MOC -import numpy as np -import sqlalchemy as sa from sqlalchemy.dialects.postgresql import INT8RANGE from .constants import HPX, LEVEL, PIXEL_AREA_LITERAL -__all__ = ('Point', 'Tile') +__all__ = ("Point", "Tile") class Point(sa.TypeDecorator): - cache_ok = True impl = sa.BigInteger @@ -30,7 +30,6 @@ def process_bind_param(self, value, dialect): class Tile(sa.TypeDecorator): - cache_ok = True impl = INT8RANGE @@ -40,11 +39,10 @@ def process_bind_param(self, value, dialect): shift = 2 * (LEVEL - level) value = (ipix << shift, (ipix + 1) << shift) if isinstance(value, Sequence) and len(value) == 2: - value = f'[{value[0]},{value[1]})' + value = f"[{value[0]},{value[1]})" return value class comparator_factory(INT8RANGE.comparator_factory): - @property def lower(self): return sa.func.lower(self, type_=Point) @@ -68,29 +66,29 @@ def tiles_from(cls, obj): elif isinstance(obj, SkyCoord): return cls.tiles_from_polygon_skycoord(obj) else: - raise TypeError('Unknown type') + raise TypeError("Unknown type") @classmethod def tiles_from_polygon_skycoord(cls, polygon): return cls.tiles_from_moc( - MOC.from_polygon_skycoord( - polygon.transform_to(HPX.frame))) + MOC.from_polygon_skycoord(polygon.transform_to(HPX.frame)) + ) @classmethod def tiles_from_moc(cls, moc): - return (f'[{lo},{hi})' for lo, hi in moc.to_depth29_ranges) + return (f"[{lo},{hi})" for lo, hi in moc.to_depth29_ranges) -@sa.event.listens_for(sa.Index, 'after_parent_attach') +@sa.event.listens_for(sa.Index, "after_parent_attach") def _create_indices(index, parent): """Set index method to SP-GiST_ for any indexed Tile or Region columns. .. _SP-GiST: https://www.postgresql.org/docs/current/spgist.html """ if ( - index._column_flag and - len(index.expressions) == 1 and - isinstance(index.expressions[0], sa.Column) and - isinstance(index.expressions[0].type, Tile) + index._column_flag + and len(index.expressions) == 1 + and isinstance(index.expressions[0], sa.Column) + and isinstance(index.expressions[0].type, Tile) ): - index.dialect_options['postgresql']['using'] = 'spgist' + index.dialect_options["postgresql"]["using"] = "spgist" diff --git a/tests/benchmarks/conftest.py b/tests/benchmarks/conftest.py index 5cfcdae..eef311d 100644 --- a/tests/benchmarks/conftest.py +++ b/tests/benchmarks/conftest.py @@ -1,6 +1,6 @@ import numpy as np -from sqlalchemy import orm import pytest +from sqlalchemy import orm from . import data, models diff --git a/tests/benchmarks/data.py b/tests/benchmarks/data.py index 0b21d38..c5fafbd 100644 --- a/tests/benchmarks/data.py +++ b/tests/benchmarks/data.py @@ -5,22 +5,21 @@ We use the psycopg ``copy`` rather than SQLAlchemy for fast insertion. """ -from astropy.coordinates import SkyCoord, uniform_spherical_random_surface -from astropy import units as u -from mocpy import MOC + import numpy as np import pytest +from astropy import units as u +from astropy.coordinates import SkyCoord, uniform_spherical_random_surface +from mocpy import MOC from healpix_alchemy.constants import HPX, LEVEL, PIXEL_AREA from healpix_alchemy.types import Tile -from .models import Galaxy, Field, FieldTile, Skymap, SkymapTile +from .models import Field, FieldTile, Galaxy, Skymap, SkymapTile -( - RANDOM_GALAXIES_SEED, - RANDOM_FIELDS_SEED, - RANDOM_SKY_MAP_SEED -) = np.random.SeedSequence(12345).spawn(3) +(RANDOM_GALAXIES_SEED, RANDOM_FIELDS_SEED, RANDOM_SKY_MAP_SEED) = ( + np.random.SeedSequence(12345).spawn(3) +) def get_ztf_footprint_corners(): @@ -67,7 +66,7 @@ def get_footprints_grid(lon, lat, offsets): def get_random_points(n, seed): with pytest.MonkeyPatch.context() as monkeypatch: - monkeypatch.setattr(np, 'random', np.random.default_rng(seed)) + monkeypatch.setattr(np, "random", np.random.default_rng(seed)) return uniform_spherical_random_surface(n) @@ -75,8 +74,8 @@ def get_random_galaxies(n, cursor): points = SkyCoord(get_random_points(n, RANDOM_GALAXIES_SEED)) hpx = HPX.skycoord_to_healpix(points) - with cursor.copy(f'COPY {Galaxy.__tablename__} (hpx) FROM STDIN') as copy: - copy.write('\n'.join(f'{i}' for i in hpx)) + with cursor.copy(f"COPY {Galaxy.__tablename__} (hpx) FROM STDIN") as copy: + copy.write("\n".join(f"{i}" for i in hpx)) return points @@ -86,14 +85,15 @@ def get_random_fields(n, cursor): footprints = get_footprints_grid(*get_ztf_footprint_corners(), centers) mocs = [MOC.from_polygon_skycoord(footprint) for footprint in footprints] - with cursor.copy(f'COPY {Field.__tablename__} FROM STDIN') as copy: - copy.write('\n'.join(f'{i}' for i in range(len(mocs)))) + with cursor.copy(f"COPY {Field.__tablename__} FROM STDIN") as copy: + copy.write("\n".join(f"{i}" for i in range(len(mocs)))) - with cursor.copy(f'COPY {FieldTile.__tablename__} FROM STDIN') as copy: + with cursor.copy(f"COPY {FieldTile.__tablename__} FROM STDIN") as copy: copy.write( - '\n'.join( - f'{i}\t{hpx}' - for i, moc in enumerate(mocs) for hpx in Tile.tiles_from(moc) + "\n".join( + f"{i}\t{hpx}" + for i, moc in enumerate(mocs) + for hpx in Tile.tiles_from(moc) ) ) @@ -104,7 +104,7 @@ 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() + 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] @@ -119,13 +119,13 @@ def get_random_sky_map(n, cursor): probdensity = rng.uniform(0, 1, size=len(tiles) - 1) probdensity /= np.sum(np.diff(tiles) * probdensity) * PIXEL_AREA - with cursor.copy(f'COPY {Skymap.__tablename__} FROM STDIN') as copy: - copy.write('1') + with cursor.copy(f"COPY {Skymap.__tablename__} FROM STDIN") as copy: + copy.write("1") - with cursor.copy(f'COPY {SkymapTile.__tablename__} FROM STDIN') as copy: + with cursor.copy(f"COPY {SkymapTile.__tablename__} FROM STDIN") as copy: copy.write( - '\n'.join( - f'1\t[{lo},{hi})\t{p}' + "\n".join( + f"1\t[{lo},{hi})\t{p}" for lo, hi, p in zip(tiles[:-1], tiles[1:], probdensity) ) ) diff --git a/tests/benchmarks/models.py b/tests/benchmarks/models.py index 8ad29e8..7937aa6 100644 --- a/tests/benchmarks/models.py +++ b/tests/benchmarks/models.py @@ -1,4 +1,5 @@ """SQLAlchemy ORM models for unit tests.""" + import sqlalchemy as sa from sqlalchemy import orm @@ -7,7 +8,6 @@ @orm.as_declarative() class Base: - @orm.declared_attr def __tablename__(cls): return cls.__name__.lower() diff --git a/tests/benchmarks/test_benchmarks.py b/tests/benchmarks/test_benchmarks.py index 8d5af3f..bf5b459 100644 --- a/tests/benchmarks/test_benchmarks.py +++ b/tests/benchmarks/test_benchmarks.py @@ -1,19 +1,18 @@ from functools import reduce import numpy as np -import sqlalchemy as sa import pytest +import sqlalchemy as sa from healpix_alchemy import func -from .models import Galaxy, FieldTile, SkymapTile +from .models import FieldTile, Galaxy, SkymapTile @pytest.fixture def bench(benchmark, session): - def _func(query): - session.execute(sa.text('ANALYZE')) + session.execute(sa.text("ANALYZE")) return benchmark(lambda: session.execute(query).all()) return _func @@ -21,7 +20,6 @@ def _func(query): @pytest.fixture def bench_and_check(bench): - def _func(query, expected): np.testing.assert_almost_equal(bench(query), expected, decimal=6) @@ -31,12 +29,8 @@ def _func(query, expected): def test_union_area(bench_and_check, random_fields): """Find the area of the union of N fields.""" # Assemble query - subquery = sa.select( - func.union(FieldTile.hpx).label('hpx') - ).subquery() - query = sa.select( - sa.func.sum(subquery.columns.hpx.area) - ) + subquery = sa.select(func.union(FieldTile.hpx).label("hpx")).subquery() + query = sa.select(sa.func.sum(subquery.columns.hpx.area)) # Expected result union = reduce(lambda a, b: a.union(b), random_fields) @@ -47,29 +41,24 @@ def test_union_area(bench_and_check, random_fields): bench_and_check(query, expected) -def test_crossmatch_galaxies_and_fields(bench_and_check, - random_fields, random_galaxies): +def test_crossmatch_galaxies_and_fields( + bench_and_check, random_fields, random_galaxies +): """Cross match N galaxies with M fields.""" # Assemble query count = sa.func.count(Galaxy.id) - query = sa.select( - count - ).filter( - FieldTile.hpx.contains(Galaxy.hpx) - ).group_by( - FieldTile.id - ).order_by( - count.desc() - ).limit( - 5 + query = ( + sa.select(count) + .filter(FieldTile.hpx.contains(Galaxy.hpx)) + .group_by(FieldTile.id) + .order_by(count.desc()) + .limit(5) ) # Expected result points = random_galaxies fields = random_fields - result = np.sum( - [moc.contains_skycoords(points) for moc in fields], - axis=1) + result = np.sum([moc.contains_skycoords(points) for moc in fields], axis=1) expected = np.flipud(np.sort(result))[:5].reshape(-1, 1) # Run benchmark, check result @@ -79,29 +68,24 @@ def test_crossmatch_galaxies_and_fields(bench_and_check, def test_fields_in_90pct_credible_region(bench, random_fields, random_sky_map): """Find which of N fields overlap the 90% credible region.""" # Assemble query - cum_prob = sa.func.sum( - SkymapTile.probdensity * SkymapTile.hpx.area - ).over( - order_by=SkymapTile.probdensity.desc() - ).label( - 'cum_prob' + cum_prob = ( + sa.func.sum(SkymapTile.probdensity * SkymapTile.hpx.area) + .over(order_by=SkymapTile.probdensity.desc()) + .label("cum_prob") + ) + subquery1 = ( + sa.select(SkymapTile.probdensity, cum_prob) + .filter(SkymapTile.id == 1) + .subquery() + ) + min_probdensity = ( + sa.select(sa.func.min(subquery1.columns.probdensity)) + .filter(subquery1.columns.cum_prob <= 0.9) + .scalar_subquery() ) - subquery1 = sa.select( - SkymapTile.probdensity, - cum_prob - ).filter( - SkymapTile.id == 1 - ).subquery() - min_probdensity = sa.select( - sa.func.min(subquery1.columns.probdensity) - ).filter( - subquery1.columns.cum_prob <= 0.9 - ).scalar_subquery() - query = sa.select( - sa.func.count(FieldTile.id.distinct()) - ).filter( + query = sa.select(sa.func.count(FieldTile.id.distinct())).filter( SkymapTile.hpx.overlaps(FieldTile.hpx), - SkymapTile.probdensity >= min_probdensity + SkymapTile.probdensity >= min_probdensity, ) # Run benchmark