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, x)

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

This function provides the VJP for the predict_array_with_grad function, which is used in reverse-mode automatic differentiation to compute gradients.

Parameters:

Name	Type	Description	Default
ans	ndarray	The output of the predict_array_with_grad function.	required
x	ndarray	The input device array for which the gradient is computed.	required

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, x: np.ndarray):
    """
    Define the vector-Jacobian product (VJP) for the prediction function.

    This function provides the VJP for the `predict_array_with_grad` function,
    which is used in reverse-mode automatic differentiation to compute gradients.

    Parameters
    ----------
    ans : np.ndarray
        The output of the `predict_array_with_grad` function.
    x : np.ndarray
        The input device array for which the gradient is computed.

    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