Kalman Filter - Non Linear Measurement Model I'm new to Kalman filters and estimation in general. I'm running a simple test of an EKF to check my understanding, but I'm getting some odd results with a particular case.

Given a state vector:
$$ \textbf{x} = 
\left(\begin{matrix}
p_x \\
p_y \\
p_z \\
\end{matrix}\right)
$$

If I have a sensor that measures these values directly, i.e.:
$$ \textbf{z} = 
\left(\begin{matrix}
p_x \\
p_y \\
p_z \\
\end{matrix}\right)
$$

Then I get a converging solution. However, if I decide that my sensor actually measures the square of each value, i.e.
$$ \textbf{z} = 
\left(\begin{matrix}
p_x^2 \\
p_y^2 \\
p_z^2 \\
\end{matrix}\right)
$$

Then following the process on the EKF wiki page causes my solution to diverge very rapidly. This seems to be irrespective of the dynamic model of the system, but entirely dependent on whether what the sensor is measuring is linear. Is there a standard approach to follow in this situation?

EDIT: Turns out it was a silly mistake on my part. I just implemented the filter incorrectly. The EKF was valid in this case.
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")