Use correct number of units for momentum

Fix for #504

Fix number of units for momentum attributes(s). Use number of components instead of simDim. Tested with 2D LWFA sim.

**Does not affect `master`**