I want to derive the diffusion flux $J_D$ of charged particles with charge $e$ moving in the z direction where we apply an electric field $\vec{E}=E\hat{z}$ in order to proof the Einstein relation that is often given without derivation or using Kubo's derivation. I am looking for an easy approach assuming thermal equilibrium. The mean number density at $z$ is $n(z)$. I tried the following:
$$\frac{n(z+dz)}{n(z)}=\frac{e^{-\beta W(z+dz)}}{e^{-\beta W(z)}},$$
where I used $W(z)=-eEz$ and $W(z+dz) = -eE(z+dz)$ so these are potential energies. From this I derive that
$$\frac{n(z+dz)}{n(z)} \approx 1+\beta eEdz.$$
Using the expansion $e^x \approx 1+x$ in the last part.
This then gives with $$J_D = -D\frac{dn}{dz}$$
$$J_D = -Dn(z)\beta eE,$$
which is exactly what I need to derive the Einstein relation using $J_\mu = n\mu E$ and $J_D + J_\mu = 0$ so that we find:
$$\frac{\mu}{D} = e\beta,$$
which is indeed the Einstein relation.
My question is: how can we get away with deriving this using the potential energy and not the total energy?