93 lines
2.6 KiB
Python
93 lines
2.6 KiB
Python
"""
|
||
环境层通用工具
|
||
=============
|
||
提供数组拼接、索引采样、tensor→numpy 转换等辅助功能。
|
||
"""
|
||
|
||
from typing import Dict, Iterable, List, Optional, Union
|
||
|
||
import numpy as np
|
||
from numpy import ndarray
|
||
from torch import Tensor
|
||
from torch_geometric.data.data import BaseData
|
||
|
||
|
||
def save_concatenate(
|
||
arrays: Iterable[np.ndarray], *args, **kwargs
|
||
) -> Optional[np.ndarray]:
|
||
"""
|
||
安全拼接多个数组。自动过滤 None 值，空列表返回 None。
|
||
|
||
Args:
|
||
arrays: 要拼接的数组列表（可能包含 None）
|
||
|
||
Returns:
|
||
拼接后的数组；若全为 None 则返回 None
|
||
|
||
Example:
|
||
>>> result = save_concatenate([arr1, None, arr2], axis=1)
|
||
"""
|
||
arrays = [array for array in arrays if array is not None]
|
||
if len(arrays) == 0:
|
||
return None
|
||
return np.concatenate(arrays, *args, **kwargs)
|
||
|
||
|
||
class IndexSampler:
|
||
"""
|
||
随机索引采样器 — 用于循环缓冲区中随机抽取 PDE 实例。
|
||
|
||
内部维护一个随机排列的索引数组，每次调用 next() 返回一个索引。
|
||
遍历完所有索引后自动重新洗牌。
|
||
|
||
Example:
|
||
>>> sampler = IndexSampler(100, np.random.RandomState(42))
|
||
>>> idx = sampler.next() # 随机抽取一个索引
|
||
"""
|
||
|
||
def __init__(self, size: int, random_state: np.random.RandomState):
|
||
self._size = size
|
||
self._indices = np.arange(size)
|
||
self._random_state = random_state
|
||
self._reset()
|
||
|
||
def next(self) -> int:
|
||
"""返回下一个随机索引，到底后自动洗牌重排。"""
|
||
if self._position == self._size:
|
||
self._reset()
|
||
index = self._indices[self._position]
|
||
self._position += 1
|
||
return index
|
||
|
||
def _reset(self):
|
||
self._position = 0
|
||
self._random_state.shuffle(self._indices)
|
||
|
||
def __len__(self):
|
||
return self._size
|
||
|
||
|
||
def detach(
|
||
tensor: Union[Tensor, Dict[str, Tensor], List[Tensor]],
|
||
) -> Union[ndarray, Dict[str, ndarray], List[ndarray], BaseData]:
|
||
"""
|
||
将 PyTorch tensor 安全转换为 numpy 数组（自动处理 GPU→CPU）。
|
||
|
||
Args:
|
||
tensor: PyTorch tensor、tensor 字典或 tensor 列表
|
||
|
||
Returns:
|
||
对应的 numpy 数组
|
||
|
||
Example:
|
||
>>> action_np = detach(actions_tensor) # → np.ndarray
|
||
"""
|
||
if isinstance(tensor, dict):
|
||
return {key: detach(value) for key, value in tensor.items()}
|
||
elif isinstance(tensor, list):
|
||
return [detach(value) for value in tensor]
|
||
if tensor.is_cuda:
|
||
return tensor.cpu().detach().numpy()
|
||
else:
|
||
return tensor.detach().numpy()
|