Kalman Filter - Non Linear Measurement Model
So let's assume your model is:
$$
 \textbf{x}_{k} = \textbf{x}_{k-1} +  \textbf{n}_k  
$$
where 
$$
\textbf{n}_k = 
\left(\begin{matrix}
n_x \\
n_y \\
n_z \\
\end{matrix}\right)
$$
and each of the $n_x,n_y,n_z$ are independent, zero-mean Gaussian noise terms with variances $\sigma^2_{nx},\sigma^2_{ny},\sigma^2_{nz}$.

The output equation is then:
$$
\textbf{z}_k = \textbf{f}(\textbf{x}_k) +  \textbf{m}_k = \left(\begin{matrix}
p_x^2 \\
p_y^2 \\
p_z^2 \\
\end{matrix}\right)  
+  \textbf{m}_k  
$$
where 
$$
\textbf{m}_k = 
\left(\begin{matrix}
m_x \\
m_y \\
m_z \\
\end{matrix}\right)
$$
and each of the $m_x,m_y,m_z$ are independent, zero-mean Gaussian noise terms with variances $\sigma^2_{mx},\sigma^2_{my},\sigma^2_{mz}$.

So following the EKF Wikipedia page, that means:
$$
\textbf{F}_{k-1} = \textbf{I}\\
\textbf{H}_{k} = {\tt diag}(2\textbf{x}_k)
$$
where ${\tt diag}$ creates a square matrix from the vector argument.

Implementing this in R (code below) shows the follow example of the true state and the state estimate.  As you can see, it has some problems because the output equation masks the sign of the state (that's why the top plot shows a mirror image around the zero axis).





R Code Below

#30103

# First, construct the signal model and generate the data.
T &lt;- 1000

sigma_n &lt;- 0.1
n_k &lt;- array(rnorm(T*3,0, sigma_n), c(3,T))

sigma_m &lt;- 0.1
m_k &lt;- array(rnorm(T*3,0, sigma_m), c(3,T))

F &lt;- array(c(1,0,0, 0, 1, 0, 0, 0,1),c(3,3))

x_0 &lt;- array(1,c(3,1))
x &lt;- array(0,c(3,T))
x[,1] &lt;- x_0

z &lt;- array(0,c(3,T))


for (k in 2:T)
{
  x[,k] &lt;- F %*% x[,k-1] + n_k[,k]
  z[,k] &lt;- `^`(x[,k],2) + m_k[,k]
}


# Next, form the EKF
H &lt;- 2* array(c(1,0,0, 0, 1, 0, 0, 0,1),c(3,3))

Q &lt;- sigma_m^2*array(c(1,0,0, 0, 1, 0, 0, 0,1),c(3,3))
R &lt;- sigma_n^2*array(c(1,0,0, 0, 1, 0, 0, 0,1),c(3,3))

library("MASS") # For pseudo inverse ginv()

xkm1km1 &lt;- matrix(rep(0,3*T+3),3,T+1)
xkm1km1[,1] &lt;- x_0

xkkm1 &lt;- matrix(rep(0,3*T),3,T)
K &lt;- array(rep(0,3*3*T),c(3,3,T))
Pkm1km1 &lt;- array(0,c(3,3,T+1))
Pkm1km1[,,1] &lt;- array(c(1000,0,0 ,0,1000,0, 0,0,1000), c(3,3))

zhat &lt;- matrix(rep(0,3*T),c(3,T))
err &lt;-  matrix(rep(0,3*T),c(3,T))

for (k in 2:T)
{  
  xkkm1[,k] &lt;- F %*% xkm1km1[,k-1]
  Pkkm1 &lt;- F %*% Pkm1km1[,,k-1] %*% t(F) + Q

  H &lt;- 2*diag(xkkm1[,k])

  K[,,k] &lt;- Pkkm1 %*% t(H) %*% ginv( H %*% Pkkm1 %*% t(H) + R)
  err[,k] &lt;- z[,k] - H %*% xkkm1[,k]
  xkm1km1[,k] &lt;- xkkm1[,k] + K[,,k] %*% err[,k]
  Pkm1km1[,,k] &lt;- (matrix(c(1,0,0,0,1,0,0,0,1),3,3) - K[,,k] %*% H) %*% Pkkm1  
  zhat[,k] &lt;- as.numeric(H %*% xkkm1[,k])
}


p1 &lt;- x
p2 &lt;- xkm1km1
lims &lt;- c(-5,5)
par(mfrow = c(3,1), pty="m")
plot(p1[1,], col="grey",ylim=lims)
lines(p2[1,], col="red")
title("True (grey) and Estimated (red) State p_x")
plot(p1[2,], col="grey",ylim=lims)
lines(p2[2,], col="red")
title("True (grey) and Estimated (red) State p_y")
plot(p1[3,], col="grey",ylim=lims)
lines(p2[3,], col="red")
title("True (grey) and Estimated (red) State p_z")