initial commit

src/data/dataset.py

@@ -0,0 +1,37 @@
+import torch
+from torch_geometric.data import Dataset, InMemoryDataset
+import os
+
+class LayoutDataset(InMemoryDataset):
+    """用于加载预处理后的版图图数据的 PyTorch Geometric 数据集。"""
+
+    def __init__(self, root, transform=None, pre_transform=None):
+        """
+        Args:
+            root: 数据集应保存的根目录。
+            transform: 一个函数/变换，作用于 `Data` 对象并返回一个转换后的版本。
+            pre_transform: 一个函数/变换，作用于 `Data` 对象并返回一个转换后的版本。
+        """
+        super(LayoutDataset, self).__init__(root, transform, pre_transform)
+        # 加载已处理的数据
+        self.data, self.slices = torch.load(self.processed_paths[0])
+
+    @property
+    def raw_file_names(self):
+        """如果 `download()` 返回一个路径列表，这里会返回它们的文件名。"""
+        return []  # 我们不从网络下载原始文件
+
+    @property
+    def processed_file_names(self):
+        """在 `processed_dir` 目录中必须存在的文件列表，用以跳过处理步骤。"""
+        return ['data.pt']
+
+    def download(self):
+        """从网上下载原始数据到 `raw_dir` 目录。"""
+        pass  # 假设数据是预先处理好的
+
+    def process(self):
+        """处理原始数据并将其保存到 `processed_dir` 目录。"""
+        # 如果希望在加载时动态处理数据，可以在这里实现 `scripts/preprocess_gds.py` 中的逻辑。
+        # 在我们的框架中，我们假设预处理是通过脚本独立完成的。
+        pass

src/data/gds_parser.py

@@ -0,0 +1,68 @@
+from typing import List, Dict, Tuple
+import gdstk
+import numpy as np
+
+class GDSParser:
+    """解析 GDSII/OASIS 文件，提取指定区块内的版图几何图形。"""
+
+    def __init__(self, gds_file: str, layer_mapping: Dict[str, int]):
+        """初始化 GDSParser。
+
+        Args:
+            gds_file: GDSII/OASIS 文件的路径。
+            layer_mapping: 一个字典，将 GDS 的层/数据类型字符串（例如 "1/0"）映射到整数索引。
+        """
+        self.gds_file = gds_file
+        self.layer_mapping = layer_mapping
+        # 使用 gdstk 读取 GDS 文件
+        self.library = gdstk.read_gds(gds_file)
+        # 获取顶层单元
+        self.top_cell = self.library.top_level()[0]
+
+    def get_patches(self, patch_size: float, patch_stride: float) -> List[Tuple[float, float, float, float]]:
+        """生成覆盖整个版图的区块坐标。
+
+        Args:
+            patch_size: 正方形区块的尺寸（单位：微米）。
+            patch_stride: 滑动窗口的步长（单位：微米）。
+
+        Returns:
+            一个包含所有区块边界框 (x_min, y_min, x_max, y_max) 的列表。
+        """
+        # 获取顶层单元的边界框
+        x_min, y_min, x_max, y_max = self.top_cell.bb()
+        patches = []
+        # 使用步长在 x 和 y 方向上生成区块
+        for x in np.arange(x_min, x_max, patch_stride):
+            for y in np.arange(y_min, y_max, patch_stride):
+                patches.append((x, y, x + patch_size, y + patch_size))
+        return patches
+
+    def extract_geometries_from_patch(self, patch_bbox: Tuple[float, float, float, float]) -> List[Dict]:
+        """从给定的区块中提取所有几何对象。
+
+        Args:
+            patch_bbox: 区块的边界框 (x_min, y_min, x_max, y_max)。
+
+        Returns:
+            一个字典列表，每个字典代表一个几何对象及其属性（多边形、层、边界框）。
+        """
+        x_min, y_min, x_max, y_max = patch_bbox
+        # 获取单元内的所有多边形
+        polygons = self.top_cell.get_polygons(by_spec=True)
+        geometries = []
+        # 遍历所有多边形
+        for (layer, datatype), poly_list in polygons.items():
+            layer_str = f"{layer}/{datatype}"
+            # 只处理在 layer_mapping 中定义的层
+            if layer_str in self.layer_mapping:
+                for poly in poly_list:
+                    # 简单的边界框相交检查
+                    p_xmin, p_ymin, p_xmax, p_ymax = poly.bb()
+                    if not (p_xmax < x_min or p_xmin > x_max or p_ymax < y_min or p_ymin > y_max):
+                        geometries.append({
+                            "polygon": poly,
+                            "layer": self.layer_mapping[layer_str],
+                            "bbox": (p_xmin, p_ymin, p_xmax, p_ymax)
+                        })
+        return geometries

src/data/graph_constructor.py

@@ -0,0 +1,83 @@
+from typing import List, Dict
+import torch
+from torch_geometric.data import Data
+from scipy.spatial import cKDTree
+import numpy as np
+
+class GraphConstructor:
+    """从几何图形列表中构建 PyTorch Geometric 的 Data 对象（即图）。"""
+
+    def __init__(self, edge_strategy: str = "knn", knn_k: int = 8, radius_d: float = 1.0):
+        """
+        Args:
+            edge_strategy: 创建边的策略（'knn' 或 'radius'）。
+            knn_k: KNN 策略中的 K（最近邻的数量）。
+            radius_d: 半径图策略中的半径大小。
+        """
+        self.edge_strategy = edge_strategy
+        self.knn_k = knn_k
+        self.radius_d = radius_d
+
+    def construct_graph(self, geometries: List[Dict], label: int = 0) -> Data:
+        """为单个区块构建一个图。
+
+        Args:
+            geometries: 来自 GDSParser 的几何图形字典列表。
+            label: 图的标签（例如，0 表示非热点，1 表示热点）。
+
+        Returns:
+            一个 PyTorch Geometric 的 Data 对象。
+        """
+        # 如果没有几何图形，则返回 None
+        if not geometries:
+            return None
+
+        node_features = []
+        node_positions = []
+        # 提取每个几何图形的特征
+        for geo in geometries:
+            x_min, y_min, x_max, y_max = geo["bbox"]
+            width = x_max - x_min
+            height = y_max - y_min
+            area = width * height
+            centroid_x = x_min + width / 2
+            centroid_y = y_min + height / 2
+
+            # 特征包括：中心点坐标、宽度、高度、面积
+            features = [centroid_x, centroid_y, width, height, area]
+            node_features.append(features)
+            node_positions.append([centroid_x, centroid_y])
+
+        # 将特征和位置转换为 PyTorch 张量
+        x = torch.tensor(node_features, dtype=torch.float)
+        pos = torch.tensor(node_positions, dtype=torch.float)
+
+        # 根据选定的策略创建边
+        edge_index = self._create_edges(pos)
+
+        # 创建图数据对象
+        data = Data(x=x, edge_index=edge_index, pos=pos, y=torch.tensor([label], dtype=torch.float))
+        return data
+
+    def _create_edges(self, node_positions: torch.Tensor) -> torch.Tensor:
+        """根据选定的策略创建边。"""
+        nodes_np = node_positions.numpy()
+        if self.edge_strategy == "knn":
+            # 使用 cKDTree 进行高效的 K 最近邻搜索
+            tree = cKDTree(nodes_np)
+            # 查询每个点的 k+1 个最近邻（包括自身）
+            dist, ind = tree.query(nodes_np, k=self.knn_k + 1)
+            # 创建边列表，排除自环
+            row = np.repeat(np.arange(len(nodes_np)), self.knn_k)
+            col = ind[:, 1:].flatten()
+            edge_index = torch.tensor([row, col], dtype=torch.long)
+
+        elif self.edge_strategy == "radius":
+            # 使用 cKDTree 查找在指定半径内的所有点对
+            tree = cKDTree(nodes_np)
+            pairs = tree.query_pairs(r=self.radius_d)
+            edge_index = torch.tensor(list(pairs), dtype=torch.long).t().contiguous()
+        else:
+            raise ValueError(f"未知的边构建策略: {self.edge_strategy}")
+
+        return edge_index