Numpy如何将linspace坐标转换为矩阵最近数组索引

坐标转最近索引的两种实现方案 #

1. 通用适配任意有序坐标的方案 #

如果你的坐标数组不是均匀采样的，可以用np.searchsorted实现，适配所有单调有序的坐标序列：

def get_nearest_idx(arr, target):
    # arr为单调有序的一维坐标数组
    idx = np.searchsorted(arr, target, side="left")
    # 处理target超过坐标最大值的边界情况
    if idx >= len(arr):
        return len(arr) - 1
    # 比较相邻两个点位的距离，取更近的索引
    if idx > 0 and (target - arr[idx-1]) < (arr[idx] - target):
        return idx - 1
    return idx

使用示例：

# 先把x、y转为一维坐标数组
x_coords = x.flatten()
y_coords = y.flatten()

# 给坐标(-0.5, 0.123)对应的最近点位赋值
ix = get_nearest_idx(x_coords, -0.5)
iy = get_nearest_idx(y_coords, 0.123)
# 注意y对应行维度、x对应列维度，不要搞反索引顺序
m[iy, ix] = val

2. 均匀采样（linspace专属）快速方案 #

因为np.linspace是固定步长的均匀采样，可直接通过数学计算得到索引，性能更高：

def get_uniform_idx(min_val, max_val, num_samples, target):
    step = (max_val - min_val) / (num_samples - 1)
    idx = round((target - min_val) / step)
    # 边界裁剪避免索引越界
    return np.clip(idx, 0, num_samples - 1).astype(int)

使用示例：

ix = get_uniform_idx(-1, 1, 100, -0.5)
iy = get_uniform_idx(-1, 1, 100, 0.123)
m[iy, ix] = val

进阶：实现m[坐标x, 坐标y]的直接调用语法 #

如果想要完全贴合你期望的写法，可以封装自定义类，重写索引相关的魔术方法：

import numpy as np

class CoordMatrix:
    def __init__(self, x_min, x_max, x_num, y_min, y_max, y_num):
        self.x_coords = np.linspace(x_min, x_max, x_num)
        self.y_coords = np.linspace(y_min, y_max, y_num)
        self.x_num = x_num
        self.y_num = y_num
        # 生成你需要的复数矩阵
        self.data = self.x_coords.reshape(1, x_num) + 1j * self.y_coords.reshape(y_num, 1)
    
    def _get_idx(self, coords):
        x_target, y_target = coords
        # 此处用均匀采样的快速算法，也可替换为通用的searchsorted方案
        ix = round((x_target - self.x_coords[0]) / (self.x_coords[-1] - self.x_coords[0]) * (self.x_num - 1))
        iy = round((y_target - self.y_coords[0]) / (self.y_coords[-1] - self.y_coords[0]) * (self.y_num - 1))
        ix = np.clip(ix, 0, self.x_num - 1)
        iy = np.clip(iy, 0, self.y_num - 1)
        return int(iy), int(ix)
    
    def __getitem__(self, coords):
        iy, ix = self._get_idx(coords)
        return self.data[iy, ix]
    
    def __setitem__(self, coords, val):
        iy, ix = self._get_idx(coords)
        self.data[iy, ix] = val

使用示例：

# 初始化参数和你原逻辑完全一致的对象
m = CoordMatrix(x_min=-1, x_max=1, x_num=100, y_min=-1, y_max=1, y_num=100)
# 直接用坐标赋值
m[-0.5, 0.123] = 10 + 2j
# 也可直接用坐标取值
print(m[-0.5, 0.123])

内容的提问来源于stack exchange，提问作者Sergey Kostrukov