Closed leliel12 closed 2 years ago
Test in tests/test_data.py
def test_ParticleSet_to_dataframe_no_potential(data_particleset):
m, x, y, z, vx, vy, vz, soft, pot = data_particleset(
seed=42, has_potential=False
)
pset = data.ParticleSet(
data.ParticleSetType.STARS,
m=m,
x=x,
y=y,
z=z,
vx=vx,
vy=vy,
vz=vz,
softening=soft,
potential=pot,
)
expected = pd.DataFrame(
{
"potential": np.full(len(pset), np.nan),
"total_energy": np.full(len(pset), np.nan),
}
)
df = pset.to_dataframe(attributes=["potential", "total_energy"])
assert df.equals(expected)