Skip to content

Predict

Prediction functions for ndarrays of device geometries.

predict_array(device_array, model, model_type, binarize, gpu=False)

Predict the nanofabrication outcome of a device array using a specified model.

This function sends the device array to a serverless prediction service, which uses a specified machine learning model to predict the outcome of the nanofabrication process. The prediction can be performed on a GPU if specified.

Parameters:

Name Type Description Default
device_array ndarray

A 2D array representing the planar geometry of the device. This array undergoes various transformations to predict the nanofabrication process.

required
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.

required
model_type str

The type of model to use (e.g., 'p' for prediction, 'c' for correction, or 's' for SEMulate).

required
binarize bool

If True, the predicted device geometry will be binarized using a threshold method. This is useful for converting probabilistic predictions into binary geometries.

required
gpu bool

If True, the prediction will be performed on a GPU. Defaults to False. Note: The GPU option has more startup overhead and will take longer for small devices, but will be faster for larger devices.

False

Returns:

Type Description
ndarray

The predicted output array.

Raises:

Type Description
RuntimeError

If the request to the prediction service fails.

Source code in prefab/predict.py
def predict_array(
    device_array: np.ndarray,
    model: Model,
    model_type: str,
    binarize: bool,
    gpu: bool = False,
) -> np.ndarray:
    """
    Predict the nanofabrication outcome of a device array using a specified model.

    This function sends the device array to a serverless prediction service, which uses
    a specified machine learning model to predict the outcome of the nanofabrication
    process. The prediction can be performed on a GPU if specified.

    Parameters
    ----------
    device_array : np.ndarray
        A 2D array representing the planar geometry of the device. This array undergoes
        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.
    model_type : str
        The type of model to use (e.g., 'p' for prediction, 'c' for correction, or 's'
        for SEMulate).
    binarize : bool
        If True, the predicted device geometry will be binarized using a threshold
        method. This is useful for converting probabilistic predictions into binary
        geometries.
    gpu : bool
        If True, the prediction will be performed on a GPU. Defaults to False. Note: The
        GPU option has more startup overhead and will take longer for small devices, but
        will be faster for larger devices.

    Returns
    -------
    np.ndarray
        The predicted output array.

    Raises
    ------
    RuntimeError
        If the request to the prediction service fails.
    """
    headers = _prepare_headers()
    predict_data = _prepare_predict_data(device_array, model, model_type, binarize)
    endpoint_url = (
        f"{BASE_ENDPOINT_URL}-gpu-v{ENDPOINT_VERSION}.modal.run"
        if gpu
        else f"{BASE_ENDPOINT_URL}-v{ENDPOINT_VERSION}.modal.run"
    )

    try:
        with requests.post(
            endpoint_url,
            data=json.dumps(predict_data),
            headers=headers,
            stream=True,
        ) as response:
            response.raise_for_status()
            result = _process_response(response, model_type, binarize)
            if result is None:
                raise RuntimeError("No prediction result received.")
            return result
    except requests.RequestException as e:
        raise RuntimeError(f"Request failed: {e}") from e

predict_array_with_grad(device_array, model)

Predict the nanofabrication outcome of a device array and compute its gradient.

This function predicts the outcome of the nanofabrication process for a given device array using a specified model. It also computes the gradient of the prediction with respect to the input device array, making it suitable for use in automatic differentiation applications (e.g., autograd).

Parameters:

Name Type Description Default
device_array ndarray

A 2D array representing the planar geometry of the device.

required
model Model

The model to use for prediction, representing a specific fabrication process.

required

Returns:

Type Description
ndarray

The predicted output array.

Source code in prefab/predict.py
@primitive
def predict_array_with_grad(device_array: np.ndarray, model: Model) -> np.ndarray:
    """
    Predict the nanofabrication outcome of a device array and compute its gradient.

    This function predicts the outcome of the nanofabrication process for a given
    device array using a specified model. It also computes the gradient of the
    prediction with respect to the input device array, making it suitable for use in
    automatic differentiation applications (e.g., autograd).

    Parameters
    ----------
    device_array : np.ndarray
        A 2D array representing the planar geometry of the device.
    model : Model
        The model to use for prediction, representing a specific fabrication process.

    Returns
    -------
    np.ndarray
        The predicted output array.
    """
    prediction_array, gradient_array = _predict_array_with_grad(
        device_array=device_array, model=model
    )
    predict_array_with_grad.gradient_array = gradient_array  # type: ignore
    return prediction_array

predict_array_with_grad_vjp(ans, device_array, *args)

Define the vector-Jacobian product (VJP) for the prediction function.

Parameters:

Name Type Description Default
ans ndarray

The output of the predict_array_with_grad function.

required
device_array ndarray

The input device array for which the gradient is computed.

required
*args

Additional arguments that aren't used in the VJP computation.

()

Returns:

Type Description
function

A function that computes the VJP given an upstream gradient g.

Source code in prefab/predict.py
def predict_array_with_grad_vjp(ans: np.ndarray, device_array: np.ndarray, *args):
    """
    Define the vector-Jacobian product (VJP) for the prediction function.

    Parameters
    ----------
    ans : np.ndarray
        The output of the `predict_array_with_grad` function.
    device_array : np.ndarray
        The input device array for which the gradient is computed.
    *args :
        Additional arguments that aren't used in the VJP computation.

    Returns
    -------
    function
        A function that computes the VJP given an upstream gradient `g`.
    """
    grad_x = predict_array_with_grad.gradient_array  # type: ignore

    def vjp(g: np.ndarray) -> np.ndarray:
        return g * grad_x

    return vjp

predict_gdstk(gdstk_cell, model, model_type, gds_layer=(1, 0), eta=0.5)

Predict the nanofabrication outcome of a gdstk cell using a specified model.

This function extracts polygons from a gdstk cell, sends them to the prediction server, and returns a new cell containing the predicted polygons.

Parameters:

Name Type Description Default
gdstk_cell Cell

The gdstk.Cell object containing polygons to predict.

required
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.

required
model_type str

The type of model to use ('p' for prediction, 'c' for correction).

required
gds_layer tuple[int, int]

The layer and datatype to use within the GDSTK cell. Defaults to (1, 0).

(1, 0)
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 binarized using a threshold value of eta.

0.5

Returns:

Type Description
Cell

A new gdstk cell containing the predicted polygons. For multi-level predictions, each level's polygons will be placed on a different layer: - Level 0: (layer, 99) - Level 1: (layer, 100)

Raises:

Type Description
ValueError

If no polygons are found in the specified layer.

Source code in prefab/predict.py
def predict_gdstk(
    gdstk_cell: gdstk.Cell,
    model: Model,
    model_type: str,
    gds_layer: tuple[int, int] = (1, 0),
    eta: float = 0.5,
) -> gdstk.Cell:
    """
    Predict the nanofabrication outcome of a gdstk cell using a specified model.

    This function extracts polygons from a gdstk cell, sends them to the prediction
    server, and returns a new cell containing the predicted polygons.

    Parameters
    ----------
    gdstk_cell : gdstk.Cell
        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.
    model_type : str
        The type of model to use ('p' for prediction, 'c' for correction).
    gds_layer : tuple[int, int]
        The layer and datatype to use within the GDSTK cell. Defaults to (1, 0).
    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
        binarized using a threshold value of eta.

    Returns
    -------
    gdstk.Cell
        A new gdstk cell containing the predicted polygons. For multi-level
        predictions, each level's polygons will be placed on a different layer:
        - Level 0: (layer, 99)
        - Level 1: (layer, 100)

    Raises
    ------
    ValueError
        If no polygons are found in the specified layer.
    """
    polygons = gdstk_cell.get_polygons(layer=gds_layer[0], datatype=gds_layer[1])
    if not polygons:
        raise ValueError("No polygons found in the specified layer")

    polygon_points = [polygon.points.tolist() for polygon in polygons]

    predicted_polygon_data = _predict_poly(
        polygon_points=polygon_points,
        model=model,
        model_type=model_type,
        eta=eta,
    )

    result_cell = gdstk.Cell(f"{gdstk_cell.name}_predicted")

    polygons_by_channel = {}
    for polygon_data in predicted_polygon_data:
        channel = polygon_data.get("channel", 0)
        points = polygon_data.get("points", [])

        if channel not in polygons_by_channel:
            polygons_by_channel[channel] = []

        polygons_by_channel[channel].append(points)

    for channel, points_list in polygons_by_channel.items():
        layer = gds_layer[0]
        datatype = 99 + channel

        for points in points_list:
            points_array = np.array(points)
            polygon = gdstk.Polygon(points_array, layer=layer, datatype=datatype)
            result_cell.add(polygon)

    return result_cell