Ok, so first I'll note that if you're interested in thinking about how the phase fronts of optical waves propagate through free space and are modified when passing through a lens then that means you are interested in Fourier Optics. I highly recommend Fundamentals of Photonics by Saleh and Teich to learn about this question and more.
It's sounds like your goal is to simulate the phase pattern of a beam of light as it hits a lens and then propagates towards its focus. To understand how to do this we need to work out some initial details.
First, our plan of attack will be to work on surfaces of constant $z$. That is, we will think about the complex amplitude of the electric field on different $xy$ planes at different values for $z$. We will then ask the question, if i know the value of the complex amplitude of the electric field at the plane located at $z_0$, what will be the value of the complex amplitude of the electric field at the plane located at $z_1$ located at some distance to the right of $z_0$? We can use a complex function $f_0(x,y)$ to describe the complex amplitude on the plane at $z_0$ and $f_1(x,y)$ to describe the field on the plane at $z_1$. The question is if we know $f_0(x,y)$, what is $f_1(x,y)$?
If the field we are describing is a plane wave travelling directly in the $z$ direction then you have already given us the answer. As the wave propagates it picks up phase at a rate of $2\pi$ radians per distance $\lambda$. This is ratio is captured by the wavevector $k = \frac{2\pi}{\lambda}$. We can then calculate $f_1(x,y)$ as a function of $f_0(x,y)$ by
$$
f_1(x,y) = f_0(x,y) e^{i k(z_1 - z_0)} = f_0(x,y) e^{i kd}
$$
Where $d = z_1-z_0$ is the free space propagation distance. Here's an important question: If we have a plane wave travelling along the $z$ direction what would $f_1(x,y)$ look like? The answer is that it will be a constant function for all of $x$ and $y$. That is why it is called a plane wave. Note we can answer the above questions by recalling that the complex amplitude of the electric field for a plane wave is known, $U(x,y,z) = \tilde{U}_0 e^{i \vec{k}\cdot \vec{r}}$. We can use this to calculate $f_{0,1}(x,y) = U\left(x,y,z_{0,1}\right)$ For example, $f_0(x,y) = \tilde{U}_0 e^{i k z_0}$, a constant function.
However, the next questions is where we start to get at the question in the OP. How must we change this picture if the wave is travelling at some angle $\theta$ with respect to the $z$ axis. Note, we do not change the direction of the $z$ axis, just the direction of the wave. Suppose the wave is pointing in the $+x$ and $+z$ directions but is still in the plane $xz$ plane. That is, we something like
$$
\vec{k} = k_x \hat{x} + k_z \hat{z}
$$
Where $k_x \ll k_z$ so that $\tan(\theta) = \frac{k_x}{k_z}$. We now consider a beam travelling with this new wavevector. If we know it's value on the plane at $z_0:$ $f_0(x,y)$, what will be its value on $f_1(x,y)$? This can again be calculated from $U(x,y,z)$ to be
\begin{align}
\frac{f_1(x,y)}{f_0(x,y)} = e^{i k_z d}
\end{align}
Note the wave has some total wavevector magnitude $\lvert \vec{k} \rvert^2 = k_0^2 = k_x^2 + k_z^2$. We also know that for light of a given frequency, $\omega_0$ that $k_0$ is fixed to be $\frac{\omega}{c}$ by the Helmholtz equation. We see that we can write
$$
k_z = \sqrt{k_0^2 - k_x^2}
$$
So we have
\begin{align}
\frac{f_1(x,y)}{f_0(x,y)} = e^{i d \sqrt{k_0^2 - k_x^2}}
\end{align}
What does this intuitively tell us? The phase factor intuitively tells us how much phase we pick up if we travel directly in the $z$ direction. If $k_x = 0$ and the wave is pointing directly in the $z$ direction we pick up a full $2\pi$ radians for each $\lambda$ we move to the right. However, as $k_x$ increases we pick up less and less phase as we travel to the right. This is because we're not moving in the exact direction of the wavevector. In fact, in the limit that $k_x = k_0$ we pick up NO phase as we travel right. This is because the plane wave is now pointed directly in the vertical direction. We will call this function the transfer function of free space. It tells us how $f_0$ transforms into $f_1$ depending on the value of $k_x$ and denote it by $T_{FS}(k_x)$.
$$
T_{FS}(k_x) = e^{i d \sqrt{k_0^2 - k_x^2}}
$$
Ok, we have one major piece to answer your question. We need one other piece and then we can put them together. We again need to think about what is the field due to this tilted plane wave on the plane at $z_0$. It is not too hard to answer. We get (see for yourself)
$$
f_0 = \tilde{U}_0 e^{i k_z z_0} e^{i k_x x}
$$
The first two terms here are constant complex numbers on the whole plane. However, the final term is a phase which varies sinusoidal in the $x$ direction with spatial wavelength given by $k_x = \frac{2\pi}{\lambda_x}$. (Note that we should think of $k_x$ as the spatial angular frequency of this function). We are see the projection of the periodicity of the plane wave onto the $z$ plane.
Ok now here comes the big kicker on which all of Fourier optics is based.
Imagine We had many plane waves (all with the same frequency) with differnt amplitudes and phases impinging upon and through the plane at $z_0$. Each one of these plane waves would contribute a modulation at a DIFFERENT spatial frequency. Let's describe the complex amplitude of the plane with transverse wavevector $k_x$ by $\tilde{F}(k_x)$. We can then write the total field as
$$
U(x,y,z) = \int_{k_x = k_0}^{+k_0} \tilde{F}(k_x) e^{i (k_x x + k_z z_0)} dk_x
$$
*Or expressing $f_0(x,y) = U(x,y,z_0)$ (setting $z_0 = 0$ for convenience and emphasis)
$$
f_0(x,y) = \int_{k_x = k_0}^{+k_0} \tilde{F}(k_x) e^{i k_x x} dk_x
$$
Ah, this is what we were looking for! We see the spatial amplitude of the field at $z_0$, $f_0(x,y)$ is exactly the spatial Fourier transform of the function $\tilde{F}(k_x)$ describing the coefficients of the plane waves travelling with different wavevectors $k_x$
This is particularly great because know each of these different plane waves complex amplitude will propagate through free space, each one will get multiplied by $T_{FS}(k_x)$ for the corresponding $k_x$.
That is, we just calculated $f_0(x,y)$ but we can already calculate $f_1(x,y)$ just by multiplying each $k_x$ term in the integral by the correct $T_{FS}(k_x)$. That is
\begin{align}
f_1(x,y) = \int_{k_x = -k_0}^{+k_0} T_{FS}(k_x) \tilde{F}(k_x) e^{i k_x x} dk_x
\end{align}
How does this all relate back to the problem at hand? Well, I just said that given a bunch of plane waves with amplitudes $\tilde{F}(k_x)$ that you can calculate $f_0(x,y)$. By Fourier inversion the opposite is true. Given a known field $f_0(x,y)$ you can calculate the plane wave which must be added up to create that field.
In your problem we consider $f_0(x,y)$ to be the field directly after the lens (which you have correctly calculated in your first paragraph but you have a typo in your last equation, there should be no $\lambda$). That is,
$$
f_0(x) = e^{-i \frac{k}{2f}x^2}
$$
Now to find the field $f_1(x)$ on planes at later $z$ you must do the following
- Fourier transform $f_0(x)$ to determine the amplitudes of the planes waves which are involved in "building it up".
- Multiply each of those plane waves by the relevant transfer function over the relevant propagation distance
- Integrate the answer over all $k_x$
This will make the simulation a bit more complicated than what you've done so far but it should be possible. Using the correct program it should be possible to numerically simulate without too much trouble.. You can get analytic expression for the Fourier transforms but you need to make some approximations at various points.
There's probably more to be said here but I feel I've said enough already. I hope this can be helpful
*Why are the limits on the integral $-k_0$ to $+k_0$? Well, in a plane wave no component of the wavevector can exceed the total length of the wavevector, $k_0$, so $k_x$ is bounded by $-k_0$ and $k_0$. It turns out this is the limitation that leads to the fundamental diffraction limit on how small you can focus a beam of light (using far-field propagating plane waves). However, the diffraction limit is a topic for a different post. As a teaser, if you allow the components of the wavevector to become complex then you can have some of the other components exceed $k_0$. However, the expense is that the fields with these complex wavevectors will decay away exponentially. This is a statement of the fact that, using the near field, you can extract higher spatial resolution information.