prefab 1.3.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

prefab/predict.py

@@ -1,12 +1,22 @@
-"""Prediction functions for ndarrays of device geometries."""
+"""
+Serverless prediction interface for nanofabrication modeling.
+
+This module provides functions for predicting nanofabrication outcomes using machine
+learning models hosted on a serverless platform. It supports multiple input formats
+(ndarrays, polygons, GDSII files) and model types (prediction, correction,
+segmentation). Gradient computation is available for inverse design applications
+using automatic differentiation.
+"""
 
 import base64
 import io
 import json
 import os
+from typing import Any
 
 import gdstk
 import numpy as np
+import numpy.typing as npt
 import requests
 import toml
 from autograd import primitive
@@ -21,11 +31,11 @@ ENDPOINT_VERSION = "3"
 
 
 def predict_array(
-    device_array: np.ndarray,
+    device_array: npt.NDArray[Any],
     model: Model,
     model_type: str,
     binarize: bool,
-) -> np.ndarray:
+) -> npt.NDArray[Any]:
     """
     Predict the nanofabrication outcome of a device array using a specified model.
 
@@ -40,14 +50,12 @@ def predict_array(
         various transformations to predict the nanofabrication process.
     model : Model
         The model to use for prediction, representing a specific fabrication process and
-        dataset. This model encapsulates details about the fabrication foundry, process,
-        material, technology, thickness, and sidewall presence, as defined in
-        `models.py`. Each model is associated with a version and dataset that detail its
-        creation and the data it was trained on, ensuring the prediction is tailored to
-        specific fabrication parameters.
+        dataset. This model encapsulates details about the fabrication foundry and
+        process, as defined in `models.py`. Each model is associated with a version and
+        dataset that detail its creation and the data it was trained on, ensuring the
+        prediction is tailored to specific fabrication parameters.
     model_type : str
-        The type of model to use (e.g., 'p' for prediction, 'c' for correction, or 's'
-        for SEMulate, 'b' for segmentation).
+        The type of model to use (e.g., 'p' for prediction or 'c' for correction).
     binarize : bool
         If True, the predicted device geometry will be binarized using a threshold
         method. This is useful for converting probabilistic predictions into binary
@@ -110,11 +118,11 @@ def predict_array(
 
 
 def _predict_poly(
-    polygon_points: list,
+    polygon_points: list[Any],
     model: Model,
     model_type: str,
     eta: float = 0.5,
-) -> list:
+) -> list[Any]:
     """
     Predict the nanofabrication outcome for a geometry (list of polygons).
 
@@ -128,13 +136,12 @@ def _predict_poly(
         List of polygon points, where each polygon is a list of [x, y] coordinates.
     model : Model
         The model to use for prediction, representing a specific fabrication process and
-        dataset. This model encapsulates details about the fabrication foundry, process,
-        material, technology, thickness, and sidewall presence, as defined in
-        `models.py`. Each model is associated with a version and dataset that detail its
-        creation and the data it was trained on, ensuring the prediction is tailored to
-        specific fabrication parameters.
+        dataset. This model encapsulates details about the fabrication foundry and
+        process, as defined in `models.py`. Each model is associated with a version and
+        dataset that detail its creation and the data it was trained on, ensuring the
+        prediction is tailored to specific fabrication parameters.
     model_type : str
-        The type of model to use ('p' for prediction, 'c' for correction).
+        The type of model to use (e.g., 'p' for prediction or 'c' for correction).
     eta : float
         The threshold value for binarization. Defaults to 0.5. Because intermediate
         values cannot be preserved in the polygon data, the predicted polygons are
@@ -198,7 +205,7 @@ def predict_gds(
     model_type: str,
     gds_layer: tuple[int, int] = (1, 0),
     eta: float = 0.5,
-    output_path: str = None,
+    output_path: str | None = None,
 ) -> None:
     """
     Predict the nanofabrication outcome for a GDS file and cell.
@@ -216,13 +223,12 @@ def predict_gds(
         The name of the cell within the GDS file to predict.
     model : Model
         The model to use for prediction, representing a specific fabrication process and
-        dataset. This model encapsulates details about the fabrication foundry, process,
-        material, technology, thickness, and sidewall presence, as defined in
-        `models.py`. Each model is associated with a version and dataset that detail its
-        creation and the data it was trained on, ensuring the prediction is tailored to
-        specific fabrication parameters.
+        dataset. This model encapsulates details about the fabrication foundry and
+        process, as defined in `models.py`. Each model is associated with a version and
+        dataset that detail its creation and the data it was trained on, ensuring the
+        prediction is tailored to specific fabrication parameters.
     model_type : str
-        The type of model to use ('p' for prediction, 'c' for correction).
+        The type of model to use (e.g., 'p' for prediction or 'c' for correction).
     gds_layer : tuple[int, int]
         The layer and datatype to use within the GDS file. Defaults to (1, 0).
     eta : float
@@ -242,7 +248,14 @@ def predict_gds(
         returns an error or invalid response.
     """
     gdstk_library = gdstk.read_gds(gds_path)
-    gdstk_cell = gdstk_library[cell_name]
+    cells = [
+        cell
+        for cell in gdstk_library.cells
+        if isinstance(cell, gdstk.Cell) and cell.name == cell_name
+    ]
+    if not cells:
+        raise ValueError(f"Cell '{cell_name}' not found in GDS file")
+    gdstk_cell = cells[0]
 
     predicted_cell = predict_gdstk(
         gdstk_cell=gdstk_cell,
@@ -261,6 +274,7 @@ def predict_gds(
     gdstk_library.add(predicted_cell)
 
     write_path = output_path if output_path is not None else gds_path
+    print(f"Writing to {write_path}")
     gdstk_library.write_gds(write_path, max_points=8190)
 
 
@@ -283,13 +297,12 @@ def predict_gdstk(
         The gdstk.Cell object containing polygons to predict.
     model : Model
         The model to use for prediction, representing a specific fabrication process and
-        dataset. This model encapsulates details about the fabrication foundry, process,
-        material, technology, thickness, and sidewall presence, as defined in
-        `models.py`. Each model is associated with a version and dataset that detail its
-        creation and the data it was trained on, ensuring the prediction is tailored to
-        specific fabrication parameters.
+        dataset. This model encapsulates details about the fabrication foundry and
+        process, as defined in `models.py`. Each model is associated with a version and
+        dataset that detail its creation and the data it was trained on, ensuring the
+        prediction is tailored to specific fabrication parameters.
     model_type : str
-        The type of model to use ('p' for prediction, 'c' for correction).
+        The type of model to use (e.g., 'p' for prediction or 'c' for correction).
     gds_layer : tuple[int, int]
         The layer and datatype to use within the GDSTK cell. Defaults to (1, 0).
     eta : float
@@ -317,7 +330,7 @@ def predict_gdstk(
     if not polygons:
         raise ValueError("No polygons found in the specified layer")
 
-    polygon_points = [polygon.points.tolist() for polygon in polygons]
+    polygon_points = [polygon.points.tolist() for polygon in polygons]  # pyright: ignore[reportAttributeAccessIssue]
 
     predicted_polygon_data = _predict_poly(
         polygon_points=polygon_points,
@@ -326,7 +339,8 @@ def predict_gdstk(
         eta=eta,
     )
 
-    result_cell = gdstk.Cell(f"{gdstk_cell.name}_predicted")
+    suffix = "corrected" if model_type == "c" else "predicted"
+    result_cell = gdstk.Cell(f"{gdstk_cell.name}_{suffix}")
 
     polygons_by_channel = {}
     for polygon_data in predicted_polygon_data:
@@ -344,15 +358,15 @@ def predict_gdstk(
 
         for points in points_list:
             points_array = np.array(points)
-            polygon = gdstk.Polygon(points_array, layer=layer, datatype=datatype)
+            polygon = gdstk.Polygon(points_array, layer=layer, datatype=datatype)  # pyright: ignore[reportArgumentType]
             result_cell.add(polygon)
 
     return result_cell
 
 
 def _predict_array_with_grad(
-    device_array: np.ndarray, model: Model
-) -> tuple[np.ndarray, np.ndarray]:
+    device_array: npt.NDArray[Any], model: Model
+) -> tuple[npt.NDArray[Any], npt.NDArray[Any]]:
     """
     Predict the nanofabrication outcome of a device array and compute its gradient.
 
@@ -360,11 +374,6 @@ def _predict_array_with_grad(
     device array using a specified model. It also computes the gradient of the
     prediction with respect to the input device array.
 
-    Notes
-    -----
-    This function is currently not used in the main `predict_array` function as
-    the main `predict_array` function (e.g., GPU support and progress bar) for now.
-
     Parameters
     ----------
     device_array : np.ndarray
@@ -422,7 +431,9 @@ def _predict_array_with_grad(
 
 
 @primitive
-def predict_array_with_grad(device_array: np.ndarray, model: Model) -> np.ndarray:
+def predict_array_with_grad(
+    device_array: npt.NDArray[Any], model: Model
+) -> npt.NDArray[Any]:
     """
     Predict the nanofabrication outcome of a device array and compute its gradient.
 
@@ -453,11 +464,13 @@ def predict_array_with_grad(device_array: np.ndarray, model: Model) -> np.ndarra
     prediction_array, gradient_array = _predict_array_with_grad(
         device_array=device_array, model=model
     )
-    predict_array_with_grad.gradient_array = gradient_array  # type: ignore
+    predict_array_with_grad.gradient_array = gradient_array  # pyright: ignore[reportFunctionMemberAccess]
     return prediction_array
 
 
-def predict_array_with_grad_vjp(ans: np.ndarray, device_array: np.ndarray, *args):
+def predict_array_with_grad_vjp(
+    ans: npt.NDArray[Any], device_array: npt.NDArray[Any], *args: Any
+) -> Any:
     """
     Define the vector-Jacobian product (VJP) for the prediction function.
 
@@ -475,10 +488,10 @@ def predict_array_with_grad_vjp(ans: np.ndarray, device_array: np.ndarray, *args
     function
         A function that computes the VJP given an upstream gradient `g`.
     """
-    grad_x = predict_array_with_grad.gradient_array  # type: ignore
+    grad_x = predict_array_with_grad.gradient_array  # pyright: ignore[reportFunctionMemberAccess]
 
-    def vjp(g: np.ndarray) -> np.ndarray:
-        return g * grad_x
+    def vjp(g: npt.NDArray[Any]) -> npt.NDArray[Any]:
+        return g * grad_x  # type: ignore[no-any-return]
 
     return vjp
 
@@ -486,7 +499,7 @@ def predict_array_with_grad_vjp(ans: np.ndarray, device_array: np.ndarray, *args
 defvjp(predict_array_with_grad, predict_array_with_grad_vjp)
 
 
-def _encode_array(array):
+def _encode_array(array: npt.NDArray[Any]) -> str:
     """Encode an ndarray as a base64 encoded image for transmission."""
     image = Image.fromarray(np.uint8(array * 255))
     buffered = io.BytesIO()
@@ -495,14 +508,14 @@ def _encode_array(array):
     return encoded_png
 
 
-def _decode_array(encoded_png):
+def _decode_array(encoded_png: str) -> npt.NDArray[Any]:
     """Decode a base64 encoded image and return an ndarray."""
     binary_data = base64.b64decode(encoded_png)
     image = Image.open(io.BytesIO(binary_data))
-    return np.array(image) / 255
+    return np.array(image) / 255  # type: ignore[no-any-return]
 
 
-def _prepare_headers():
+def _prepare_headers() -> dict[str, str]:
     """Prepare HTTP headers for a server request."""
     token_file_path = os.path.expanduser("~/.prefab.toml")
     try:
@@ -519,7 +532,7 @@ def _prepare_headers():
     except FileNotFoundError:
         raise FileNotFoundError(
             "Could not validate user.\n"
-            "Please update prefab using: pip install --upgrade prefab.\n"
-            "Signup/login and generate a new token.\n"
-            "See https://docs.prefabphotonics.com/."
+            + "Please update prefab using: pip install --upgrade prefab.\n"
+            + "Signup/login and generate a new token.\n"
+            + "See https://docs.prefabphotonics.com/."
         ) from None