The argument of periapsis is the angle between the node vector $\vec{n}$ (which points towards the ascending node) and the eccentricity vector $\vec{e}$ which points towards periapsis. By the dot product formula $\vec{a} \cdot \vec{b} = \|\vec{a}\|\|\vec{b}\|\cos(\theta)$ where $\theta$ is the angle between the two vectors, $\frac{\vec{n} \cdot \vec{e}}{ne} = \cos(\omega)$. The angle is $\omega$ by definition as it is the angle between the ascending node and the periapsis. This is the derivation of the original formula, with the coterminal case required due to the quadrant ambiguity of $\mathrm{cos}^{-1}$.
Remember, $\vec{h}$ points in the direction perpendicular to the plane of the orbit. Thus, $\vec{h} \times \vec{n}$, by the right hand rule, is a vector coplanar to the orbital plane, perpendicular to the ascending node.
Note that
$$\|\vec{h} \times \vec{n}\| = \|\vec{h}\|\|\vec{n}\|\sin(\theta)$$
where $\theta$ is the angle between $\vec{h}$ and $\vec{n}$ in the plane containing them. However, remember that $\vec{n} = [0,0,1] \times \vec{h}$ and thus $\vec{h}$ and $\vec{n}$ are orthogonal. Thus, $\theta = \frac{\pi}{2}$ and $\sin(\frac{\pi}{2}) = 1$ and thus the length of this vector is simply $\|\vec{h}\|\|\vec{n}\|$.
Thus, $\vec{r} \cdot (\vec{h} \times \vec{n}) = \|\vec{r}\|\|\vec{h} \times \vec{n}\|\cos(\theta_1) = \|\vec{r}\|\|\vec{h}\|\|\vec{n}\|\cos(\theta_1)$. Which makes $p_y = \|\vec{r}\|\|\vec{n}\|\cos(\theta_1)$. Note that $\theta_1$, the angle between $\vec{r}$ and $\vec{h} \times \vec{n}$, will be $\omega + \nu - \frac{\pi}{2}$. Thus,
$$\cos(\theta_1) = \cos(\omega + \nu - \frac{\pi}{2}) = \sin(\omega + \nu)$$
This is because the angle between $\vec{r}$ and $\vec{h} \times \vec{n}$, a vector perpendicular to $\vec{n}$ (in the same plane as $\vec{r}$), is the same as the angle between $\vec{r}$ and $\vec{n}$ (which is $\omega + \nu$) minus the $\frac{\pi}{2}$ angle between $\vec{n}$ and $\vec{h} \times \vec{n}$.
Additionally, $p_x = \vec{r} \cdot \vec{n} = \|\vec{r}\|\|\vec{n}\|\cos(\theta_2)$. And $\theta_2 = \omega + \nu$, as it is the angle between the ascending node and the current radius, which is the argument of periapsis plus the true anomaly.
Dividing these two:
$$\frac{p_y}{p_x} = \frac{\|\vec{r}\|\|\vec{n}\|\sin(\omega + \nu)}{\|\vec{r}\|\|\vec{n}\|\cos(\omega + \nu)} = \frac{\sin(\omega + \nu)}{\cos(\omega + \nu)} = \tan(\omega + \nu)$$
Thus we get $\mathrm{arctan2}(p_y, p_x) = \tan^{-1}(\frac{p_y}{p_x}) = \omega + \nu$ and thus, finally,
$$\omega = \mathrm{arctan2}(p_y, p_x) - \nu \mod (2\pi)$$
I believe this derivation is correct, but if you spot any mistakes, please don't fret to let me know or correct them yourself!
@
is matrix multiplication, so your mapping onto dot product for vectors should be correct. In case others aren't familiar with the two-argument arctangent arctan2/atan2 I'll leave that here as well. $\endgroup$