Python rewrite

.gitattributes

@@ -0,0 +1 @@
+testdata/leveldbs/** filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,3 +1 @@
-alerts.db/
-build/
-.gobincache
+virtualenv

LICENSE

@@ -0,0 +1,11 @@
+Copyright 2021 Spencer Nelson <[email protected]>
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

Makefile

@@ -0,0 +1,55 @@
+virtualenv:
+	python -m venv virtualenv
+
+VENV = virtualenv/bin
+$(VENV): virtualenv
+$(VENV)/pip: $(VENV)
+
+.git/hooks/pre-commit:
+	ln -sf ../../devconfig/pre-commit.sh .git/hooks/pre-commit
+
+
+# Development tools
+.PHONY: dev-setup
+dev-setup: .git/hooks/pre-commit $(VENV)/flake8 $(VENV)/black $(VENV)/mypy
+deps: $(VENV)
+	$(VENV)/pip install -e '.[dev]'
+
+$(VENV)/flake8 $(VENV)/black $(VENV)/mypy $(VENV)/pytest &: $(VENV)/pip
+	$(VENV)/pip install -e '.[dev]'
+
+.PHONY: lint check-format format typecheck precommit test
+lint: $(VENV)/flake8
+	$(VENV)/flake8 ./src
+	$(VENV)/flake8 ./test
+
+check-format: $(VENV)/black
+	$(VENV)/black \
+		--target-version py38 \
+		--check \
+		./src/alertbase
+	$(VENV)/black \
+		--target-version py38 \
+		--check \
+		./test
+
+format: $(VENV)/black
+	$(VENV)/black \
+		--target-version py38 \
+		./src/alertbase
+	$(VENV)/black \
+		--target-version py38 \
+		./test
+
+typecheck: $(VENV)/mypy
+	$(VENV)/mypy ./src
+
+test: $(VENV)/pytest
+	$(VENV)/pytest .
+
+precommit: check-format typecheck lint test
+
+
+.PHONY: clean
+clean:
+	rm -rf virtualenv

README.md

@@ -0,0 +1,15 @@
+# Atlas
+
+## Developing
+
+To start, you'll need Python 3.8.
+
+Install [`git-lfs`](https://git-lfs.github.com/) so that you can download test
+datasets.
+
+Run `make dev-setup`.
+
+You're good to go!
+
+## Name
+Atlas: The **A**tlas **T**wo-**L**evel **A**lert **S**earch tool.

devconfig/pre-commit.sh

@@ -0,0 +1,3 @@
+#!/bin/sh
+
+make precommit

go/.gitignore

@@ -0,0 +1,3 @@
+alerts.db/
+build/
+.gobincache

pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"

setup.cfg

@@ -0,0 +1,42 @@
+[metadata]
+name = alertbase
+version = 0.0.1-alpha
+description = ZTF Alert Database
+long_description = file: README.md
+license = BSD 3-Clause License
+license_file = LICENSE
+classifiers =
+    License :: OSI Approved :: BSD License
+    Programming Language :: Python :: 3
+author = Spencer Nelson
+author_email = [email protected]
+
+[options]
+package_dir=
+    =src
+packages=find:
+install_requires =
+    healpy==1.14.0
+    plyvel==1.3.0
+
+[options.extras_require]
+dev = flake8; black; mypy; pytest
+
+[options.packages.find]
+where=src
+
+[options.package_data]
+alertbase = py.typed
+
+[mypy]
+python_version = 3.8
+warn_return_any = True
+
+[mypy-plyvel]
+ignore_missing_imports = True
+
+[mypy-astropy.*]
+ignore_missing_imports = True
+
+[flake8]
+max-line-length = 88

setup.py

@@ -0,0 +1,2 @@
+import setuptools
+setuptools.setup()

src/alertbase/__init__.py

@@ -0,0 +1,3 @@
+from alertbase.db import open_db
+
+__all__ = ["open_db"]

src/alertbase/db.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+from typing import Iterator
+
+import pathlib
+import plyvel
+import astropy.coordinates
+
+
+def open_db(path: str) -> Database:
+    return Database(path)
+
+
+class Database:
+    db_root: pathlib.Path
+    objects: plyvel.DB
+    candidates: plyvel.DB
+    healpixels: plyvel.DB
+    timestamps: plyvel.DB
+
+    def __init__(self, db_path: str):
+        self.db_root = pathlib.Path(db_path)
+        self.objects = plyvel.DB(str(self.db_root / "objects"))
+        self.candidates = plyvel.DB(str(self.db_root / "candidates"))
+        self.healpixels = plyvel.DB(str(self.db_root / "healpixels"))
+        self.timestamps = plyvel.DB(str(self.db_root / "timestamps"))
+
+    def cone_search(self, center: astropy.coordinates.SkyCoord) -> Iterator[str]:
+        pass
+
+    def count_objects(self) -> int:
+        """count_objects iterates over all the objects in the database to count how
+        many there are.
+
+        """
+        return sum(1 for _ in self.objects.iterator())
+
+    def count_candidates(self) -> int:
+        """count_candidates iterates over all the candidates in the database to count
+        how many there are.
+
+        """
+        return sum(1 for _ in self.candidates.iterator())
+
+    def count_healpixels(self) -> int:
+        """count_candidates iterates over all the HEALPix pixels in the database to
+        count how many have data.
+
+        """
+        return sum(1 for _ in self.healpixels.iterator())
+
+    def count_timestamps(self) -> int:
+        """count_timestamps iterates over all the HEALPix pixels in the database to
+        count how many unique timestamps have data.
+
+        """
+        return sum(1 for _ in self.timestamps.iterator())

src/alertbase/encoding.py

@@ -0,0 +1,123 @@
+from typing import List, Tuple, Iterator
+import struct
+
+
+# This file provides a variety of utilities for converting between python data
+# types and byte arrays.
+
+
+def pack_uint64(data: int) -> bytes:
+    """Pack an integer as a fixed-size 64-bit unsigned integer. This is more
+    efficient (both in space and compute) than pack_varint for large integers."""
+    return struct.pack(">Q", data)
+
+
+def pack_uint64s(data: List[int]) -> bytes:
+    """Pack a sequence of integers using pack_uint64. """
+    result = b""
+    for i in data:
+        result += pack_uint64(i)
+    return result
+
+
+def unpack_uint64s(data: bytes) -> List[int]:
+    """ Unpack a series of integers that were packed with pack_uint64."""
+    if len(data) == 0:
+        return []
+    return list(x[0] for x in struct.iter_unpack(">Q", data))
+
+
+def pack_jd_timestamp(jd: float) -> bytes:
+    """Pack a julian date as a unix nanosecond timestamp (doesn't attempt to handle
+    leap seconds)
+    """
+    return pack_uint64(int((jd - 2440587.5) * 86400000000000))
+
+
+def _pack_uvarint(n: int) -> bytes:
+    """Pack an unsigned variable-length integer into bytes. """
+    result = b""
+    while True:
+        chunk = n & 0x7F
+        n >>= 7
+        if n:
+            result += bytes((chunk | 0x80,))
+        else:
+            result += bytes((chunk,))
+            break
+    return result
+
+
+def pack_varint(n):
+    """Pack a zig-zag encoded, signed integer into bytes."""
+    return _pack_uvarint(_zigzag_encode(n))
+
+
+def pack_varint_list(data: List[int]) -> bytes:
+    """Pack a series of integers into bytes using pack_varint. """
+    result = b""
+    for value in data:
+        result += pack_varint(value)
+    return result
+
+
+def _zigzag_encode(x):
+    if x >= 0:
+        return x << 1
+    return (x << 1) ^ (~0)
+
+
+def _zigzag_decode(x):
+    if not x & 0x1:
+        return x >> 1
+    return (x >> 1) ^ (~0)
+
+
+def _unpack_uvarint(data: bytes) -> Tuple[int, int]:
+    """Unpacks a variable-length integer stored in given byte buffer.
+
+    Returns the integer and the number of bytes that were read."""
+    shift = 0
+    result = 0
+    n = 0
+    for b in data:
+        n += 1
+        result |= (b & 0x7F) << shift
+        if not (b & 0x80):
+            break
+        shift += 7
+    return result, n
+
+
+def unpack_varint(data: bytes) -> Tuple[int, int]:
+    """Unpacks a variable-length, zig-zag-encoded integer from a given byte buffer.
+
+    Returns the integer and the number of bytes that were read.
+    """
+    result, n = _unpack_uvarint(data)
+    return _zigzag_decode(result), n
+
+
+def unpack_varint_list(data: bytes) -> List[int]:
+    """Calls unpack_varint repeatedly on data, returning the complete list of all
+    integers encoded therein.
+
+    """
+    result = []
+    pos = 0
+    while pos < len(data):
+        val, n_read = unpack_varint(data[pos:])
+        pos += n_read
+        result.append(val)
+    return result
+
+
+def iter_varints(data: bytes) -> Iterator[int]:
+    """Calls unpack_varint repeatedly on data, iterating over the integers encoded
+    therein.
+    """
+    pos = 0
+    while pos < len(data):
+        val, n_read = unpack_varint(data[pos:])
+        pos += n_read
+        yield val

test/test_db.py

@@ -0,0 +1,48 @@
+import pytest
+import shutil
+import tempfile
+import pathlib
+
+import alertbase
+
+
+@pytest.fixture(scope="function")
+def leveldb_5k():
+    """Copy the testdata/leveldbs/alerts.db.5k database to a temporary directory,
+    scoped to a single test invocation.
+
+    """
+    db_path = "testdata/leveldbs/alerts.db.5k"
+    with tempfile.TemporaryDirectory(prefix="test-alerts-5k-") as tmp_dir:
+        tmp_db = pathlib.Path(tmp_dir) / "alerts.db.5k"
+        shutil.copytree(db_path, tmp_db)
+        yield tmp_db
+
+
+class TestDatabase:
+    def test_open_database(self, leveldb_5k):
+        alertbase.open_db(leveldb_5k)
+
+    def test_count_candidates(self, leveldb_5k):
+        db = alertbase.open_db(leveldb_5k)
+        n = db.count_candidates()
+        assert n == 5000
+
+    def test_count_objects(self, leveldb_5k):
+        db = alertbase.open_db(leveldb_5k)
+        n = db.count_objects()
+        assert n == 4848
+
+    def test_count_timestamps(self, leveldb_5k):
+        db = alertbase.open_db(leveldb_5k)
+        n = db.count_timestamps()
+        assert n == 11
+
+    def test_count_healpixels(self, leveldb_5k):
+        db = alertbase.open_db(leveldb_5k)
+        n = db.count_healpixels()
+        assert n == 4216
+
+    def test_open_missing_db(self):
+        with pytest.raises(Exception):
+            alertbase.open_db("bogus")