prefab 1.4.2__py3-none-any.whl → 1.6.0__py3-none-any.whl

prefab/__init__.py

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

prefab/__main__.py

@@ -18,8 +18,8 @@ def store_jwt(jwt, refresh_token):
     print(
         f"Token successfully stored in {prefab_file_path}.\n\n"
         "🎉 Welcome to PreFab.\n"
-        "See our examples at https://docs.prefabphotonics.com/examples to start.\n"
-        "Reach out to us at hi@prefabphotonics.com for support."
+        "See our examples at https://docs.prefabphotonics.com to start.\n"
+        "Reach out to us at support@prefabphotonics.com for support."
     )
 
 

prefab/predict.py

@@ -3,12 +3,13 @@ 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.
+(ndarrays, polygons, GDSII files) and model types (prediction and correction). Gradient
+computation is available for inverse design applications using automatic
+differentiation.
 """
 
 import base64
+import gzip
 import io
 import json
 import os
@@ -26,8 +27,17 @@ from PIL import Image
 from .geometry import binarize_hard
 from .models import Model
 
-BASE_ENDPOINT_URL = "https://prefab-photonics--predict"
-ENDPOINT_VERSION = "3"
+BASE_URL = "https://prefab-photonics"
+
+# Endpoint versions
+PREDICT_VERSION = "3"
+PREDICT_POLY_VERSION = "3"
+VJP_VERSION = "3"
+
+# Endpoint URLs
+PREDICT_ENDPOINT = f"{BASE_URL}--predict-v{PREDICT_VERSION}.modal.run"
+PREDICT_POLY_ENDPOINT = f"{BASE_URL}--predict-poly-v{PREDICT_POLY_VERSION}.modal.run"
+VJP_ENDPOINT = f"{BASE_URL}--vjp-v{VJP_VERSION}.modal.run"
 
 
 def predict_array(
@@ -75,7 +85,7 @@ def predict_array(
     ValueError
         If the server returns an error or invalid response.
     """
-    endpoint_url = f"{BASE_ENDPOINT_URL}-v{ENDPOINT_VERSION}.modal.run"
+    endpoint_url = PREDICT_ENDPOINT
     predict_data = {
         "device_array": _encode_array(np.squeeze(device_array)),
         "model": model.to_json(),
@@ -167,7 +177,7 @@ def _predict_poly(
         "eta": eta,
     }
 
-    endpoint_url = f"{BASE_ENDPOINT_URL}-poly-v{ENDPOINT_VERSION}.modal.run"
+    endpoint_url = PREDICT_POLY_ENDPOINT
     headers = _prepare_headers()
 
     try:
@@ -364,47 +374,32 @@ def predict_gdstk(
     return result_cell
 
 
-def _predict_array_with_grad(
-    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.
-
-    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.
+# Storage for caching prediction results for VJP computation
+_diff_cache: dict[int, tuple[npt.NDArray[Any], Model]] = {}
 
-    Returns
-    -------
-    tuple[np.ndarray, np.ndarray]
-        The predicted output array and gradient array.
 
-    Raises
-    ------
-    RuntimeError
-        If the request to the prediction service fails.
-    ValueError
-        If the server returns an error or invalid response.
-    """
+def _compute_vjp(
+    device_array: npt.NDArray[Any],
+    upstream_gradient: npt.NDArray[Any],
+    model: Model,
+) -> npt.NDArray[Any]:
+    """Compute J.T @ upstream_gradient via the server-side VJP endpoint."""
     headers = _prepare_headers()
-    predict_data = {
+    upstream_arr = np.squeeze(upstream_gradient).astype(np.float32)
+    vjp_data = {
         "device_array": _encode_array(np.squeeze(device_array)),
+        "upstream_gradient": base64.b64encode(
+            gzip.compress(upstream_arr.tobytes(), compresslevel=1)
+        ).decode("utf-8"),
+        "upstream_gradient_shape": list(upstream_arr.shape),
         "model": model.to_json(),
         "model_type": "p",
-        "binary": False,
     }
-    endpoint_url = f"{BASE_ENDPOINT_URL}-with-grad-v{ENDPOINT_VERSION}.modal.run"
+    endpoint_url = VJP_ENDPOINT
 
     try:
         response = requests.post(
-            endpoint_url, data=json.dumps(predict_data), headers=headers
+            endpoint_url, data=json.dumps(vjp_data), headers=headers
         )
         response.raise_for_status()
 
@@ -414,89 +409,152 @@ def _predict_array_with_grad(
         response_data = response.json()
 
         if "error" in response_data:
-            raise ValueError(f"Prediction error: {response_data['error']}")
+            raise ValueError(f"VJP error: {response_data['error']}")
 
-        prediction_array = _decode_array(response_data["prediction_array"])
-        gradient_array = _decode_array(response_data["gradient_array"])
-        gradient_min = response_data["gradient_min"]
-        gradient_max = response_data["gradient_max"]
-        gradient_range = gradient_max - gradient_min
-        gradient_array = gradient_array * gradient_range + gradient_min
-        return (prediction_array, gradient_array)
+        vjp_array = _decode_array(response_data["vjp_array"])
+        vjp_min = response_data["vjp_min"]
+        vjp_max = response_data["vjp_max"]
+        vjp_range = vjp_max - vjp_min
+        vjp_array = vjp_array * vjp_range + vjp_min
+        return vjp_array
 
     except requests.exceptions.RequestException as e:
-        raise RuntimeError(f"Request failed: {e}") from e
+        raise RuntimeError(f"VJP request failed: {e}") from e
     except json.JSONDecodeError as e:
         raise ValueError(f"JSON decode error: {e}") from e
 
 
 @primitive
-def predict_array_with_grad(
+def predict_array_diff(
     device_array: npt.NDArray[Any], model: Model
 ) -> npt.NDArray[Any]:
     """
-    Predict the nanofabrication outcome of a device array and compute its gradient.
+    Differentiable fab prediction with exact gradient support.
 
-    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).
+    Compatible with autograd for automatic differentiation. Gradients are computed via
+    a server-side VJP endpoint during the backward pass.
 
     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.
+        The model to use for prediction.
 
     Returns
     -------
     np.ndarray
-        The predicted output array.
-
-    Raises
-    ------
-    RuntimeError
-        If the request to the prediction service fails.
-    ValueError
-        If the server returns an error or invalid response.
+        The predicted fabrication outcome array.
     """
-    prediction_array, gradient_array = _predict_array_with_grad(
-        device_array=device_array, model=model
+    # Use standard forward pass
+    prediction_array = predict_array(
+        device_array=device_array,
+        model=model,
+        model_type="p",
+        binarize=False,
     )
-    predict_array_with_grad.gradient_array = gradient_array  # pyright: ignore[reportFunctionMemberAccess]
+
+    # Cache the input for VJP computation during backward pass
+    _diff_cache[id(prediction_array)] = (device_array.copy(), model)
+
     return prediction_array
 
 
-def predict_array_with_grad_vjp(
-    ans: npt.NDArray[Any], device_array: npt.NDArray[Any], *args: Any
+def _predict_array_diff_vjp(
+    ans: npt.NDArray[Any], device_array: npt.NDArray[Any], model: Model
 ) -> Any:
+    """Define the exact VJP for predict_array_diff using server-side computation."""
+    cache_key = id(ans)
+    cached_device_array, cached_model = _diff_cache.get(
+        cache_key, (device_array, model)
+    )
+
+    def vjp(g: npt.NDArray[Any]) -> tuple[npt.NDArray[Any], None]:
+        # Compute exact VJP: J.T @ g via server endpoint
+        vjp_result = _compute_vjp(
+            device_array=cached_device_array,
+            upstream_gradient=g,
+            model=cached_model,
+        )
+        # Clean up cache
+        _diff_cache.pop(cache_key, None)
+        # Ensure gradient shape matches input shape
+        vjp_result = vjp_result.reshape(cached_device_array.shape)
+        # Return gradient for device_array, None for model (not differentiable)
+        return (vjp_result, None)
+
+    return vjp
+
+
+defvjp(predict_array_diff, _predict_array_diff_vjp)
+
+
+# Alias for backward compatibility with existing code
+predict_array_with_grad = predict_array_diff
+"""Alias for predict_array_diff. Deprecated, use predict_array_diff directly."""
+
+
+def differentiable(model: Model):
     """
-    Define the vector-Jacobian product (VJP) for the prediction function.
+    Create a model-bound differentiable predictor for clean autograd integration.
+
+    Returns a function that takes only `device_array` as input, enabling seamless
+    composition with other differentiable functions. The VJP returns a single
+    gradient array (not a tuple), making it compatible with standard autograd workflows.
 
     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.
+    model : Model
+        The model to use for prediction.
 
     Returns
     -------
-    function
-        A function that computes the VJP given an upstream gradient `g`.
+    callable
+        A differentiable prediction function that takes `device_array` and returns
+        the predicted fabrication outcome.
+
+    Examples
+    --------
+    >>> predictor = pf.predict.differentiable(model)
+    >>> def loss_fn(x):
+    ...     pred = predictor(x)
+    ...     return np.mean((pred - target) ** 2)
+    >>> gradient = grad(loss_fn)(device_array)  # Returns array, not tuple
     """
-    grad_x = predict_array_with_grad.gradient_array  # pyright: ignore[reportFunctionMemberAccess]
 
-    def vjp(g: npt.NDArray[Any]) -> npt.NDArray[Any]:
-        return g * grad_x  # type: ignore[no-any-return]
+    @primitive
+    def predict(device_array: npt.NDArray[Any]) -> npt.NDArray[Any]:
+        prediction = predict_array(
+            device_array=device_array,
+            model=model,
+            model_type="p",
+            binarize=False,
+        )
+        _diff_cache[id(prediction)] = (device_array.copy(), model)
+        return prediction
+
+    def predict_vjp(
+        ans: npt.NDArray[Any], device_array: npt.NDArray[Any]
+    ) -> Any:
+        cache_key = id(ans)
+        cached_device_array, cached_model = _diff_cache.get(
+            cache_key, (device_array, model)
+        )
 
-    return vjp
+        def vjp(g: npt.NDArray[Any]) -> npt.NDArray[Any]:
+            vjp_result = _compute_vjp(
+                device_array=cached_device_array,
+                upstream_gradient=g,
+                model=cached_model,
+            )
+            _diff_cache.pop(cache_key, None)
+            # Ensure gradient shape matches input shape
+            return vjp_result.reshape(device_array.shape)
 
+        return vjp
 
-defvjp(predict_array_with_grad, predict_array_with_grad_vjp)
+    defvjp(predict, predict_vjp)
+    return predict
 
 
 def _encode_array(array: npt.NDArray[Any]) -> str:
@@ -512,7 +570,7 @@ 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  # type: ignore[no-any-return]
+    return np.array(image) / 255
 
 
 def _prepare_headers() -> dict[str, str]:

{prefab-1.4.2.dist-info → prefab-1.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prefab
-Version: 1.4.2
+Version: 1.6.0
 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
@@ -511,14 +511,13 @@ License:                   GNU LESSER GENERAL PUBLIC LICENSE
         
         That's all there is to it!
 License-File: LICENSE
-Keywords: computer-vision,deep-learning,electronic-design-automation,integrated-photonics,machine-learning,nanofabrication,semiconductor-manufacturing
+Keywords: electronic-design-automation,integrated-photonics,inverse-design,lithography,nanofabrication,photonics,semiconductor-manufacturing,silicon-photonics
 Classifier: License :: OSI Approved :: GNU Lesser General Public License v2 (LGPLv2)
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.9
 Requires-Dist: autograd
 Requires-Dist: gdstk>=0.9.55
-Requires-Dist: ipykernel>=6.31.0
 Requires-Dist: matplotlib
 Requires-Dist: numpy
 Requires-Dist: opencv-python-headless
@@ -528,9 +527,6 @@ Requires-Dist: requests
 Requires-Dist: scikit-image
 Requires-Dist: scipy
 Requires-Dist: toml
-Provides-Extra: dev
-Requires-Dist: jupyter; extra == 'dev'
-Requires-Dist: jupyterlab; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: mkdocs; extra == 'docs'
 Requires-Dist: mkdocs-material; extra == 'docs'

{prefab-1.4.2.dist-info → prefab-1.6.0.dist-info}/RECORD

@@ -1,17 +1,17 @@
-prefab/__init__.py,sha256=XffqgQ1AexVSwGAgt7Wbfo-4raEefrzmXwE2EEG3WUI,425
-prefab/__main__.py,sha256=M5ulRfgSMdPjeHPoMGQ3P86zFMJDII5I18NGrpFUcsM,3458
+prefab/__init__.py,sha256=WmddNmU7gmyH4cqkT6BW4iQ3RDG_l5SpKxg4poS6wsU,425
+prefab/__main__.py,sha256=uL6AdCeimPbXiWv0gnq9TDeZhAIrxGwJsEKNYbg9MZg,3454
 prefab/compare.py,sha256=aX7nr9tznSebYeeztvqIPz57npnJ4-iUeKEedrZdksE,3676
 prefab/device.py,sha256=1O6vTOq4wQRGVYvFWLH0uj1XhhYCfnDnIapDEYnBKHw,47996
 prefab/geometry.py,sha256=-nTaGjdw3KN1SVoyvqdcrE2GJP7OqPF6ivUhrO78rUk,11244
-prefab/predict.py,sha256=LSOf-lo7l8JRzRSXTDWosGfr2baE-D5zZjinkObeJV4,19182
+prefab/predict.py,sha256=gxxHCcaMv2TE5O639mQzO3IzFRsLGjY-FQowVp5vLkI,20659
 prefab/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 prefab/read.py,sha256=5BmvFemobA72urxs4j2VZRVvanZZGu1mDB1Uh-epyvI,8635
 prefab/shapes.py,sha256=mRGwsPS-A9XsW3jgvUuMCEeNv9BLXsEnJkytlIHUAKE,28802
 prefab/models/__init__.py,sha256=rRrjcOcHPcob98Coksc0tbvYcXbm6SoLEb-Md233Jvo,1391
 prefab/models/base.py,sha256=t4VNMsOztPedj3kN5fZ1-4tk0SRHWrMuqnIVHztsCs4,1514
 prefab/models/evaluation.py,sha256=2_Klui6tY8xPvOSVD8VpZCVAnT1RX15FONqWG-_x-J8,484
-prefab-1.4.2.dist-info/METADATA,sha256=JFb3uFdmPuf8thzsPnOim8F47qqSjVajVG9plmBuonQ,33880
-prefab-1.4.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-prefab-1.4.2.dist-info/entry_points.txt,sha256=h1_A9O9F3NAIoKXD1RPb3Eo-WCSiHhMB_AnagBi6XTQ,48
-prefab-1.4.2.dist-info/licenses/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
-prefab-1.4.2.dist-info/RECORD,,
+prefab-1.6.0.dist-info/METADATA,sha256=CvuP1Khblh_LA-vx5NLgTH5A_zaewTkklkaPEsIY4hI,33754
+prefab-1.6.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+prefab-1.6.0.dist-info/entry_points.txt,sha256=h1_A9O9F3NAIoKXD1RPb3Eo-WCSiHhMB_AnagBi6XTQ,48
+prefab-1.6.0.dist-info/licenses/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
+prefab-1.6.0.dist-info/RECORD,,