def time_derivative(
        coords, velocity_field, timestep, initial_velocity=None):
    """Computes the time derivative (du/dt).

    Args:
        coords (torch.Tensor) : coordinates of shape [N, D]
        velocity_field (neuralff.Field) : velocity field function (aka u)
        timestep (float) : delta t of the simulation
        initial_velocity (torch.Tensor) : if provided, will use this instead of sampling the initial velocity.
                                          This is useful for preconditioning the velocity field.

    Returns:
        (torch.Tensor, torch.Tensor) : 
            - time derivative of shape [N, D]
            - velocity of shape [N, D]
    """
    u = velocity_field.sample(coords)
    if initial_velocity is None:
        dudt = (u - u.detach()) / timestep
    else:
        dudt = (u - initial_velocity) / timestep
    return dudt, u