from __future__ import annotations
import json
import math
from pathlib import Path
import numpy as np
_FLAT_PAYLOAD_KINDS = {
"tensor",
"tensor_2d",
"matrix",
"vector_list",
"grid",
"grid_periodic",
"grid_clamped",
"square_lattice_periodic",
"square_lattice_reflective",
}
_SPARSE_PAYLOAD_KINDS = {"tensor_sparse", "matrix_sparse"}
_JSON_SCHEMA_VERSION = 1
def to_ndarray(path) -> np.ndarray:
payload = json.loads(Path(path).read_text())
return _payload_to_ndarray(payload)
def _payload_to_ndarray(payload) -> np.ndarray:
if isinstance(payload, dict):
kind = payload.get("kind")
if kind in _SPARSE_PAYLOAD_KINDS:
return _sparse_payload_to_ndarray(payload)
if kind in _FLAT_PAYLOAD_KINDS and {"shape", "data"}.issubset(payload.keys()):
return _flat_payload_to_ndarray(payload)
if kind == "tensor" and "storage" in payload:
return _legacy_tensor_payload_to_ndarray(payload)
if kind == "vector_list" and "storage" in payload:
return _legacy_vector_list_payload_to_ndarray(payload)
if kind == "grid" and "storage" in payload:
return _legacy_grid_payload_to_ndarray(payload)
if _looks_like_legacy_compact_grid(payload):
return _legacy_compact_grid_payload_to_ndarray(payload)
return _object_scalar_array(_json_to_python(payload))
if isinstance(payload, list):
return np.asarray([_json_to_python(item) for item in payload], dtype=object)
return np.asarray(payload)
def _flat_payload_to_ndarray(payload) -> np.ndarray:
_require_keys(payload, {"version", "scalar", "shape", "data"})
_require_current_version(payload)
shape = _normalize_shape(payload["shape"])
array = _numeric_values(payload["data"], payload["scalar"])
expected_size = math.prod(shape)
if array.size != expected_size:
raise ValueError(
f"flat payload data length mismatch: expected {expected_size}, got {array.size}"
)
return array.reshape(shape)
def _sparse_payload_to_ndarray(payload) -> np.ndarray:
_require_keys(payload, {"version", "scalar", "shape", "indices", "values"})
_require_current_version(payload)
shape = _normalize_shape(payload["shape"])
indices = payload["indices"]
values = payload["values"]
if not isinstance(indices, list) or not isinstance(values, list):
raise ValueError("sparse PiP indices and values must be arrays")
if len(indices) != len(values):
raise ValueError(
f"sparse PiP entry length mismatch: {len(indices)} indices, {len(values)} values"
)
logical_size = math.prod(shape)
previous = -1
for index in indices:
if not isinstance(index, int) or isinstance(index, bool):
raise ValueError(f"invalid sparse PiP index: {index!r}")
if index <= previous:
raise ValueError("sparse PiP indices must be strictly increasing")
if index >= logical_size:
raise ValueError(
f"sparse PiP index out of bounds: {index} >= {logical_size}"
)
previous = index
numeric_values = _numeric_values(values, payload["scalar"])
if np.any(numeric_values == 0):
raise ValueError("sparse PiP values must not contain explicit zeros")
array = np.zeros(logical_size, dtype=numeric_values.dtype)
if indices:
array[np.asarray(indices, dtype=np.intp)] = numeric_values
return array.reshape(shape)
def _numeric_values(values, scalar) -> np.ndarray:
if not isinstance(values, list):
raise ValueError("PiP numeric payload values must be an array")
if scalar in {"complex_f32", "complex_f64"}:
converted = []
for value in values:
if isinstance(value, dict) and set(value) == {"re", "im"}:
real, imaginary = value["re"], value["im"]
elif isinstance(value, list) and len(value) == 2:
real, imaginary = value
else:
raise ValueError(f"invalid PiP complex scalar: {value!r}")
if not _is_finite_number(real) or not _is_finite_number(imaginary):
raise ValueError(f"invalid PiP complex scalar: {value!r}")
converted.append(complex(real, imaginary))
dtype = np.complex64 if scalar == "complex_f32" else np.complex128
array = np.asarray(converted, dtype=dtype)
if not np.all(np.isfinite(array)):
raise ValueError(f"PiP {scalar} values overflow their declared dtype")
return array
integer_dtypes = {
"i8": np.int8,
"i16": np.int16,
"i32": np.int32,
"i64": np.int64,
"isize": np.intp,
"u8": np.uint8,
"u16": np.uint16,
"u32": np.uint32,
"u64": np.uint64,
"usize": np.uintp,
}
if scalar in integer_dtypes:
dtype = integer_dtypes[scalar]
_validate_integer_values(values, scalar, np.iinfo(dtype).min, np.iinfo(dtype).max)
return np.asarray(values, dtype=dtype)
if scalar in {"i128", "u128"}:
minimum = -(1 << 127) if scalar == "i128" else 0
maximum = (1 << 127) - 1 if scalar == "i128" else (1 << 128) - 1
_validate_integer_values(values, scalar, minimum, maximum)
return np.asarray(values, dtype=object)
float_dtypes = {
"f32": np.float32,
"f64": np.float64,
}
try:
dtype = float_dtypes[scalar]
except KeyError as error:
raise ValueError(f"unsupported PiP scalar kind: {scalar!r}") from error
if not all(_is_finite_number(value) for value in values):
raise ValueError(f"PiP {scalar} values must contain only finite numbers")
array = np.asarray(values, dtype=dtype)
if not np.all(np.isfinite(array)):
raise ValueError(f"PiP {scalar} values overflow their declared dtype")
return array
def _validate_integer_values(values, scalar, minimum, maximum) -> None:
for value in values:
if (
not isinstance(value, int)
or isinstance(value, bool)
or value < minimum
or value > maximum
):
raise ValueError(f"invalid PiP {scalar} scalar: {value!r}")
def _is_finite_number(value) -> bool:
if not isinstance(value, (int, float)) or isinstance(value, bool):
return False
try:
return math.isfinite(value)
except OverflowError:
return False
def _legacy_tensor_payload_to_ndarray(payload) -> np.ndarray:
_require_keys(payload, {"shape", "storage", "data"})
shape = _normalize_shape(payload["shape"])
storage = payload["storage"]
if storage == "dense":
array = np.asarray(payload["data"])
return array.reshape(shape)
if storage == "sparse":
entries = payload["data"]["entries"]
dtype = _numpy_dtype_from_scalar_type(payload.get("scalar_type"))
array = np.zeros(shape, dtype=dtype)
for entry in entries:
array.flat[entry["index"]] = entry["value"]
return array
raise ValueError(f"unsupported legacy PiP tensor storage: {storage!r}")
def _legacy_vector_list_payload_to_ndarray(payload) -> np.ndarray:
_require_keys(payload, {"shape", "data"})
dim, n = _normalize_shape(payload["shape"])
array = np.asarray(payload["data"])
if array.shape != (n, dim):
raise ValueError(
f"legacy vector_list payload data shape mismatch: expected {(n, dim)}, got {array.shape}"
)
return array.T
def _legacy_grid_payload_to_ndarray(payload) -> np.ndarray:
_require_keys(payload, {"shape", "data"})
return _reshape_legacy_grid_data(payload["shape"], payload["data"])
def _legacy_compact_grid_payload_to_ndarray(payload) -> np.ndarray:
return _reshape_legacy_grid_data(payload["shape"], payload["data"])
def _reshape_legacy_grid_data(shape_metadata, data) -> np.ndarray:
d, l = _normalize_shape(shape_metadata)
array = np.asarray(data)
return array.reshape((l,) * d)
def _looks_like_legacy_compact_grid(payload) -> bool:
if set(payload.keys()) != {"shape", "data"}:
return False
shape = payload["shape"]
return (
isinstance(shape, list)
and len(shape) == 2
and all(isinstance(value, int) for value in shape)
)
def _normalize_shape(shape) -> tuple[int, ...]:
if (
not isinstance(shape, list)
or not shape
or not all(isinstance(dim, int) and not isinstance(dim, bool) for dim in shape)
):
raise ValueError(f"invalid PiP shape metadata: {shape!r}")
if any(dim <= 0 for dim in shape):
raise ValueError(f"PiP shape dimensions must be > 0: {shape!r}")
return tuple(shape)
def _require_keys(payload, required_keys) -> None:
missing = required_keys.difference(payload.keys())
if missing:
raise ValueError(f"missing required PiP payload keys: {sorted(missing)}")
def _require_current_version(payload) -> None:
version = payload.get("version")
if version != _JSON_SCHEMA_VERSION:
raise ValueError(
f"unsupported PiP JSON schema version: expected {_JSON_SCHEMA_VERSION}, got {version!r}"
)
def _numpy_dtype_from_scalar_type(scalar_type):
if scalar_type in {"f32"}:
return np.float32
if scalar_type in {"f64"}:
return np.float64
if scalar_type in {"i8"}:
return np.int8
if scalar_type in {"i16"}:
return np.int16
if scalar_type in {"i32"}:
return np.int32
if scalar_type in {"i64", "isize"}:
return np.int64
if scalar_type in {"u8"}:
return np.uint8
if scalar_type in {"u16"}:
return np.uint16
if scalar_type in {"u32"}:
return np.uint32
if scalar_type in {"u64", "usize"}:
return np.uint64
if isinstance(scalar_type, str) and "Complex<f32>" in scalar_type:
return np.complex64
if isinstance(scalar_type, str) and "Complex<f64>" in scalar_type:
return np.complex128
return np.float64
def _json_to_python(value):
if isinstance(value, dict):
return {key: _json_to_python(item) for key, item in value.items()}
if isinstance(value, list):
return [_json_to_python(item) for item in value]
return value
def _object_scalar_array(value) -> np.ndarray:
array = np.empty((), dtype=object)
array[()] = value
return array