How to use the plonk.analysis.particles function in plonk

def momentum(snap: SnapLike, ignore_accreted: bool = True) -> ndarray:
    """Calculate the total momentum.

    Parameters
    ----------
    snap
        The Snap object.
    ignore_accreted : optional
        Ignore accreted particles. Default is True.

    Returns
    -------
    ndarray
        The total linear momentum like (px, py, pz).
    """
    return particles.momentum(snap=snap, ignore_accreted=ignore_accreted).sum(axis=0)

def _add_array(snap, name):
    return set(particles.array_requires[name]).issubset(snap.available_arrays())

def kinetic_energy(snap: SnapLike, ignore_accreted: bool = True) -> float:
    """Calculate the total kinetic energy.

    Parameters
    ----------
    snap
        The Snap object.
    ignore_accreted : optional
        Ignore accreted particles. Default is True.

    Returns
    -------
    float
        The total kinetic energy.
    """
    return particles.kinetic_energy(snap=snap, ignore_accreted=ignore_accreted).sum(
        axis=0
    )

def stokes_number(snap) -> ndarray:
            """Stokes number."""
            gravitational_parameter = snap.properties.get('gravitational_parameter')
            if gravitational_parameter is None:
                raise ValueError(
                    'To get eccentricity, first set the gravitational parameter\n'
                    'via snap.set_gravitational_parameter.'
                )
            origin = (
                snap.translation if snap.translation is not None else (0.0, 0.0, 0.0)
            )
            return particles.stokes_number(
                snap=snap,
                gravitational_parameter=gravitational_parameter,
                origin=origin,
            )