prefab 1.1.5__py3-none-any.whl → 1.1.6__py3-none-any.whl

prefab/geometry.py

@@ -1,8 +1,9 @@
 """Provides functions for manipulating ndarrays of device geometries."""
 
+from typing import Optional
+
 import cv2
 import numpy as np
-from skimage.morphology import closing, disk, opening, square
 
 
 def normalize(device_array: np.ndarray) -> np.ndarray:
@@ -34,9 +35,9 @@ def binarize(
     ----------
     device_array : np.ndarray
         The input array to be binarized.
-    eta : float, optional
+    eta : float
         The threshold value for binarization. Defaults to 0.5.
-    beta : float, optional
+    beta : float
         The scaling factor for the binarization process. A higher value makes the
         transition sharper. Defaults to np.inf, which results in a hard threshold.
 
@@ -60,7 +61,7 @@ def binarize_hard(device_array: np.ndarray, eta: float = 0.5) -> np.ndarray:
     ----------
     device_array : np.ndarray
         The input array to be binarized.
-    eta : float, optional
+    eta : float
         The threshold value for binarization. Defaults to 0.5.
 
     Returns
@@ -95,11 +96,12 @@ def binarize_sem(sem_array: np.ndarray) -> np.ndarray:
 
 def binarize_monte_carlo(
     device_array: np.ndarray,
-    threshold_noise_std: float,
-    threshold_blur_std: float,
+    noise_magnitude: float,
+    blur_radius: float,
 ) -> np.ndarray:
     """
-    Binarize the input ndarray using a Monte Carlo approach with Gaussian blurring.
+    Binarize the input ndarray using a dynamic thresholding approach to simulate surface
+    roughness.
 
     This function applies a dynamic thresholding technique where the threshold value is
     determined by a base value perturbed by Gaussian-distributed random noise. The
@@ -107,14 +109,19 @@ def binarize_monte_carlo(
     simulating a potentially more realistic scenario where the threshold is not uniform
     across the device.
 
+    Notes
+    -----
+    This is a temporary solution, where the defaults are chosen based on what looks
+    good. A better, data-driven approach is needed.
+
     Parameters
     ----------
     device_array : np.ndarray
         The input array to be binarized.
-    threshold_noise_std : float
+    noise_magnitude : float
         The standard deviation of the Gaussian distribution used to generate noise for
         the threshold values. This controls the amount of randomness in the threshold.
-    threshold_blur_std : float
+    blur_radius : float
         The standard deviation for the Gaussian kernel used in blurring the threshold
         map. This controls the spatial variation of the threshold across the array.
 
@@ -127,10 +134,10 @@ def binarize_monte_carlo(
     device_array = np.squeeze(device_array)
     base_threshold = np.random.normal(loc=0.5, scale=0.1)
     threshold_noise = np.random.normal(
-        loc=0, scale=threshold_noise_std, size=device_array.shape
+        loc=0, scale=noise_magnitude, size=device_array.shape
     )
     spatial_threshold = cv2.GaussianBlur(
-        threshold_noise, ksize=(0, 0), sigmaX=threshold_blur_std
+        threshold_noise, ksize=(0, 0), sigmaX=blur_radius
     )
     dynamic_threshold = base_threshold + spatial_threshold
     binarized_array = np.where(device_array < dynamic_threshold, 0.0, 1.0)
@@ -149,9 +156,9 @@ def ternarize(
     ----------
     device_array : np.ndarray
         The input array to be ternarized.
-    eta1 : float, optional
+    eta1 : float
         The first threshold value for ternarization. Defaults to 1/3.
-    eta2 : float, optional
+    eta2 : float
         The second threshold value for ternarization. Defaults to 2/3.
 
     Returns
@@ -162,7 +169,9 @@ def ternarize(
     return np.where(device_array < eta1, 0.0, np.where(device_array >= eta2, 1.0, 0.5))
 
 
-def trim(device_array: np.ndarray, buffer_thickness: dict = None) -> np.ndarray:
+def trim(
+    device_array: np.ndarray, buffer_thickness: Optional[dict[str, int]] = None
+) -> np.ndarray:
     """
     Trim the input ndarray by removing rows and columns that are completely zero.
 
@@ -170,7 +179,7 @@ def trim(device_array: np.ndarray, buffer_thickness: dict = None) -> np.ndarray:
     ----------
     device_array : np.ndarray
         The input array to be trimmed.
-    buffer_thickness : dict, optional
+    buffer_thickness : Optional[dict[str, int]]
         A dictionary specifying the thickness of the buffer to leave around the non-zero
         elements of the array. Should contain keys 'top', 'bottom', 'left', 'right'.
         Defaults to None, which means no buffer is added.
@@ -200,6 +209,30 @@ def trim(device_array: np.ndarray, buffer_thickness: dict = None) -> np.ndarray:
     ]
 
 
+def pad(device_array: np.ndarray, pad_width: int) -> np.ndarray:
+    """
+    Pad the input ndarray uniformly with a specified width on all sides.
+
+    Parameters
+    ----------
+    device_array : np.ndarray
+        The input array to be padded.
+    pad_width : int
+        The number of pixels to pad on each side.
+
+    Returns
+    -------
+    np.ndarray
+        The padded array.
+    """
+    return np.pad(
+        device_array,
+        pad_width=((pad_width, pad_width), (pad_width, pad_width), (0, 0)),
+        mode="constant",
+        constant_values=0,
+    )
+
+
 def blur(device_array: np.ndarray, sigma: float = 1.0) -> np.ndarray:
     """
     Apply Gaussian blur to the input ndarray and normalize the result.
@@ -208,7 +241,7 @@ def blur(device_array: np.ndarray, sigma: float = 1.0) -> np.ndarray:
     ----------
     device_array : np.ndarray
         The input array to be blurred.
-    sigma : float, optional
+    sigma : float
         The standard deviation for the Gaussian kernel. This controls the amount of
         blurring. Defaults to 1.0.
 
@@ -318,13 +351,18 @@ def enforce_feature_size(
     device geometry are at least the specified minimum size. It uses either a disk
     or square structuring element for the operations.
 
+    Notes
+    -----
+    This function does not guarantee that the minimum feature size is enforced in all
+    cases. A better process is needed.
+
     Parameters
     ----------
     device_array : np.ndarray
         The input array representing the device geometry.
     min_feature_size : int
         The minimum feature size to enforce, in nanometers.
-    strel : str, optional
+    strel : str
         The type of structuring element to use. Can be either "disk" or "square".
         Defaults to "disk".
 
@@ -339,12 +377,18 @@ def enforce_feature_size(
         If an invalid structuring element type is specified.
     """
     if strel == "disk":
-        structuring_element = disk(radius=min_feature_size / 2)
+        kernel = cv2.getStructuringElement(
+            cv2.MORPH_ELLIPSE, (min_feature_size, min_feature_size)
+        )
     elif strel == "square":
-        structuring_element = square(width=min_feature_size)
+        kernel = cv2.getStructuringElement(
+            cv2.MORPH_RECT, (min_feature_size, min_feature_size)
+        )
     else:
         raise ValueError(f"Invalid structuring element: {strel}")
 
-    modified_geometry = closing(device_array[:, :, 0], structuring_element)
-    modified_geometry = opening(modified_geometry, structuring_element)
-    return np.expand_dims(modified_geometry, axis=-1)
+    device_array_2d = (device_array[:, :, 0] * 255).astype(np.uint8)
+    modified_geometry = cv2.morphologyEx(device_array_2d, cv2.MORPH_CLOSE, kernel)
+    modified_geometry = cv2.morphologyEx(modified_geometry, cv2.MORPH_OPEN, kernel)
+
+    return np.expand_dims(modified_geometry.astype(float) / 255, axis=-1)

prefab/predict.py

@@ -1,4 +1,4 @@
-"""Provides prediction functions for ndarrays of device geometries."""
+"""Prediction functions for ndarrays of device geometries."""
 
 import base64
 import io
@@ -16,7 +16,8 @@ from tqdm import tqdm
 from .geometry import binarize_hard
 from .models import Model
 
-BASE_URL = "https://prefab-photonics--predict"
+BASE_ENDPOINT_URL = "https://prefab-photonics--predict"
+ENDPOINT_VERSION = 1
 
 
 def predict_array(
@@ -46,15 +47,16 @@ def predict_array(
         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').
+        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, optional
+    gpu : bool
         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.
+        GPU option has more startup overhead and will take longer for small devices, but
+        will be faster for larger devices.
 
     Returns
     -------
@@ -68,7 +70,11 @@ def predict_array(
     """
     headers = _prepare_headers()
     predict_data = _prepare_predict_data(device_array, model, model_type, binarize)
-    endpoint_url = f"{BASE_URL}-gpu-v1.modal.run" if gpu else f"{BASE_URL}-v1.modal.run"
+    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(
@@ -78,7 +84,10 @@ def predict_array(
             stream=True,
         ) as response:
             response.raise_for_status()
-            return _process_response(response, model_type, binarize)
+            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
 
@@ -86,9 +95,33 @@ def predict_array(
 def _predict_array_with_grad(
     device_array: np.ndarray, model: Model
 ) -> tuple[np.ndarray, 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.
+
+    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
+        A 2D array representing the planar geometry of the device.
+    model : Model
+        The model to use for prediction, representing a specific fabrication process.
+
+    Returns
+    -------
+    tuple[np.ndarray, np.ndarray]
+        The predicted output array and gradient array.
+    """
     headers = _prepare_headers()
     predict_data = _prepare_predict_data(device_array, model, "p", False)
-    endpoint_url = f"{BASE_URL}-with-grad-v1.modal.run"
+    endpoint_url = f"{BASE_ENDPOINT_URL}-with-grad-v{ENDPOINT_VERSION}.modal.run"
 
     response = requests.post(
         endpoint_url, data=json.dumps(predict_data), headers=headers
@@ -99,7 +132,6 @@ def _predict_array_with_grad(
     gradient_max = response.json()["gradient_max"]
     gradient_range = gradient_max - gradient_min
     gradient_array = gradient_array * gradient_range + gradient_min
-
     return (prediction_array, gradient_array)
 
 
@@ -110,7 +142,8 @@ def predict_array_with_grad(device_array: np.ndarray, model: Model) -> np.ndarra
 
     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.
+    prediction with respect to the input device array, making it suitable for use in
+    automatic differentiation applications (e.g., autograd).
 
     Parameters
     ----------
@@ -127,32 +160,29 @@ 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
+    predict_array_with_grad.gradient_array = gradient_array  # type: ignore
     return prediction_array
 
 
-def predict_array_with_grad_vjp(ans: np.ndarray, x: np.ndarray, model: Model):
+def predict_array_with_grad_vjp(ans: np.ndarray, device_array: np.ndarray, *args):
     """
     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
+    device_array : np.ndarray
         The input device array for which the gradient is computed.
-    model : Model
-        The model used for prediction.
+    *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
+    grad_x = predict_array_with_grad.gradient_array  # type: ignore
 
     def vjp(g: np.ndarray) -> np.ndarray:
         return g * grad_x
@@ -195,7 +225,7 @@ def _read_tokens():
             "Could not validate user.\n"
             "Please update prefab using: pip install --upgrade prefab.\n"
             "Signup/login and generate a new token.\n"
-            "See https://www.prefabphotonics.com/docs/guides/quickstart."
+            "See https://docs.prefabphotonics.com/."
         ) from None