prefab 1.1.3__tar.gz → 1.1.4__tar.gz

{prefab-1.1.3 → prefab-1.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: prefab
-Version: 1.1.3
+Version: 1.1.4
 Summary: Artificial nanofabrication of integrated photonic circuits using deep learning
 Project-URL: Homepage, https://prefabphotonics.com
 Project-URL: Repository, https://github.com/PreFab-Photonics/PreFab

{prefab-1.1.3 → prefab-1.1.4}/docs/CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## [1.1.4](https://github.com/PreFab-Photonics/PreFab/releases/tag/v1.1.4) - 2024-11-01
+
+- Added custom vector-Jacobian product (VJP) for the `predict.predict_array_with_grad` function.
+- Changed some of the docstrings in `prefab.predict` to be more consistent and clear.
+
 ## [1.1.3](https://github.com/PreFab-Photonics/PreFab/releases/tag/v1.1.3) - 2024-10-26
 
 - Moved prediction logic to `prefab.predict` module.
@@ -19,7 +24,7 @@
 
 ## [1.1.1](https://github.com/PreFab-Photonics/PreFab/releases/tag/v1.1.1) - 2024-09-24
 
-- Manually adding small random noise to the semulated images to better match the real data. This is ideally included in the model training, but for now this is a quick fix.
+- Manually adding small random noise to the "SEMulated" images to better match the real data. This is ideally included in the model training, but for now this is a quick fix.
 - Added z-padding to the device array before exporting to STL with `Device.to_stl` to ensure that the exported device is closed.
 - Removed buffer from `Device.device_array` before exporting to with `Device.to_gdsfactory`.
 - The additions from `1.0.3` and `1.0.4` releases, which should be considered part of this release. Release planning a work in progress.

prefab-1.1.4/docs/reference/predict.md

@@ -0,0 +1,3 @@
+# Predict
+
+::: prefab.predict

{prefab-1.1.3 → prefab-1.1.4}/mkdocs.yml

@@ -21,6 +21,7 @@ nav:
     - Device: reference/device.md
     - Geometry: reference/geometry.md
     - Models: reference/models.md
+    - Predict: reference/predict.md
     - Read: reference/read.md
     - Shapes: reference/shapes.md
 

{prefab-1.1.3 → prefab-1.1.4}/prefab/__init__.py

@@ -5,7 +5,7 @@ Usage:
     import prefab as pf
 """
 
-__version__ = "1.1.3"
+__version__ = "1.1.4"
 
 from . import compare, geometry, predict, read, shapes
 from .device import BufferSpec, Device

{prefab-1.1.3 → prefab-1.1.4}/prefab/predict.py

@@ -1,3 +1,5 @@
+"""Provides prediction functions for ndarrays of device geometries."""
+
 import base64
 import io
 import json
@@ -6,6 +8,8 @@ import os
 import numpy as np
 import requests
 import toml
+from autograd import primitive
+from autograd.extend import defvjp
 from PIL import Image
 from tqdm import tqdm
 
@@ -23,24 +27,34 @@ def predict_array(
     gpu: bool = False,
 ) -> np.ndarray:
     """
-    Predicts the output array for a given device array using a specified model.
+    Predict the nanofabrication outcome of a device array using a specified model.
 
-    This function sends the device array to a prediction service, which uses a machine
-    learning model to predict the outcome of the nanofabrication process. The prediction
-    can be performed on a GPU if specified.
+    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
-        The input device array to be predicted.
+        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.
+        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', 'c', 's').
     binarize : bool
-        Whether to binarize the output.
+        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, optional
-        Whether to use GPU for prediction. Defaults to False.
+        If True, the prediction will be performed on a GPU. Defaults to False. Note: The
+        GPU option has more overhead and will take longer for small devices, but will be
+        faster for larger devices.
 
     Returns
     -------
@@ -69,38 +83,11 @@ def predict_array(
         raise RuntimeError(f"Request failed: {e}") from e
 
 
-def predict_array_with_grad(
-    device_array: np.ndarray, model: Model, model_type: str
+def _predict_array_with_grad(
+    device_array: np.ndarray, model: Model
 ) -> tuple[np.ndarray, np.ndarray]:
-    """
-    Predicts the output array and its gradient for a given device array using a
-    specified model.
-
-    This function sends the device array to a prediction service, which uses a machine
-    learning model to predict both the outcome and the gradient of the nanofabrication
-    process.
-
-    Parameters
-    ----------
-    device_array : np.ndarray
-        The input device array to be predicted.
-    model : Model
-        The model to use for prediction.
-    model_type : str
-        The type of model to use (e.g., 'p', 'c', 's').
-
-    Returns
-    -------
-    tuple[np.ndarray, np.ndarray]
-        A tuple containing the predicted output array and its gradient.
-
-    Raises
-    ------
-    RuntimeError
-        If the request to the prediction service fails.
-    """
     headers = _prepare_headers()
-    predict_data = _prepare_predict_data(device_array, model, model_type, False)
+    predict_data = _prepare_predict_data(device_array, model, "p", False)
     endpoint_url = f"{BASE_URL}-with-grad-v1.modal.run"
 
     response = requests.post(
@@ -116,8 +103,68 @@ def predict_array_with_grad(
     return (prediction_array, gradient_array)
 
 
+@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.
+
+    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
+    return prediction_array
+
+
+def predict_array_with_grad_vjp(ans: np.ndarray, x: np.ndarray, model: Model):
+    """
+    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.
+    model : Model
+        The model used for prediction.
+
+    Returns
+    -------
+    function
+        A function that computes the VJP given an upstream gradient `g`.
+    """
+    grad_x = predict_array_with_grad.gradient_array
+
+    def vjp(g: np.ndarray) -> np.ndarray:
+        return g * grad_x
+
+    return vjp
+
+
+defvjp(predict_array_with_grad, predict_array_with_grad_vjp)
+
+
 def _encode_array(array):
-    """Encode a numpy array as a PNG image and return the base64 encoded string."""
+    """Encode an ndarray as a base64 encoded image for transmission."""
     image = Image.fromarray(np.uint8(array * 255))
     buffered = io.BytesIO()
     image.save(buffered, format="PNG")
@@ -126,7 +173,7 @@ def _encode_array(array):
 
 
 def _decode_array(encoded_png):
-    """Decode a base64 encoded PNG image and return a numpy array."""
+    """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

{prefab-1.1.3 → prefab-1.1.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "prefab"
-version = "1.1.3"
+version = "1.1.4"
 description = "Artificial nanofabrication of integrated photonic circuits using deep learning"
 authors = [{ name = "Dusan Gostimirovic", email = "dusan@prefabphotonics.com" }]
 keywords = ["photonics", "nanofabrication", "machine-learning"]