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 <- 1000
sigma_n <- 0.1
n_k <- array(rnorm(T*3,0, sigma_n), c(3,T))
sigma_m <- 0.1
m_k <- array(rnorm(T*3,0, sigma_m), c(3,T))
F <- array(c(1,0,0, 0, 1, 0, 0, 0,1),c(3,3))
x_0 <- array(1,c(3,1))
x <- array(0,c(3,T))
x[,1] <- x_0
z <- array(0,c(3,T))
for (k in 2:T)
{
x[,k] <- F %*% x[,k-1] + n_k[,k]
z[,k] <- `^`(x[,k],2) + m_k[,k]
}
# Next, form the EKF
H <- 2* array(c(1,0,0, 0, 1, 0, 0, 0,1),c(3,3))
Q <- sigma_m^2*array(c(1,0,0, 0, 1, 0, 0, 0,1),c(3,3))
R <- sigma_n^2*array(c(1,0,0, 0, 1, 0, 0, 0,1),c(3,3))
library("MASS") # For pseudo inverse ginv()
xkm1km1 <- matrix(rep(0,3*T+3),3,T+1)
xkm1km1[,1] <- x_0
xkkm1 <- matrix(rep(0,3*T),3,T)
K <- array(rep(0,3*3*T),c(3,3,T))
Pkm1km1 <- array(0,c(3,3,T+1))
Pkm1km1[,,1] <- array(c(1000,0,0 ,0,1000,0, 0,0,1000), c(3,3))
zhat <- matrix(rep(0,3*T),c(3,T))
err <- matrix(rep(0,3*T),c(3,T))
for (k in 2:T)
{
xkkm1[,k] <- F %*% xkm1km1[,k-1]
Pkkm1 <- F %*% Pkm1km1[,,k-1] %*% t(F) + Q
H <- 2*diag(xkkm1[,k])
K[,,k] <- Pkkm1 %*% t(H) %*% ginv( H %*% Pkkm1 %*% t(H) + R)
err[,k] <- z[,k] - H %*% xkkm1[,k]
xkm1km1[,k] <- xkkm1[,k] + K[,,k] %*% err[,k]
Pkm1km1[,,k] <- (matrix(c(1,0,0,0,1,0,0,0,1),3,3) - K[,,k] %*% H) %*% Pkkm1
zhat[,k] <- as.numeric(H %*% xkkm1[,k])
}
p1 <- x
p2 <- xkm1km1
lims <- 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")