vcristiani / galaxy-chop

MIT License
5 stars 2 forks source link

Recalculated potential dont work to_dataframe(attributed=["potential"]) #145

Closed leliel12 closed 2 years ago

leliel12 commented 2 years ago
leliel12 commented 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)