524 lines
17 KiB
Python
524 lines
17 KiB
Python
"""
|
|
Lazy GDSII readers and writers backed by native Arrow scan/materialize paths.
|
|
|
|
This module is intentionally separate from `gdsii.arrow` so the eager read path
|
|
keeps its current behavior and performance profile.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import IO, TYPE_CHECKING, Any, cast
|
|
from collections import defaultdict
|
|
import gzip
|
|
import logging
|
|
import mmap
|
|
import pathlib
|
|
|
|
import numpy
|
|
|
|
from . import arrow
|
|
from .lazy_write import write as write, writefile as writefile
|
|
from ..utils import is_gzipped
|
|
from ...library import (
|
|
ILibraryView,
|
|
LibraryView,
|
|
PortsLibraryView,
|
|
dangling_mode_t,
|
|
)
|
|
from ...utils import apply_transforms
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterator, Sequence
|
|
|
|
from numpy.typing import NDArray
|
|
import pyarrow
|
|
|
|
from ...pattern import Pattern
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _StructRange:
|
|
start: int
|
|
end: int
|
|
|
|
|
|
@dataclass
|
|
class _SourceBuffer:
|
|
path: pathlib.Path
|
|
data: bytes | mmap.mmap
|
|
handle: IO[bytes] | None = None
|
|
|
|
def raw_slice(self, start: int, end: int) -> bytes:
|
|
return self.data[start:end]
|
|
|
|
|
|
@dataclass
|
|
class _ScanRefs:
|
|
offsets: NDArray[numpy.integer[Any]]
|
|
targets: NDArray[numpy.integer[Any]]
|
|
xy: NDArray[numpy.int32]
|
|
xy0: NDArray[numpy.int32]
|
|
xy1: NDArray[numpy.int32]
|
|
counts: NDArray[numpy.int64]
|
|
invert_y: NDArray[numpy.bool_ | numpy.bool]
|
|
angle_rad: NDArray[numpy.floating[Any]]
|
|
scale: NDArray[numpy.floating[Any]]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _CellScan:
|
|
cell_id: int
|
|
struct_range: _StructRange
|
|
ref_start: int
|
|
ref_stop: int
|
|
children: set[str]
|
|
|
|
|
|
@dataclass
|
|
class _ScanPayload:
|
|
libarr: pyarrow.StructScalar
|
|
library_info: dict[str, Any]
|
|
cell_names: list[str]
|
|
cell_order: list[str]
|
|
cells: dict[str, _CellScan]
|
|
refs: _ScanRefs
|
|
|
|
def is_available() -> bool:
|
|
return arrow.is_available()
|
|
|
|
|
|
def _open_source_buffer(path: pathlib.Path) -> _SourceBuffer:
|
|
if is_gzipped(path):
|
|
with gzip.open(path, mode='rb') as stream:
|
|
data = stream.read()
|
|
return _SourceBuffer(path=path, data=data)
|
|
|
|
handle = path.open(mode='rb', buffering=0)
|
|
mapped = mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ)
|
|
return _SourceBuffer(path=path, data=mapped, handle=handle)
|
|
|
|
|
|
def _extract_scan_payload(libarr: pyarrow.StructScalar) -> _ScanPayload:
|
|
library_info = arrow._read_header(libarr)
|
|
cell_names = libarr['cell_names'].as_py()
|
|
|
|
cells = libarr['cells']
|
|
cell_values = cells.values
|
|
cell_ids = cell_values.field('id').to_numpy()
|
|
struct_starts = cell_values.field('struct_start_offset').to_numpy()
|
|
struct_ends = cell_values.field('struct_end_offset').to_numpy()
|
|
|
|
refs = cell_values.field('refs')
|
|
ref_values = refs.values
|
|
ref_offsets = refs.offsets.to_numpy()
|
|
targets = ref_values.field('target').to_numpy()
|
|
xy = arrow._packed_xy_u64_to_pairs(ref_values.field('xy').to_numpy())
|
|
xy0 = arrow._packed_xy_u64_to_pairs(ref_values.field('xy0').to_numpy())
|
|
xy1 = arrow._packed_xy_u64_to_pairs(ref_values.field('xy1').to_numpy())
|
|
counts = arrow._packed_counts_u32_to_pairs(ref_values.field('counts').to_numpy())
|
|
invert_y = ref_values.field('invert_y').to_numpy(zero_copy_only=False)
|
|
angle_rad = ref_values.field('angle_rad').to_numpy()
|
|
scale = ref_values.field('scale').to_numpy()
|
|
|
|
ref_payload = _ScanRefs(
|
|
offsets=ref_offsets,
|
|
targets=targets,
|
|
xy=xy,
|
|
xy0=xy0,
|
|
xy1=xy1,
|
|
counts=counts,
|
|
invert_y=invert_y,
|
|
angle_rad=angle_rad,
|
|
scale=scale,
|
|
)
|
|
|
|
cell_order = [cell_names[int(cell_id)] for cell_id in cell_ids]
|
|
cell_scan: dict[str, _CellScan] = {}
|
|
for cc, name in enumerate(cell_order):
|
|
ref_start = int(ref_offsets[cc])
|
|
ref_stop = int(ref_offsets[cc + 1])
|
|
children = {
|
|
cell_names[int(target)]
|
|
for target in targets[ref_start:ref_stop]
|
|
}
|
|
cell_scan[name] = _CellScan(
|
|
cell_id=int(cell_ids[cc]),
|
|
struct_range=_StructRange(int(struct_starts[cc]), int(struct_ends[cc])),
|
|
ref_start=ref_start,
|
|
ref_stop=ref_stop,
|
|
children=children,
|
|
)
|
|
|
|
return _ScanPayload(
|
|
libarr=libarr,
|
|
library_info=library_info,
|
|
cell_names=cell_names,
|
|
cell_order=cell_order,
|
|
cells=cell_scan,
|
|
refs=ref_payload,
|
|
)
|
|
|
|
def _make_ref_rows(
|
|
xy: NDArray[numpy.integer[Any]],
|
|
angle_rad: NDArray[numpy.floating[Any]],
|
|
invert_y: NDArray[numpy.bool_ | numpy.bool],
|
|
scale: NDArray[numpy.floating[Any]],
|
|
) -> NDArray[numpy.float64]:
|
|
rows = numpy.empty((len(xy), 5), dtype=float)
|
|
rows[:, :2] = xy
|
|
rows[:, 2] = angle_rad
|
|
rows[:, 3] = invert_y.astype(float)
|
|
rows[:, 4] = scale
|
|
return rows
|
|
|
|
|
|
def _expand_aref_row(
|
|
xy: NDArray[numpy.integer[Any]],
|
|
xy0: NDArray[numpy.integer[Any]],
|
|
xy1: NDArray[numpy.integer[Any]],
|
|
counts: NDArray[numpy.integer[Any]],
|
|
angle_rad: float,
|
|
invert_y: bool,
|
|
scale: float,
|
|
) -> NDArray[numpy.float64]:
|
|
a_count = int(counts[0])
|
|
b_count = int(counts[1])
|
|
aa, bb = numpy.meshgrid(numpy.arange(a_count), numpy.arange(b_count), indexing='ij')
|
|
displacements = aa.reshape(-1, 1) * xy0[None, :] + bb.reshape(-1, 1) * xy1[None, :]
|
|
rows = numpy.empty((displacements.shape[0], 5), dtype=float)
|
|
rows[:, :2] = xy + displacements
|
|
rows[:, 2] = angle_rad
|
|
rows[:, 3] = float(invert_y)
|
|
rows[:, 4] = scale
|
|
return rows
|
|
|
|
|
|
class ArrowLibrary(ILibraryView):
|
|
"""
|
|
Read-only library backed by the native lazy Arrow scan schema.
|
|
|
|
Materializing a cell via `__getitem__` caches a real `Pattern` for that cell.
|
|
Cached cells are treated as edited for future writes from this module.
|
|
"""
|
|
|
|
path: pathlib.Path
|
|
library_info: dict[str, Any]
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
path: pathlib.Path,
|
|
payload: _ScanPayload,
|
|
source: _SourceBuffer,
|
|
) -> None:
|
|
self.path = path
|
|
self.library_info = payload.library_info
|
|
self._payload = payload
|
|
self._source = source
|
|
self._cache: dict[str, Pattern] = {}
|
|
|
|
@classmethod
|
|
def from_file(cls, filename: str | pathlib.Path) -> ArrowLibrary:
|
|
path = pathlib.Path(filename).expanduser().resolve()
|
|
source = _open_source_buffer(path)
|
|
scan_arr = arrow._scan_buffer_to_arrow(source.data)
|
|
assert len(scan_arr) == 1
|
|
payload = _extract_scan_payload(scan_arr[0])
|
|
return cls(path=path, payload=payload, source=source)
|
|
|
|
def __getitem__(self, key: str) -> Pattern:
|
|
return self._materialize_pattern(key, persist=True)
|
|
|
|
def __iter__(self) -> Iterator[str]:
|
|
return iter(self._payload.cell_order)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._payload.cell_order)
|
|
|
|
def __contains__(self, key: object) -> bool:
|
|
return key in self._payload.cells
|
|
|
|
def source_order(self) -> tuple[str, ...]:
|
|
return tuple(self._payload.cell_order)
|
|
|
|
def raw_struct_bytes(self, name: str) -> bytes:
|
|
struct_range = self._payload.cells[name].struct_range
|
|
return self._source.raw_slice(struct_range.start, struct_range.end)
|
|
|
|
def can_copy_raw_struct(self, name: str) -> bool:
|
|
return name not in self._cache
|
|
|
|
def materialize_many(
|
|
self,
|
|
names: Sequence[str],
|
|
*,
|
|
persist: bool = True,
|
|
) -> LibraryView:
|
|
mats = self._materialize_patterns(names, persist=persist)
|
|
return LibraryView(mats)
|
|
|
|
def _materialize_patterns(
|
|
self,
|
|
names: Sequence[str],
|
|
*,
|
|
persist: bool,
|
|
) -> dict[str, Pattern]:
|
|
ordered_names = list(dict.fromkeys(names))
|
|
missing = [name for name in ordered_names if name not in self._payload.cells]
|
|
if missing:
|
|
raise KeyError(missing[0])
|
|
|
|
materialized: dict[str, Pattern] = {}
|
|
uncached = [name for name in ordered_names if name not in self._cache]
|
|
if uncached:
|
|
ranges = numpy.asarray(
|
|
[
|
|
[
|
|
self._payload.cells[name].struct_range.start,
|
|
self._payload.cells[name].struct_range.end,
|
|
]
|
|
for name in uncached
|
|
],
|
|
dtype=numpy.uint64,
|
|
)
|
|
arrow_arr = arrow._read_selected_cells_to_arrow(self._source.data, ranges)
|
|
assert len(arrow_arr) == 1
|
|
selected_lib, _info = arrow.read_arrow(arrow_arr[0])
|
|
for name in uncached:
|
|
pat = selected_lib[name]
|
|
materialized[name] = pat
|
|
if persist:
|
|
self._cache[name] = pat
|
|
|
|
for name in ordered_names:
|
|
if name in self._cache:
|
|
materialized[name] = self._cache[name]
|
|
return materialized
|
|
|
|
def _materialize_pattern(self, name: str, *, persist: bool) -> Pattern:
|
|
return self._materialize_patterns((name,), persist=persist)[name]
|
|
|
|
def _raw_children(self, name: str) -> set[str]:
|
|
return set(self._payload.cells[name].children)
|
|
|
|
def _collect_raw_transforms(self, cell: _CellScan, target_id: int) -> list[NDArray[numpy.float64]]:
|
|
refs = self._payload.refs
|
|
start = cell.ref_start
|
|
stop = cell.ref_stop
|
|
if stop <= start:
|
|
return []
|
|
|
|
targets = refs.targets[start:stop]
|
|
mask = targets == target_id
|
|
if not mask.any():
|
|
return []
|
|
|
|
rows: list[NDArray[numpy.float64]] = []
|
|
counts = refs.counts[start:stop]
|
|
unit_mask = mask & (counts[:, 0] == 1) & (counts[:, 1] == 1)
|
|
if unit_mask.any():
|
|
rows.append(_make_ref_rows(
|
|
refs.xy[start:stop][unit_mask],
|
|
refs.angle_rad[start:stop][unit_mask],
|
|
refs.invert_y[start:stop][unit_mask],
|
|
refs.scale[start:stop][unit_mask],
|
|
))
|
|
|
|
aref_indices = numpy.nonzero(mask & ~unit_mask)[0]
|
|
for idx in aref_indices:
|
|
abs_idx = start + int(idx)
|
|
rows.append(_expand_aref_row(
|
|
xy=refs.xy[abs_idx],
|
|
xy0=refs.xy0[abs_idx],
|
|
xy1=refs.xy1[abs_idx],
|
|
counts=refs.counts[abs_idx],
|
|
angle_rad=float(refs.angle_rad[abs_idx]),
|
|
invert_y=bool(refs.invert_y[abs_idx]),
|
|
scale=float(refs.scale[abs_idx]),
|
|
))
|
|
return rows
|
|
|
|
def child_graph(
|
|
self,
|
|
dangling: dangling_mode_t = 'error',
|
|
) -> dict[str, set[str]]:
|
|
graph: dict[str, set[str]] = {}
|
|
for name in self._payload.cell_order:
|
|
if name in self._cache:
|
|
graph[name] = {child for child, refs in self._cache[name].refs.items() if child is not None and refs}
|
|
else:
|
|
graph[name] = self._raw_children(name)
|
|
|
|
existing = set(graph)
|
|
dangling_refs = set().union(*(children - existing for children in graph.values()))
|
|
if dangling == 'error':
|
|
if dangling_refs:
|
|
raise self._dangling_refs_error(cast('set[str]', dangling_refs), 'building child graph')
|
|
return graph
|
|
if dangling == 'ignore':
|
|
return {name: {child for child in children if child in existing} for name, children in graph.items()}
|
|
|
|
for child in dangling_refs:
|
|
graph.setdefault(cast('str', child), set())
|
|
return graph
|
|
|
|
def parent_graph(
|
|
self,
|
|
dangling: dangling_mode_t = 'error',
|
|
) -> dict[str, set[str]]:
|
|
child_graph = self.child_graph(dangling='include' if dangling == 'include' else 'ignore')
|
|
existing = set(self.keys())
|
|
igraph: dict[str, set[str]] = {name: set() for name in child_graph}
|
|
for parent, children in child_graph.items():
|
|
for child in children:
|
|
if child in existing or dangling == 'include':
|
|
igraph.setdefault(child, set()).add(parent)
|
|
if dangling == 'error':
|
|
raw = self.child_graph(dangling='include')
|
|
dangling_refs = set().union(*(children - existing for children in raw.values()))
|
|
if dangling_refs:
|
|
raise self._dangling_refs_error(cast('set[str]', dangling_refs), 'building parent graph')
|
|
return igraph
|
|
|
|
def subtree(
|
|
self,
|
|
tops: str | Sequence[str],
|
|
) -> ILibraryView:
|
|
if isinstance(tops, str):
|
|
tops = (tops,)
|
|
keep = cast('set[str]', self.referenced_patterns(tops) - {None})
|
|
keep |= set(tops)
|
|
return self.materialize_many(tuple(keep), persist=True)
|
|
|
|
def tops(self) -> list[str]:
|
|
graph = self.child_graph(dangling='ignore')
|
|
names = set(graph)
|
|
not_toplevel: set[str] = set()
|
|
for children in graph.values():
|
|
not_toplevel |= children
|
|
return list(names - not_toplevel)
|
|
|
|
def with_ports_from_data(
|
|
self,
|
|
*,
|
|
layers: Sequence[tuple[int, int] | int],
|
|
max_depth: int = 0,
|
|
skip_subcells: bool = True,
|
|
) -> PortsLibraryView:
|
|
return PortsLibraryView(
|
|
self,
|
|
layers=layers,
|
|
max_depth=max_depth,
|
|
skip_subcells=skip_subcells,
|
|
)
|
|
|
|
def close(self) -> None:
|
|
data = self._source.data
|
|
if isinstance(data, mmap.mmap):
|
|
data.close()
|
|
if self._source.handle is not None:
|
|
self._source.handle.close()
|
|
self._source.handle = None
|
|
|
|
def __enter__(self) -> ArrowLibrary:
|
|
return self
|
|
|
|
def __exit__(self, *_args: object) -> None:
|
|
self.close()
|
|
|
|
def find_refs_local(
|
|
self,
|
|
name: str,
|
|
parent_graph: dict[str, set[str]] | None = None,
|
|
dangling: dangling_mode_t = 'error',
|
|
) -> dict[str, list[NDArray[numpy.float64]]]:
|
|
instances: dict[str, list[NDArray[numpy.float64]]] = defaultdict(list)
|
|
if parent_graph is None:
|
|
graph_mode = 'ignore' if dangling == 'ignore' else 'include'
|
|
parent_graph = self.parent_graph(dangling=graph_mode)
|
|
|
|
if name not in self:
|
|
if name not in parent_graph:
|
|
return instances
|
|
if dangling == 'error':
|
|
raise self._dangling_refs_error({name}, f'finding local refs for {name!r}')
|
|
if dangling == 'ignore':
|
|
return instances
|
|
|
|
target_id = self._payload.cells.get(name)
|
|
for parent in parent_graph.get(name, set()):
|
|
if parent in self._cache:
|
|
for ref in self._cache[parent].refs.get(name, []):
|
|
instances[parent].append(ref.as_transforms())
|
|
continue
|
|
|
|
if target_id is None or parent not in self._payload.cells:
|
|
continue
|
|
rows = self._collect_raw_transforms(self._payload.cells[parent], target_id.cell_id)
|
|
if rows:
|
|
instances[parent].extend(rows)
|
|
return instances
|
|
|
|
def find_refs_global(
|
|
self,
|
|
name: str,
|
|
order: list[str] | None = None,
|
|
parent_graph: dict[str, set[str]] | None = None,
|
|
dangling: dangling_mode_t = 'error',
|
|
) -> dict[tuple[str, ...], NDArray[numpy.float64]]:
|
|
graph_mode = 'ignore' if dangling == 'ignore' else 'include'
|
|
if order is None:
|
|
order = self.child_order(dangling=graph_mode)
|
|
if parent_graph is None:
|
|
parent_graph = self.parent_graph(dangling=graph_mode)
|
|
|
|
if name not in self:
|
|
if name not in parent_graph:
|
|
return {}
|
|
if dangling == 'error':
|
|
raise self._dangling_refs_error({name}, f'finding global refs for {name!r}')
|
|
if dangling == 'ignore':
|
|
return {}
|
|
|
|
self_keys = set(self.keys())
|
|
transforms: dict[str, list[tuple[tuple[str, ...], NDArray[numpy.float64]]]]
|
|
transforms = defaultdict(list)
|
|
for parent, vals in self.find_refs_local(name, parent_graph=parent_graph, dangling=dangling).items():
|
|
transforms[parent] = [((name,), numpy.concatenate(vals))]
|
|
|
|
for next_name in order:
|
|
if next_name not in transforms:
|
|
continue
|
|
if not parent_graph.get(next_name, set()) & self_keys:
|
|
continue
|
|
|
|
outers = self.find_refs_local(next_name, parent_graph=parent_graph, dangling=dangling)
|
|
inners = transforms.pop(next_name)
|
|
for parent, outer in outers.items():
|
|
outer_tf = numpy.concatenate(outer)
|
|
for path, inner in inners:
|
|
combined = apply_transforms(outer_tf, inner)
|
|
transforms[parent].append(((next_name,) + path, combined))
|
|
|
|
result = {}
|
|
for parent, targets in transforms.items():
|
|
for path, instances in targets:
|
|
full_path = (parent,) + path
|
|
result[full_path] = instances
|
|
return result
|
|
|
|
|
|
def readfile(
|
|
filename: str | pathlib.Path,
|
|
) -> tuple[ArrowLibrary, dict[str, Any]]:
|
|
lib = ArrowLibrary.from_file(filename)
|
|
return lib, lib.library_info
|
|
|
|
|
|
def load_libraryfile(
|
|
filename: str | pathlib.Path,
|
|
) -> tuple[ArrowLibrary, dict[str, Any]]:
|
|
return readfile(filename)
|