Skip to content

Object-oriented API: Mask, LitheEngine, Report

The OO façade is a thin wrapper over the existing functional API. It exists so fab-/EDA-shaped callers can think in masks and engines without going through the tensor + registry plumbing directly.

The functional API is unchanged — compute_epe, models.registry, and workflow.tiling.tile_layout all work exactly as before. Use whichever shape fits the task.

Quick start

from openlithohub import Mask, LitheEngine

mask     = Mask.from_oasis("design.oas", layer="1:0", pixel_size_nm=1.0)
engine   = LitheEngine(model="neural-ilt", node="3nm-euv")
optimized = engine.optimize(mask)
report    = engine.evaluate(optimized, target=mask)

print(report.epe_mean_nm, report.pvband_mean_nm, report.drc_violations)
optimized.to_oasis("optimized.oas")

What goes where

Concept Class Wraps
(tensor, pixel_size_nm, layer) Mask _load_layout_as_tensor, workflow.export.export_oasis
Run a model on a layout LitheEngine models.registry, workflow.tiling, workflow.halo
Aggregate metric & compliance Report benchmark.metrics.*, benchmark.compliance.*

Constructors

Mask offers explicit and suffix-sniffing constructors:

Mask.from_tensor(t, pixel_size_nm=0.5, layer="1:0")
Mask.from_pt("design.pt")
Mask.from_npy("design.npy")
Mask.from_oasis("design.oas", layer="1:0")
Mask.from_gds("design.gds", layer="1:0")
Mask.load("design.oas", layer="1:0")    # dispatches by file suffix

Mask is a frozen dataclass — once constructed, the (tensor, pixel_size_nm, layer) triplet cannot be mutated. Build a new Mask if you need a different pitch.

Backward compatibility

engine.optimize(...) and engine.evaluate(...) accept either a Mask or a raw torch.Tensor. Existing tensor-first callers do not need to change.

Reference

openlithohub.api.mask

Mask — frozen dataclass bundling (tensor, pixel_size_nm, layer).

Wraps OpenLithoHub's existing layout I/O so callers can write Mask.from_oasis("design.oas", layer="1:0") instead of going through load_layout and export_oasis directly.

Mask dataclass

A 2-D mask tensor with its physical pixel pitch and (optional) source layer.

The fab/EDA-facing handle that LitheEngine consumes and produces. Construction is via classmethod constructors. The dataclass is frozen, so the three fields cannot be rebound after construction — but note that frozen=True does not protect against in-place mutation of the underlying tensor (mask.tensor[i, j] = 0 still mutates storage). Treat Mask as a structural binding of (tensor, pixel_size_nm, layer), not as a deep-immutable value.

Source code in src/openlithohub/api/mask.py
@dataclass(frozen=True)
class Mask:
    """A 2-D mask tensor with its physical pixel pitch and (optional) source layer.

    The fab/EDA-facing handle that ``LitheEngine`` consumes and produces.
    Construction is via classmethod constructors. The dataclass is frozen,
    so the three fields cannot be rebound after construction — but note
    that ``frozen=True`` does not protect against in-place mutation of the
    underlying tensor (``mask.tensor[i, j] = 0`` still mutates storage).
    Treat ``Mask`` as a structural binding of ``(tensor, pixel_size_nm,
    layer)``, not as a deep-immutable value.
    """

    tensor: torch.Tensor
    pixel_size_nm: float = 1.0
    layer: str | None = None

    @property
    def shape(self) -> tuple[int, int]:
        h, w = self.tensor.shape
        return int(h), int(w)

    def __array__(self, dtype: object = None) -> NDArray[np.float32]:
        arr = self.tensor.detach().cpu().numpy()
        if dtype is not None:
            # numpy hands `dtype` in via the protocol as an arbitrary object
            # (numpy's own type isn't exposed for this hook), so we narrow
            # it to what `astype` accepts before calling.
            arr = arr.astype(dtype)  # type: ignore[call-overload]
        return arr

    @classmethod
    def from_tensor(
        cls,
        tensor: torch.Tensor,
        *,
        pixel_size_nm: float = 1.0,
        layer: str | None = None,
    ) -> Mask:
        if not isinstance(tensor, torch.Tensor):
            raise TypeError(f"expected torch.Tensor, got {type(tensor).__name__}")
        if tensor.ndim != 2:
            raise ValueError(f"Mask tensor must be 2-D (H, W), got ndim={tensor.ndim}")
        return cls(tensor=tensor.float(), pixel_size_nm=pixel_size_nm, layer=layer)

    @classmethod
    def from_pt(cls, path: str | Path, *, pixel_size_nm: float = 1.0) -> Mask:
        t = load_layout(Path(path), pixel_size_nm)
        return cls(tensor=t, pixel_size_nm=pixel_size_nm, layer=None)

    @classmethod
    def from_npy(cls, path: str | Path, *, pixel_size_nm: float = 1.0) -> Mask:
        t = load_layout(Path(path), pixel_size_nm)
        return cls(tensor=t, pixel_size_nm=pixel_size_nm, layer=None)

    @classmethod
    def from_oasis(
        cls,
        path: str | Path,
        *,
        pixel_size_nm: float = 1.0,
        layer: str | None = None,
    ) -> Mask:
        t = load_layout(Path(path), pixel_size_nm, layer=layer)
        return cls(tensor=t, pixel_size_nm=pixel_size_nm, layer=layer)

    @classmethod
    def from_gds(
        cls,
        path: str | Path,
        *,
        pixel_size_nm: float = 1.0,
        layer: str | None = None,
    ) -> Mask:
        t = load_layout(Path(path), pixel_size_nm, layer=layer)
        return cls(tensor=t, pixel_size_nm=pixel_size_nm, layer=layer)

    @classmethod
    def load(
        cls,
        path: str | Path,
        *,
        pixel_size_nm: float = 1.0,
        layer: str | None = None,
    ) -> Mask:
        """Suffix-sniffing constructor.

        ``.pt`` / ``.npy`` ignore ``layer``. ``.oas`` / ``.gds`` honour it.
        """
        suffix = Path(path).suffix.lower()
        if suffix in {".pt", ".npy"}:
            if layer is not None:
                raise ValueError(f"layer is meaningless for {suffix} inputs")
            t = load_layout(Path(path), pixel_size_nm)
            return cls(tensor=t, pixel_size_nm=pixel_size_nm, layer=None)
        if suffix in {".oas", ".gds"}:
            t = load_layout(Path(path), pixel_size_nm, layer=layer)
            return cls(tensor=t, pixel_size_nm=pixel_size_nm, layer=layer)
        raise ValueError(f"unsupported extension {suffix!r} — expected .pt / .npy / .oas / .gds")

    def to_pt(self, path: str | Path) -> None:
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        torch.save(self.tensor, str(path))

    def to_npy(self, path: str | Path) -> None:
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        np.save(str(path), self.tensor.detach().cpu().numpy())

    def to_oasis(self, path: str | Path, *, mode: str = "curvilinear") -> None:
        export_oasis(self.tensor, path, mode=mode, pixel_size_nm=self.pixel_size_nm)

    def to_gds(self, path: str | Path, *, mode: str = "curvilinear") -> None:
        # `export_oasis` writes via klayout, which sniffs the format from the
        # filename suffix — `.gds` produces GDSII, `.oas` produces OASIS.
        export_oasis(self.tensor, path, mode=mode, pixel_size_nm=self.pixel_size_nm)

load(path, *, pixel_size_nm=1.0, layer=None) classmethod

Suffix-sniffing constructor.

.pt / .npy ignore layer. .oas / .gds honour it.

Source code in src/openlithohub/api/mask.py
@classmethod
def load(
    cls,
    path: str | Path,
    *,
    pixel_size_nm: float = 1.0,
    layer: str | None = None,
) -> Mask:
    """Suffix-sniffing constructor.

    ``.pt`` / ``.npy`` ignore ``layer``. ``.oas`` / ``.gds`` honour it.
    """
    suffix = Path(path).suffix.lower()
    if suffix in {".pt", ".npy"}:
        if layer is not None:
            raise ValueError(f"layer is meaningless for {suffix} inputs")
        t = load_layout(Path(path), pixel_size_nm)
        return cls(tensor=t, pixel_size_nm=pixel_size_nm, layer=None)
    if suffix in {".oas", ".gds"}:
        t = load_layout(Path(path), pixel_size_nm, layer=layer)
        return cls(tensor=t, pixel_size_nm=pixel_size_nm, layer=layer)
    raise ValueError(f"unsupported extension {suffix!r} — expected .pt / .npy / .oas / .gds")

openlithohub.api.engine

LitheEngine — thin wrapper over registry + tile/halo/stitch pipeline.

Mirrors the body of server.app._run_optimize minus filesystem I/O so callers can drive the engine in-process without touching the HTTP server or the CLI helpers.

LitheEngine

Object-oriented driver for the OpenLithoHub optimization pipeline.

engine = LitheEngine(model="neural-ilt", node="3nm-euv") optimized = engine.optimize(mask) report = engine.evaluate(optimized, target=mask)

Source code in src/openlithohub/api/engine.py
class LitheEngine:
    """Object-oriented driver for the OpenLithoHub optimization pipeline.

    ``engine = LitheEngine(model="neural-ilt", node="3nm-euv")``
    ``optimized = engine.optimize(mask)``
    ``report = engine.evaluate(optimized, target=mask)``
    """

    def __init__(
        self,
        model: str | LithographyModel,
        *,
        node: str | None = None,
        tile_size: int = 2048,
        pretrained: bool = False,
        **model_kwargs: Any,
    ) -> None:
        register_builtin_models()

        if isinstance(model, str):
            kwargs: dict[str, Any] = dict(model_kwargs)
            if pretrained:
                kwargs.setdefault("pretrained", True)
            self._model: LithographyModel = registry.get(model, **kwargs)
            # Engine constructed the instance, so it owns the setup() call —
            # and, symmetrically, the teardown() in __exit__.
            self._model.setup()
            self._owns_model = True
        elif isinstance(model, LithographyModel):
            if model_kwargs or pretrained:
                raise ValueError(
                    "model_kwargs / pretrained only apply when `model` is a name; "
                    "pass a fully constructed LithographyModel without them."
                )
            # Caller-supplied instance: assume the caller has already called
            # setup(). Calling it again would re-load weights / re-init GPU
            # state in non-idempotent models like NeuralILTModel. The caller
            # also owns teardown — we must not close resources we did not open.
            self._model = model
            self._owns_model = False
        else:
            raise TypeError(
                f"`model` must be a name (str) or LithographyModel, got {type(model).__name__}"
            )

        # Let `get_node` raise KeyError on typos; silently coercing unknown
        # node names to None hides physics-affecting misconfiguration.
        self._node_config: ProcessNodeConfig | None = get_node(node) if node is not None else None
        self._tile_size = tile_size

    @property
    def model(self) -> LithographyModel:
        return self._model

    @property
    def node(self) -> ProcessNodeConfig | None:
        return self._node_config

    @staticmethod
    def list_models() -> list[str]:
        register_builtin_models()
        return registry.list_models()

    def close(self) -> None:
        """Tear down the underlying model if the engine constructed it.

        Safe to call multiple times. No-op for caller-supplied models
        (the caller owns those — closing them here would yank resources
        out from under code the engine never owned).
        """
        if self._owns_model and self._model is not None:
            self._model.teardown()
            self._owns_model = False

    def __enter__(self) -> LitheEngine:
        return self

    def __exit__(self, *_exc: object) -> None:
        self.close()

    def _resolve_pixel_size(self, supplied: float) -> float:
        # Mirror server/app.py:67–69 — when the caller did not override the
        # default 1.0 nm/px and a process node is active, prefer the node's
        # native pitch.
        if self._node_config is not None and supplied == 1.0:
            return self._node_config.pixel_size_nm
        return supplied

    def _coerce_to_mask(self, design: Mask | torch.Tensor) -> Mask:
        if isinstance(design, Mask):
            return design
        if isinstance(design, torch.Tensor):
            pixel_nm = self._node_config.pixel_size_nm if self._node_config is not None else 1.0
            return Mask.from_tensor(design, pixel_size_nm=pixel_nm)
        raise TypeError(f"expected Mask or torch.Tensor, got {type(design).__name__}")

    def optimize(self, design: Mask | torch.Tensor) -> Mask:
        """Run the model over ``design`` with tiling + halo + stitching.

        Returns a binarised ``Mask`` matching the input shape and pixel pitch.
        """
        in_mask = self._coerce_to_mask(design)
        pixel_nm = self._resolve_pixel_size(in_mask.pixel_size_nm)
        tensor = in_mask.tensor

        halo_px = compute_halo_px(
            node=self._node_config,
            model=self._model,
            pixel_nm=pixel_nm,
            tile_size=self._tile_size,
        )

        tiles = tile_layout(tensor, tile_size=self._tile_size, overlap=halo_px)
        tile_results: list[tuple[Any, torch.Tensor]] = []
        for tile in tiles:
            result = self._model.predict(tile.tensor)
            tile_results.append((tile, result.mask))

        h, w = tensor.shape
        stitched = stitch_tiles(tile_results, (int(h), int(w)))
        binarised = (stitched > 0.5).float()

        return Mask(tensor=binarised, pixel_size_nm=pixel_nm, layer=in_mask.layer)

    def evaluate(
        self,
        predicted: Mask | torch.Tensor,
        target: Mask | torch.Tensor,
    ) -> Report:
        """Compute the canonical metric / compliance battery on ``predicted`` vs ``target``."""
        pred = self._coerce_to_mask(predicted)
        tgt = self._coerce_to_mask(target)
        if pred.shape != tgt.shape:
            raise ValueError(f"shape mismatch: predicted {pred.shape} vs target {tgt.shape}")
        if pred.pixel_size_nm != tgt.pixel_size_nm:
            raise ValueError(
                f"pixel_size_nm mismatch: predicted {pred.pixel_size_nm} vs "
                f"target {tgt.pixel_size_nm}. EPE is reported in nanometers, so "
                f"masks at different pitches cannot be compared without resampling."
            )

        pixel_nm = self._resolve_pixel_size(pred.pixel_size_nm)

        epe = compute_epe(pred.tensor, tgt.tensor, pixel_size_nm=pixel_nm)
        pvband = compute_pvband(pred.tensor, pixel_size_nm=pixel_nm)
        drc = check_drc(pred.tensor, pixel_size_nm=pixel_nm)
        mrc = check_mrc(pred.tensor, pixel_size_nm=pixel_nm)
        shots = estimate_shot_count(pred.tensor, pixel_size_nm=pixel_nm)
        curvilinear_mrc = (
            check_curvilinear_mrc(pred.tensor, pixel_size_nm=pixel_nm)
            if self._model.supports_curvilinear
            else None
        )

        # Recompute the same halo `optimize()` would have used at this pitch,
        # so the report documents the tile/halo configuration that produced
        # (or would produce) `predicted` — useful for reproducing a run.
        halo_px = compute_halo_px(
            node=self._node_config,
            model=self._model,
            pixel_nm=pixel_nm,
            tile_size=self._tile_size,
        )

        return Report(
            epe_mean_nm=float(epe["epe_mean_nm"]),
            epe_max_nm=float(epe["epe_max_nm"]),
            epe_std_nm=float(epe["epe_std_nm"]),
            pvband_mean_nm=float(pvband["pvband_mean_nm"]),
            pvband_max_nm=float(pvband["pvband_max_nm"]),
            drc_violations=int(drc.violation_count),
            drc_passed=bool(drc.passed),
            mrc_violations=int(mrc.violation_count),
            mrc_passed=bool(mrc.passed),
            shot_count=int(shots["shot_count"]),
            estimated_write_time_s=float(shots["estimated_write_time_s"]),
            model_name=self._model.name,
            pixel_size_nm=pixel_nm,
            tile_size=int(self._tile_size),
            halo_px=int(halo_px),
            raw_epe=epe,
            raw_drc=drc,
            raw_mrc=mrc,
            raw_pvband=pvband,
            raw_shot_count=shots,
            raw_curvilinear_mrc=curvilinear_mrc,
        )

close()

Tear down the underlying model if the engine constructed it.

Safe to call multiple times. No-op for caller-supplied models (the caller owns those — closing them here would yank resources out from under code the engine never owned).

Source code in src/openlithohub/api/engine.py
def close(self) -> None:
    """Tear down the underlying model if the engine constructed it.

    Safe to call multiple times. No-op for caller-supplied models
    (the caller owns those — closing them here would yank resources
    out from under code the engine never owned).
    """
    if self._owns_model and self._model is not None:
        self._model.teardown()
        self._owns_model = False

optimize(design)

Run the model over design with tiling + halo + stitching.

Returns a binarised Mask matching the input shape and pixel pitch.

Source code in src/openlithohub/api/engine.py
def optimize(self, design: Mask | torch.Tensor) -> Mask:
    """Run the model over ``design`` with tiling + halo + stitching.

    Returns a binarised ``Mask`` matching the input shape and pixel pitch.
    """
    in_mask = self._coerce_to_mask(design)
    pixel_nm = self._resolve_pixel_size(in_mask.pixel_size_nm)
    tensor = in_mask.tensor

    halo_px = compute_halo_px(
        node=self._node_config,
        model=self._model,
        pixel_nm=pixel_nm,
        tile_size=self._tile_size,
    )

    tiles = tile_layout(tensor, tile_size=self._tile_size, overlap=halo_px)
    tile_results: list[tuple[Any, torch.Tensor]] = []
    for tile in tiles:
        result = self._model.predict(tile.tensor)
        tile_results.append((tile, result.mask))

    h, w = tensor.shape
    stitched = stitch_tiles(tile_results, (int(h), int(w)))
    binarised = (stitched > 0.5).float()

    return Mask(tensor=binarised, pixel_size_nm=pixel_nm, layer=in_mask.layer)

evaluate(predicted, target)

Compute the canonical metric / compliance battery on predicted vs target.

Source code in src/openlithohub/api/engine.py
def evaluate(
    self,
    predicted: Mask | torch.Tensor,
    target: Mask | torch.Tensor,
) -> Report:
    """Compute the canonical metric / compliance battery on ``predicted`` vs ``target``."""
    pred = self._coerce_to_mask(predicted)
    tgt = self._coerce_to_mask(target)
    if pred.shape != tgt.shape:
        raise ValueError(f"shape mismatch: predicted {pred.shape} vs target {tgt.shape}")
    if pred.pixel_size_nm != tgt.pixel_size_nm:
        raise ValueError(
            f"pixel_size_nm mismatch: predicted {pred.pixel_size_nm} vs "
            f"target {tgt.pixel_size_nm}. EPE is reported in nanometers, so "
            f"masks at different pitches cannot be compared without resampling."
        )

    pixel_nm = self._resolve_pixel_size(pred.pixel_size_nm)

    epe = compute_epe(pred.tensor, tgt.tensor, pixel_size_nm=pixel_nm)
    pvband = compute_pvband(pred.tensor, pixel_size_nm=pixel_nm)
    drc = check_drc(pred.tensor, pixel_size_nm=pixel_nm)
    mrc = check_mrc(pred.tensor, pixel_size_nm=pixel_nm)
    shots = estimate_shot_count(pred.tensor, pixel_size_nm=pixel_nm)
    curvilinear_mrc = (
        check_curvilinear_mrc(pred.tensor, pixel_size_nm=pixel_nm)
        if self._model.supports_curvilinear
        else None
    )

    # Recompute the same halo `optimize()` would have used at this pitch,
    # so the report documents the tile/halo configuration that produced
    # (or would produce) `predicted` — useful for reproducing a run.
    halo_px = compute_halo_px(
        node=self._node_config,
        model=self._model,
        pixel_nm=pixel_nm,
        tile_size=self._tile_size,
    )

    return Report(
        epe_mean_nm=float(epe["epe_mean_nm"]),
        epe_max_nm=float(epe["epe_max_nm"]),
        epe_std_nm=float(epe["epe_std_nm"]),
        pvband_mean_nm=float(pvband["pvband_mean_nm"]),
        pvband_max_nm=float(pvband["pvband_max_nm"]),
        drc_violations=int(drc.violation_count),
        drc_passed=bool(drc.passed),
        mrc_violations=int(mrc.violation_count),
        mrc_passed=bool(mrc.passed),
        shot_count=int(shots["shot_count"]),
        estimated_write_time_s=float(shots["estimated_write_time_s"]),
        model_name=self._model.name,
        pixel_size_nm=pixel_nm,
        tile_size=int(self._tile_size),
        halo_px=int(halo_px),
        raw_epe=epe,
        raw_drc=drc,
        raw_mrc=mrc,
        raw_pvband=pvband,
        raw_shot_count=shots,
        raw_curvilinear_mrc=curvilinear_mrc,
    )

openlithohub.api.report

Report — flat view over the existing metric / compliance outputs.

No new math. The flat fields are projections of fields already computed by benchmark.metrics and benchmark.compliance; the raw underlying results are kept on the dataclass for power users who want every field.

Report dataclass

Aggregated mask-quality report produced by LitheEngine.evaluate.

Source code in src/openlithohub/api/report.py
@dataclass(frozen=True)
class Report:
    """Aggregated mask-quality report produced by ``LitheEngine.evaluate``."""

    epe_mean_nm: float
    epe_max_nm: float
    epe_std_nm: float
    pvband_mean_nm: float
    pvband_max_nm: float

    drc_violations: int
    drc_passed: bool
    mrc_violations: int
    mrc_passed: bool

    shot_count: int
    estimated_write_time_s: float

    model_name: str
    pixel_size_nm: float
    tile_size: int
    halo_px: int

    raw_epe: EPEResult
    raw_drc: DRCResult
    raw_mrc: MRCResult
    raw_pvband: dict[str, float]
    raw_shot_count: dict[str, int | float]
    raw_curvilinear_mrc: CurvilinearMRCResult | None = field(default=None)

    def to_dict(self) -> dict[str, Any]:
        """JSON-serializable view. Raw payloads are recursively flattened."""
        return asdict(self)

    def __repr__(self) -> str:
        return (
            f"Report(model={self.model_name!r} "
            f"epe_mean={self.epe_mean_nm:.3f} nm "
            f"pvband_mean={self.pvband_mean_nm:.3f} nm "
            f"drc={self.drc_violations} mrc={self.mrc_violations} "
            f"shots={self.shot_count} "
            f"tile={self.tile_size} halo={self.halo_px})"
        )

to_dict()

JSON-serializable view. Raw payloads are recursively flattened.

Source code in src/openlithohub/api/report.py
def to_dict(self) -> dict[str, Any]:
    """JSON-serializable view. Raw payloads are recursively flattened."""
    return asdict(self)