import sys
sys.path.append("/mnt/aureum/kcchan_rec/ding_post_mess/jupyter_mess/aureum")

import wx_iter_funcs as wif
import numpy as np 
from matplotlib.colors import LogNorm
import matplotlib.pyplot as plt
import matplotlib
from importlib import reload

import os
import matplotlib

os.environ['PATH'] = '/software/texlive-20240311/bin/x86_64-linux:' + os.environ['PATH']
matplotlib.rc_file('/software/mess_config/beauty/matplotlibrc_dingBeauty_v3.9.3')

---

如何从粒子数据开始，计算某个snap的功率谱？

功率谱就是Fourier空间中的k-bin内对k-mode的振幅的平均，所以需要对场来计算，而不是粒子；如果你没有场的数据，那么需要先制作；如果有的话那就更好了可以直接跳到步骤2即可

0. 读取粒子数据(这里我以我自己跑的一个模拟为例，你如果是tng的那就把对应的数据置换即可)

1. 粒子转场(这里最简单用CIC，非补偿，如果需要补偿后续再说)，我们使用pylians即可，底层是c写的函数包，所以比nbodykit要快

2. 计算功率谱(为了方便，直接用nbodykit来计算即可)

PS: 所有的依赖，在conda env:nbody-env中我都已装好，可以直接用就行

---

0. 读取粒子数据(内存占用较大，后续你最好丢slurm去跑就行)

PS: 这里我只是用我自己跑的一个N体举例，对DM做简单CIC，我们Weighting全部设置1就行，因为它们的质量都相同

import h5py

_snap_path = "/Raid6/2/luantch/sim_workspace/1024/BAO/L1000_N1024/snap_BAO_020.hdf5"
with h5py.File(_snap_path, 'r') as snap:
    _header  = snap['Header']
    z        = _header.attrs['Redshift']
    boxsize  = _header.attrs['BoxSize']
    _data_dm = snap['PartType1']
    pos      = _data_dm['Coordinates'][:]

print(f'Redshift = {z:.4f}')
print(f'BoxSize = {boxsize:.2f}')
print(f'Particle Num = {pos.shape[0]}')

Redshift = 1.1144
BoxSize = 1000.00
Particle Num = 1073741824

1. 粒子转场(pylians完成)

import MAS_library as MASL

def pylians_CIC_weight(pos, grid, BoxSize, W, MAS='CIC', use_mean=True, verbose=True):
    _pos = pos.astype(np.float32)
    _cube = np.zeros((grid,grid,grid), dtype=np.float32)
    MASL.MA(_pos.astype(np.float32), _cube, BoxSize, MAS, W, verbose=verbose)
    if use_mean:
        _cube /= np.mean(_cube, dtype=np.float32)
    if verbose: print(_cube.min(), _cube.max(), _cube.mean())
    return _cube

## Construct Weighting arr, here we set all to unity.
w_arr = np.zeros_like(pos[:,0])
w_arr += 1.

%%time

Ng = 512 # Here we do CIC on (512, 512, 512) GridMesh

cube_mean = wif.pylians_CIC_weight(pos, Ng, boxsize, w_arr, MAS='CIC', use_mean=True, verbose=False)

CPU times: user 49.6 s, sys: 2min 44s, total: 3min 33s
Wall time: 3min 33s

👆这里我们因为取了平均，也就是计算的是所谓的👇：

$$ cube = \delta + 1, $$

其中$\delta$就是所谓的overdensity

print(cube_mean.mean())
print(cube_mean.min())
print(cube_mean.max())

1.0000001
0.017697573
343.68887

👆 你可以对所有你需要的snap，都先做转场的操作，然后保存出去，后面就方便多了可以直接使用了，给你一个我运行batch转snap到场的实操路径，你可以参考一下， /Raid6/2/luantch/sim_workspace/1024/BAO/L1000_N1024/field_out/R4.3_gen_field_Gadget4_LEAN.py，这个脚本负责按照snap号读取粒子数据，然后在对应的文件夹里生成场的数据保存，而负责执行它的slurm提交脚本就是： /Raid6/2/luantch/sim_workspace/1024/BAO/L1000_N1024/field_out/run_genfield.sh,

你可以看一下，看如果有需要的话改成满足自己需求的脚本执行即可

2. 计算功率谱(nbodykit完成)

import nbodykit.lab as nb

def cal_pk(cube, boxsize, kmin, dk, mode='1d'):
    # _field = nb.ArrayMesh(cube/cube.mean()-1., BoxSize=boxsize) 
    _field = nb.ArrayMesh(cube, BoxSize=boxsize)  # Here we use the original cube, but not delta cube.
    _pk=nb.FFTPower(_field, mode=mode,  dk=dk, kmin=kmin).power
    return {'k': _pk['k'], 'power': _pk['power'].real}

%%time

# Configure kmin and dk according to the requirements of your simulation snapshot.

_pk_dict = cal_pk(cube_mean, boxsize=boxsize, kmin=0.00628, dk=0.01)

CPU times: user 17.7 s, sys: 6.59 s, total: 24.3 s
Wall time: 24.2 s

plt.figure(figsize=(6,5), dpi=100)

plt.loglog(_pk_dict['k'], _pk_dict['power'])
plt.xlabel(r'$k \ [{h\,\rm Mpc^{-1}}]$',fontsize=18)
plt.ylabel(r'$P(k) \ [(\mathrm{Mpc}/h)^3]$', fontsize=18)

Text(0, 0.5, '$P(k) \\ [(\\mathrm{Mpc}/h)^3]$')

👆 至此大功告成，有问题再随时问我😊