Skip to content

Commit

Permalink
Do not use numpy.ndarray for sqlite3 converters and adapters on compl…
Browse files Browse the repository at this point in the history
…ex scalars
  • Loading branch information
Florian Blanchet committed Sep 21, 2022
1 parent 02d2bd5 commit 82b2743
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 11 deletions.
52 changes: 42 additions & 10 deletions qcodes/dataset/sqlite/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,13 @@
perform_db_upgrade,
)
from qcodes.dataset.sqlite.initial_schema import init_db
from qcodes.utils.types import complex_types, numpy_floats, numpy_ints
from qcodes.utils.types import (
complex_type_union,
complex_types,
numpy_complex_map_size2type,
numpy_floats,
numpy_ints,
)

JournalMode = Literal["DELETE", "TRUNCATE", "PERSIST", "MEMORY", "WAL", "OFF"]

Expand Down Expand Up @@ -54,10 +60,32 @@ def _convert_array(text: bytes) -> np.ndarray:
return np.lib.format.read_array(io.BytesIO(text), allow_pickle=False)


def _convert_complex(text: bytes) -> np.complexfloating:
out = io.BytesIO(text)
out.seek(0)
return np.load(out)[0]
def _convert_complex(text: bytes) -> complex_type_union:
"""
See this:
https://numpy.org/devdocs/reference/generated/numpy.lib.format.html#format-version-3-0
np.load and np.lib.format.read_array parse the npy header
Skipping the header is 70 times faster
For backward compatibility with qcodes <=0.34:
npy format has an header len evenly divisible by 64, and QCoDeS saved complexes as
single value arrays
HEADER HEADER binary value
|--- n x 64 --||- len % 64-|
For qcodes >0.34:
QCoDeS saves complexes directly as binary value: there is no header
NOTE: Will break with complex512 whose lenght is 64 bytes
"""
try:
value_size = len(text) % 64
return np.frombuffer(
text[-value_size:], dtype=numpy_complex_map_size2type[value_size]
).item()
except KeyError as exc:
raise ValueError(f"Cannot parse {str(text)}") from exc


this_session_default_encoding = sys.getdefaultencoding()
Expand Down Expand Up @@ -108,11 +136,15 @@ def _adapt_float(fl: float) -> float | str:
return float(fl)


def _adapt_complex(value: complex | np.complexfloating) -> sqlite3.Binary:
out = io.BytesIO()
np.save(out, np.array([value]))
out.seek(0)
return sqlite3.Binary(out.read())
def _adapt_complex(value: complex_type_union) -> sqlite3.Binary:
# np.save and np.lib.format.write_array add an header that is useless for a single
# value
# Avoiding the header is 15 times faster.
return sqlite3.Binary(
(
value if isinstance(value, np.complexfloating) else np.complex_(value)
).tobytes()
)


def connect(name: str | Path, debug: bool = False, version: int = -1) -> ConnectionPlus:
Expand Down
15 changes: 14 additions & 1 deletion qcodes/tests/dataset/test_dataset_basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,11 @@
from qcodes.dataset.descriptions.rundescriber import RunDescriber
from qcodes.dataset.guids import parse_guid
from qcodes.dataset.sqlite.connection import atomic, path_to_dbfile
from qcodes.dataset.sqlite.database import _convert_array, get_DB_location
from qcodes.dataset.sqlite.database import (
_convert_array,
_convert_complex,
get_DB_location,
)
from qcodes.dataset.sqlite.queries import _rewrite_timestamps, _unicode_categories
from qcodes.tests.common import error_caused_by
from qcodes.tests.dataset.helper_functions import verify_data_dict
Expand Down Expand Up @@ -615,6 +619,15 @@ def test_backward_compat__adapt_array_v0_33():
assert arr == _convert_array(out.read())


def test_backward_compat__adapt_complex_v0_33():
for dtype in complex_types:
value = dtype(1j)
out = io.BytesIO()
np.save(out, np.array([value]))
out.seek(0)
assert value == _convert_complex(out.read())


def test_missing_keys(dataset):
"""
Test that we can now have partial results with keys missing. This is for
Expand Down
1 change: 1 addition & 0 deletions qcodes/utils/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
"""

numpy_concrete_complex = (np.complex64, np.complex128)
numpy_complex_map_size2type = {np.dtype(t).itemsize: t for t in numpy_concrete_complex}
"""
Complex types with fixed sizes.
"""
Expand Down

0 comments on commit 82b2743

Please sign in to comment.