prefab 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

prefab/__init__.py

@@ -5,7 +5,9 @@ Usage:
     import prefab as pf
 """
 
-from . import compare, geometry, read
+__version__ = "1.0.4"
+
+from . import compare, geometry, read, shapes
 from .device import BufferSpec, Device
 from .models import models
 
@@ -14,6 +16,8 @@ __all__ = [
     "BufferSpec",
     "geometry",
     "read",
+    "shapes",
     "compare",
     "models",
+    "__version__",
 ]

prefab/compare.py

@@ -7,40 +7,40 @@ from .device import Device
 
 def mean_squared_error(device_a: Device, device_b: Device) -> float:
     """
-    Calculate the mean squared error (MSE) between two non-binarized devices.
+    Calculate the mean squared error (MSE) between two non-binarized devices. A lower
+    value indicates more similarity.
 
     Parameters
     ----------
     device_a : Device
-        The first device.
+        The first device (non-binarized).
     device_b : Device
-        The second device.
+        The second device (non-binarized).
 
     Returns
     -------
     float
-        The mean squared error between two devices. A lower value indicates more
-        similarity.
+        The mean squared error between two devices.
     """
     return np.mean((device_a.device_array - device_b.device_array) ** 2)
 
 
 def intersection_over_union(device_a: Device, device_b: Device) -> float:
     """
-    Calculates the Intersection over Union (IoU) between two binary devices.
+    Calculates the Intersection over Union (IoU) between two binary devices. A value
+    closer to 1 indicates more similarity (more overlap).
 
     Parameters
     ----------
     device_a : Device
-        The first device.
+        The first device (binarized).
     device_b : Device
-        The second device.
+        The second device (binarized).
 
     Returns
     -------
     float
-        The Intersection over Union between two devices. A value closer to 1 indicates
-        more similarity (more overlap).
+        The Intersection over Union between two devices.
     """
     return np.sum(
         np.logical_and(device_a.device_array, device_b.device_array)
@@ -49,40 +49,42 @@ def intersection_over_union(device_a: Device, device_b: Device) -> float:
 
 def hamming_distance(device_a: Device, device_b: Device) -> int:
     """
-    Calculates the Hamming distance between two binary devices.
+    Calculates the Hamming distance between two binary devices. A lower value indicates
+    more similarity. The Hamming distance is calculated as the number of positions at
+    which the corresponding pixels are different.
 
     Parameters
     ----------
     device_a : Device
-        The first device.
+        The first device (binarized).
     device_b : Device
-        The second device.
+        The second device (binarized).
 
     Returns
     -------
     int
-        The Hamming distance between two devices. A lower value indicates more
-        similarity.
+        The Hamming distance between two devices.
     """
     return np.sum(device_a.device_array != device_b.device_array)
 
 
 def dice_coefficient(device_a: Device, device_b: Device) -> float:
     """
-    Calculates the Dice coefficient between two binary devices.
+    Calculates the Dice coefficient between two binary devices. A value closer to 1
+    indicates more similarity. The Dice coefficient is calculated as twice the number of
+    pixels in common divided by the total number of pixels in the two devices.
 
     Parameters
     ----------
     device_a : Device
-        The first device.
+        The first device (binarized).
     device_b : Device
-        The second device.
+        The second device (binarized).
 
     Returns
     -------
     float
-        The Dice coefficient between two devices. A value closer to 1 indicates more
-        similarity.
+        The Dice coefficient between two devices.
     """
     intersection = 2.0 * np.sum(
         np.logical_and(device_a.device_array, device_b.device_array)

prefab/device.py

@@ -40,14 +40,16 @@ class BufferSpec(BaseModel):
         ('top', 'bottom', 'left', 'right'), where 'constant' is used for isolated
         structures and 'edge' is utilized for preserving the edge, such as for waveguide
         connections.
-    thickness : conint(gt=0)
-        The thickness of the buffer zone around the device. Must be greater than 0.
+    thickness : dict[str, conint(gt=0)]
+        A dictionary that defines the thickness of the buffer zone for each side of the
+        device ('top', 'bottom', 'left', 'right'). Each value must be greater than 0.
 
     Raises
     ------
     ValueError
         If any of the modes specified in the 'mode' dictionary are not one of the
-        allowed values ('constant', 'edge'). Or if the thickness is not greater than 0.
+        allowed values ('constant', 'edge'). Or if any of the thickness values are not
+        greater than 0.
 
     Example
     -------
@@ -60,7 +62,12 @@ class BufferSpec(BaseModel):
                 "left": "constant",
                 "right": "edge",
             },
-            thickness=150,
+            thickness={
+                "top": 150,
+                "bottom": 100,
+                "left": 200,
+                "right": 250,
+            },
         )
     """
 
@@ -72,7 +79,14 @@ class BufferSpec(BaseModel):
             "right": "constant",
         }
     )
-    thickness: conint(gt=0) = 128
+    thickness: dict[str, conint(gt=0)] = Field(
+        default_factory=lambda: {
+            "top": 128,
+            "bottom": 128,
+            "left": 128,
+            "right": 128,
+        }
+    )
 
     @validator("mode", pre=True)
     def check_mode(cls, v):
@@ -146,28 +160,26 @@ class Device(BaseModel):
 
         self.device_array = np.pad(
             self.device_array,
-            pad_width=((buffer_thickness, 0), (0, 0)),
+            pad_width=((buffer_thickness["top"], 0), (0, 0)),
             mode=buffer_mode["top"],
         )
         self.device_array = np.pad(
             self.device_array,
-            pad_width=((0, buffer_thickness), (0, 0)),
+            pad_width=((0, buffer_thickness["bottom"]), (0, 0)),
             mode=buffer_mode["bottom"],
         )
         self.device_array = np.pad(
             self.device_array,
-            pad_width=((0, 0), (buffer_thickness, 0)),
+            pad_width=((0, 0), (buffer_thickness["left"], 0)),
             mode=buffer_mode["left"],
         )
         self.device_array = np.pad(
             self.device_array,
-            pad_width=((0, 0), (0, buffer_thickness)),
+            pad_width=((0, 0), (0, buffer_thickness["right"])),
             mode=buffer_mode["right"],
         )
 
-        self.device_array = np.expand_dims(
-            self.device_array.astype(np.float32), axis=-1
-        )
+        self.device_array = np.expand_dims(self.device_array, axis=-1)
 
     @root_validator(pre=True)
     def check_device_array(cls, values):
@@ -212,6 +224,7 @@ class Device(BaseModel):
         model: Model,
         model_type: str,
         binarize: bool,
+        gpu: bool = False,
     ) -> "Device":
         try:
             with open(os.path.expanduser("~/.prefab.toml")) as file:
@@ -233,6 +246,7 @@ class Device(BaseModel):
                 "Signup/login and generate a new token.\n"
                 "See https://www.prefabphotonics.com/docs/guides/quickstart."
             ) from None
+
         headers = {
             "Authorization": f"Bearer {access_token}",
             "X-Refresh-Token": refresh_token,
@@ -246,84 +260,99 @@ class Device(BaseModel):
         }
         json_data = json.dumps(predict_data)
 
-        endpoint_url = "https://prefab-photonics--predict-v1.modal.run"
-
-        with requests.post(
-            endpoint_url, data=json_data, headers=headers, stream=True
-        ) as response:
-            response.raise_for_status()
-            event_type = None
-            model_descriptions = {"p": "Prediction", "c": "Correction", "s": "SEMulate"}
-            progress_bar = tqdm(
-                total=100,
-                desc=f"{model_descriptions[model_type]}",
-                unit="%",
-                colour="green",
-                bar_format="{l_bar}{bar:30}{r_bar}{bar:-10b}",
-            )
+        endpoint_url = (
+            "https://prefab-photonics--predict-gpu-v1.modal.run"
+            if gpu
+            else "https://prefab-photonics--predict-v1.modal.run"
+        )
 
-            for line in response.iter_lines():
-                if line:
-                    decoded_line = line.decode("utf-8").strip()
-                    if decoded_line.startswith("event:"):
-                        event_type = decoded_line.split(":")[1].strip()
-                    elif decoded_line.startswith("data:"):
-                        try:
-                            data_content = json.loads(decoded_line.split("data: ")[1])
-                            if event_type == "progress":
-                                progress = round(100 * data_content["progress"])
-                                progress_bar.update(progress - progress_bar.n)
-                            elif event_type == "result":
-                                results = []
-                                for key in sorted(data_content.keys()):
-                                    if key.startswith("result"):
-                                        decoded_image = self._decode_array(
-                                            data_content[key]
-                                        )
-                                        results.append(decoded_image)
+        try:
+            with requests.post(
+                endpoint_url, data=json_data, headers=headers, stream=True
+            ) as response:
+                response.raise_for_status()
+                event_type = None
+                model_descriptions = {
+                    "p": "Prediction",
+                    "c": "Correction",
+                    "s": "SEMulate",
+                }
+                progress_bar = tqdm(
+                    total=100,
+                    desc=f"{model_descriptions[model_type]}",
+                    unit="%",
+                    colour="green",
+                    bar_format="{l_bar}{bar:30}{r_bar}{bar:-10b}",
+                )
 
-                                if results:
-                                    prediction = np.stack(results, axis=-1)
-                                    if binarize:
-                                        prediction = geometry.binarize_hard(prediction)
+                for line in response.iter_lines():
+                    if line:
+                        decoded_line = line.decode("utf-8").strip()
+                        if decoded_line.startswith("event:"):
+                            event_type = decoded_line.split(":")[1].strip()
+                        elif decoded_line.startswith("data:"):
+                            try:
+                                data_content = json.loads(
+                                    decoded_line.split("data: ")[1]
+                                )
+                                if event_type == "progress":
+                                    progress = round(100 * data_content["progress"])
+                                    progress_bar.update(progress - progress_bar.n)
+                                elif event_type == "result":
+                                    results = []
+                                    for key in sorted(data_content.keys()):
+                                        if key.startswith("result"):
+                                            decoded_image = self._decode_array(
+                                                data_content[key]
+                                            )
+                                            results.append(decoded_image)
+
+                                    if results:
+                                        prediction = np.stack(results, axis=-1)
+                                        if binarize:
+                                            prediction = geometry.binarize_hard(
+                                                prediction
+                                            )
+                                        progress_bar.close()
+                                        return prediction
+                                elif event_type == "end":
+                                    print("Stream ended.")
                                     progress_bar.close()
-                                    return prediction
-                            elif event_type == "end":
-                                print("Stream ended.")
-                                progress_bar.close()
-                                break
-                            elif event_type == "auth":
-                                if "new_refresh_token" in data_content["auth"]:
-                                    prefab_file_path = os.path.expanduser(
-                                        "~/.prefab.toml"
-                                    )
-                                    with open(
-                                        prefab_file_path, "w", encoding="utf-8"
-                                    ) as toml_file:
-                                        toml.dump(
-                                            {
-                                                "access_token": data_content["auth"][
-                                                    "new_access_token"
-                                                ],
-                                                "refresh_token": data_content["auth"][
-                                                    "new_refresh_token"
-                                                ],
-                                            },
-                                            toml_file,
+                                    break
+                                elif event_type == "auth":
+                                    if "new_refresh_token" in data_content["auth"]:
+                                        prefab_file_path = os.path.expanduser(
+                                            "~/.prefab.toml"
                                         )
-                            elif event_type == "error":
-                                print(f"Error: {data_content['error']}")
-                                progress_bar.close()
-                        except json.JSONDecodeError:
-                            print(
-                                "Failed to decode JSON:",
-                                decoded_line.split("data: ")[1],
-                            )
+                                        with open(
+                                            prefab_file_path, "w", encoding="utf-8"
+                                        ) as toml_file:
+                                            toml.dump(
+                                                {
+                                                    "access_token": data_content[
+                                                        "auth"
+                                                    ]["new_access_token"],
+                                                    "refresh_token": data_content[
+                                                        "auth"
+                                                    ]["new_refresh_token"],
+                                                },
+                                                toml_file,
+                                            )
+                                elif event_type == "error":
+                                    raise ValueError(f"{data_content['error']}")
+                            except json.JSONDecodeError:
+                                raise ValueError(
+                                    "Failed to decode JSON:",
+                                    decoded_line.split("data: ")[1],
+                                ) from None
+        except requests.RequestException as e:
+            raise RuntimeError(f"Request failed: {e}") from e
 
     def predict(
         self,
         model: Model,
         binarize: bool = False,
+        gpu: bool = False,
     ) -> "Device":
         """
         Predict the nanofabrication outcome of the device using a specified model.
@@ -345,6 +374,10 @@ class Device(BaseModel):
             If True, the predicted device geometry will be binarized using a threshold
             method. This is useful for converting probabilistic predictions into binary
             geometries. Defaults to False.
+        gpu : bool, optional
+            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
         -------
@@ -361,6 +394,7 @@ class Device(BaseModel):
             model=model,
             model_type="p",
             binarize=binarize,
+            gpu=gpu,
         )
         return self.model_copy(update={"device_array": prediction_array})
 
@@ -368,6 +402,7 @@ class Device(BaseModel):
         self,
         model: Model,
         binarize: bool = True,
+        gpu: bool = False,
     ) -> "Device":
         """
         Correct the nanofabrication outcome of the device using a specified model.
@@ -391,6 +426,10 @@ class Device(BaseModel):
             If True, the corrected device geometry will be binarized using a threshold
             method. This is useful for converting probabilistic corrections into binary
             geometries. Defaults to True.
+        gpu : bool, optional
+            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
         -------
@@ -407,12 +446,14 @@ class Device(BaseModel):
             model=model,
             model_type="c",
             binarize=binarize,
+            gpu=gpu,
         )
         return self.model_copy(update={"device_array": correction_array})
 
     def semulate(
         self,
         model: Model,
+        gpu: bool = False,
     ) -> "Device":
         """
         Simulate the appearance of the device as if viewed under a scanning electron
@@ -432,6 +473,10 @@ class Device(BaseModel):
             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 SEMulation
             is tailored to specific fabrication parameters.
+        gpu : bool, optional
+            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
         -------
@@ -443,6 +488,7 @@ class Device(BaseModel):
             model=model,
             model_type="s",
             binarize=False,
+            gpu=gpu,
         )
         return self.model_copy(update={"device_array": semulated_array})
 
@@ -464,10 +510,12 @@ class Device(BaseModel):
         buffer_thickness = self.buffer_spec.thickness
         buffer_mode = self.buffer_spec.mode
 
-        crop_top = buffer_thickness if buffer_mode["top"] == "edge" else 0
-        crop_bottom = buffer_thickness if buffer_mode["bottom"] == "edge" else 0
-        crop_left = buffer_thickness if buffer_mode["left"] == "edge" else 0
-        crop_right = buffer_thickness if buffer_mode["right"] == "edge" else 0
+        crop_top = buffer_thickness["top"] if buffer_mode["top"] == "edge" else 0
+        crop_bottom = (
+            buffer_thickness["bottom"] if buffer_mode["bottom"] == "edge" else 0
+        )
+        crop_left = buffer_thickness["left"] if buffer_mode["left"] == "edge" else 0
+        crop_right = buffer_thickness["right"] if buffer_mode["right"] == "edge" else 0
 
         ndarray = device_array[
             crop_top : device_array.shape[0] - crop_bottom,
@@ -629,7 +677,7 @@ class Device(BaseModel):
         bounds: Optional[tuple[tuple[int, int], tuple[int, int]]],
         ax: Optional[Axes],
         **kwargs,
-    ) -> Axes:
+    ) -> tuple[plt.cm.ScalarMappable, Axes]:
         if ax is None:
             _, ax = plt.subplots()
         ax.set_ylabel("y (nm)")
@@ -673,6 +721,13 @@ class Device(BaseModel):
         if show_buffer:
             self._add_buffer_visualization(ax)
 
+        # # Adjust colorbar font size if a colorbar is added
+        # if "cmap" in kwargs:
+        #     cbar = plt.colorbar(mappable, ax=ax)
+        #     cbar.ax.tick_params(labelsize=14)
+        #     if "label" in kwargs:
+        #         cbar.set_label(kwargs["label"], fontsize=16)
+
         return mappable, ax
 
     def plot(
@@ -937,9 +992,9 @@ class Device(BaseModel):
         buffer_hatch = "/"
 
         mid_rect = Rectangle(
-            (buffer_thickness, buffer_thickness),
-            plot_array.shape[1] - 2 * buffer_thickness,
-            plot_array.shape[0] - 2 * buffer_thickness,
+            (buffer_thickness["left"], buffer_thickness["top"]),
+            plot_array.shape[1] - buffer_thickness["left"] - buffer_thickness["right"],
+            plot_array.shape[0] - buffer_thickness["top"] - buffer_thickness["bottom"],
             facecolor="none",
             edgecolor="black",
             linewidth=1,
@@ -949,25 +1004,25 @@ class Device(BaseModel):
         top_rect = Rectangle(
             (0, 0),
             plot_array.shape[1],
-            buffer_thickness,
+            buffer_thickness["top"],
             facecolor=buffer_fill,
             hatch=buffer_hatch,
         )
         ax.add_patch(top_rect)
 
         bottom_rect = Rectangle(
-            (0, plot_array.shape[0] - buffer_thickness),
+            (0, plot_array.shape[0] - buffer_thickness["bottom"]),
             plot_array.shape[1],
-            buffer_thickness,
+            buffer_thickness["bottom"],
             facecolor=buffer_fill,
             hatch=buffer_hatch,
         )
         ax.add_patch(bottom_rect)
 
         left_rect = Rectangle(
-            (0, buffer_thickness),
-            buffer_thickness,
-            plot_array.shape[0] - 2 * buffer_thickness,
+            (0, buffer_thickness["top"]),
+            buffer_thickness["left"],
+            plot_array.shape[0] - buffer_thickness["top"] - buffer_thickness["bottom"],
             facecolor=buffer_fill,
             hatch=buffer_hatch,
         )
@@ -975,11 +1030,11 @@ class Device(BaseModel):
 
         right_rect = Rectangle(
             (
-                plot_array.shape[1] - buffer_thickness,
-                buffer_thickness,
+                plot_array.shape[1] - buffer_thickness["right"],
+                buffer_thickness["top"],
             ),
-            buffer_thickness,
-            plot_array.shape[0] - 2 * buffer_thickness,
+            buffer_thickness["right"],
+            plot_array.shape[0] - buffer_thickness["top"] - buffer_thickness["bottom"],
             facecolor=buffer_fill,
             hatch=buffer_hatch,
         )
@@ -1017,7 +1072,9 @@ class Device(BaseModel):
         binarized_device_array = geometry.binarize(
             device_array=self.device_array, eta=eta, beta=beta
         )
-        return self.model_copy(update={"device_array": binarized_device_array})
+        return self.model_copy(
+            update={"device_array": binarized_device_array.astype(np.uint8)}
+        )
 
     def binarize_hard(self, eta: float = 0.5) -> "Device":
         """
@@ -1038,12 +1095,14 @@ class Device(BaseModel):
         binarized_device_array = geometry.binarize_hard(
             device_array=self.device_array, eta=eta
         )
-        return self.model_copy(update={"device_array": binarized_device_array})
+        return self.model_copy(
+            update={"device_array": binarized_device_array.astype(np.uint8)}
+        )
 
     def binarize_monte_carlo(
         self,
         threshold_noise_std: float = 2.0,
-        threshold_blur_std: float = 9.0,
+        threshold_blur_std: float = 8.0,
     ) -> "Device":
         """
         Binarize the device geometry using a Monte Carlo approach with Gaussian
@@ -1106,9 +1165,10 @@ class Device(BaseModel):
 
         Parameters
         ----------
-        buffer_thickness : int, optional
-            The thickness of the buffer to leave around the empty space. Defaults to 0,
-            which means no buffer is added.
+        buffer_thickness : dict, optional
+            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.
 
         Returns
         -------
@@ -1233,3 +1293,5 @@ class Device(BaseModel):
             with higher values indicating greater uncertainty.
         """
         return 1 - 2 * np.abs(0.5 - self.device_array)
+
+        return 1 - 2 * np.abs(0.5 - self.device_array)