将txt文件中素数索引序列导入SageMath并绘制相关函数

我有一个txt格式的数据文件，内容如下：

2     ;      4
                       3     ;      -8
                       5     ;       4
                    next prime      ...

我尝试将这些数据作为素数p索引的序列$(a_p)$导入SageMath以进行绘图，为此编写了如下file.sage代码：

a = open('/home/user/tmp/DataFile', 'r')

ap=[]
line=a.readline()
while(line !=''): 
    b=line.split(';')
    ap.append(integer(b[1]))
    line=f.readline()

但问题是生成的列表ap并非以素数为索引，这使得绘制函数$f(x) := |{p≤x : a_p满足某些约束}|$的基数这一任务变得困难，希望能得到技术帮助。

1. 用字典存储素数与对应值 #

使用字典可以直接通过素数p索引获取a_p，完美匹配你的需求：

# 初始化空字典，键为素数p，值为对应的a_p
ap_dict = {}

# 用with语句自动管理文件资源，避免手动关闭
with open('/home/user/tmp/DataFile', 'r') as f:
    for line in f:
        # 去除行首尾空白字符
        line_clean = line.strip()
        # 跳过空行和标记行（比如包含"next prime"的行）
        if not line_clean or 'next prime' in line_clean:
            continue
        # 分割素数和对应值部分
        prime_part, val_part = line_clean.split(';')
        # 转换为整数并去除多余空白
        p = integer(prime_part.strip())
        a_p = integer(val_part.strip())
        # 存入字典
        ap_dict[p] = a_p

现在你可以直接通过ap_dict[2]、ap_dict[5]获取对应素数的a_p值。

2. 实现约束计数函数 #

要计算$f(x)$（即≤x的素数中满足约束的数量），可以编写如下函数：

def count_valid_primes(x, constraint_func):
    count = 0
    # 遍历字典中的所有素数
    for p in ap_dict:
        if p <= x and constraint_func(ap_dict[p]):
            count += 1
    return count

# 示例：定义约束为"a_p为正数"
positive_constraint = lambda val: val > 0
# 计算x=10时的符合条件的素数数量
print(count_valid_primes(10, positive_constraint))

3. 绘制$f(x)$的图像 #

基于上面的函数，生成绘图数据并调用SageMath的绘图工具：

# 获取所有素数并排序，确定x的范围
sorted_primes = sorted(ap_dict.keys())
max_x = sorted_primes[-1]

# 生成x值序列和对应的f(x)值
x_range = range(2, max_x + 1)
y_values = [count_valid_primes(x, positive_constraint) for x in x_range]

# 绘制折线图
plot(
    list(zip(x_range, y_values)),
    plotjoined=True,
    title="满足约束的素数计数",
    xlabel="x",
    ylabel="f(x)",
    gridlines=True
)

4. 原代码的错误修正 #

原代码存在一个变量名错误：line=f.readline()中的f未定义，应该改为line=a.readline()。另外，使用with语句处理文件更安全，无需手动调用close()方法。

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