numpy.array指数通过argparse - 如何正确地做到这一点?

时间:2013-09-05 09:14:06

标签: python numpy argparse

我想实现一个小模块,它从HDF5文件加载数据,选择一个特定的数据集并对每个子集进行一些计算。我决定使用argparse选项来选择子集。 有没有比使用字符串更好的方法并手动解析它?

下面的代码应该做我想要的但它不是很漂亮,并且还不能处理像':-3'这样的输入。肯定有更好的办法!

谢谢!

parser.add_argument('--filename', required=True, 
               help='the file to open')
parser.add_argument('--snapshots', default="-1", type=str, 
               help="select snapshots (default: %(default)s) ")
[...]
H5 = tables.openFile(filename)

snapshot = args.snapshots
maxval = len(H5.root.trajectory[:])
if snapshot ==":" or snapshot.lower() == "all":
    snapshot = np.arange(maxval)
elif ':' in snapshot:
    if snapshot.index(':') == 0:
        snapshot = np.arange(int(snapshot[1:]))
    elif snapshot.index(':') == len(snapshot) - 1:
        snapshot = np.arange(int(snapshot[:-1]),maxval)
    else:
        from_to = snapshot.split(':')
        snapshot = np.arange(int(from_to[0]), int(from_to[1]))
else:
    snapshot = int(snapshot)
selected = H5.root.trajectory[snapshot]

1 个答案:

答案 0 :(得分:0)

我也倾向于解析字符串,但是使用slice来获取切片。

已更新(请参阅下面的评论)

以下几行::

def parsesection(snapshot):
    if snapshot == 'all':
        return slice(None)
    try:
        section = int(snapshot)
    except ValueError:
        section = [int(s) if s else None for s in snapshot.split(':')]
        if len(section) > 3:
            raise ValueError('snapshots input incorrect')
        section = slice(*section)
    return section

print(parsesection(':5'))
print(parsesection('5:15'))
print(parsesection('5:15:3'))
print(parsesection('3'))
print(parsesection('5::3'))
print(parsesection('::2'))
print(parsesection('all'))

结果:

slice(None, 5, None)
slice(5, 15, None)
slice(5, 15, 3)
3
slice(5, None, 3)
slice(None, None, 2)
slice(None, None, None)

section = slice(*section)可能感觉像是一个技巧,但如果你仔细看,它应该不难理解(我这样做主要是因为如果我省了if elif else系列来检查长度section

这不适用于多维索引;在这种情况下首先分成逗号。如果你想允许省略号,你必须分别捕捉三个点,并返回一个省略号对象而不是切片。