58 lines
No EOL
2.2 KiB
Python
58 lines
No EOL
2.2 KiB
Python
#!/usr/bin/env python
|
|
import sys
|
|
import numpy as np, pandas as pd
|
|
|
|
import importlib.util
|
|
spec = importlib.util.spec_from_file_location('scf', '../etics/src/scf.py')
|
|
scf = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(scf)
|
|
|
|
n_max = 15
|
|
l_max = 2
|
|
length_scale = 10
|
|
|
|
pos_filename = '../../../sobolenko/TNG50-copy/for_yohai/98990136-weight.dat'
|
|
grav_filename = 'out.dat'
|
|
coeff_filename = 'out.coeff'
|
|
|
|
pos_data = np.genfromtxt(pos_filename, skip_header=1, names=['idx', 'm', 'x', 'y', 'z', 'vx', 'vy', 'vz', 'eps', 'type', 'name', 't'], dtype=[int, float, float, float, float, float, float, float, float, int, int, int])
|
|
grav_data = np.genfromtxt(grav_filename, names=['name', 'pot', 'ax', 'ay', 'az'], dtype=[int, float, float, float, float])
|
|
|
|
coeff_count = (n_max+1)*(l_max+1)*(l_max+2)//2
|
|
A = np.fromfile(coeff_filename, dtype=np.complex128)[:coeff_count]
|
|
expansion = scf.Scf(n_max, l_max, length_scale=length_scale)
|
|
expansion.set_coefficients(A)
|
|
|
|
r = np.sqrt(pos_data['x']**2 + pos_data['y']**2 + pos_data['z']**2)
|
|
r[r==0] = 1 # Will be ignored later
|
|
density_all = expansion.calc_radial_density(r)
|
|
|
|
n = len(pos_data['idx'])
|
|
i2 = 0
|
|
out_file = open('results.dat', 'w')
|
|
for i1 in range(10):
|
|
template = '%08d%15.8e%16.8e%16.8e%16.8e%16.8e%16.8e%16.8e%11.4e %d %015d%5.04d%16.8e%16.8e%16.8e%16.8e%15.8e\n'
|
|
if pos_data['name'][i1] == grav_data['name'][i2]:
|
|
pot, ax, ay, az = grav_data['pot'][i2], grav_data['ax'][i2], grav_data['ay'][i2], grav_data['az'][i2]
|
|
den = density_all[i1]
|
|
i2 += 1
|
|
else: # Missing particle
|
|
pot, ax, ay, az = 0, 0, 0, 0
|
|
den = 0
|
|
out = template % (pos_data['idx'][i1],
|
|
pos_data['m'][i1],
|
|
pos_data['x'][i1],
|
|
pos_data['y'][i1],
|
|
pos_data['z'][i1],
|
|
pos_data['vx'][i1],
|
|
pos_data['vy'][i1],
|
|
pos_data['vz'][i1],
|
|
pos_data['eps'][i1],
|
|
pos_data['type'][i1],
|
|
pos_data['name'][i1],
|
|
pos_data['t'][i1],
|
|
pot, ax, ay, az, den
|
|
)
|
|
out_file.write(out)
|
|
|
|
out_file.close() |