How to transform lattice dots in a non-linear transformation?

[RInterested]

I am trying to reproduce just for the sheer beauty of it, the transformation of the lattice grids and points in the video about Pythagorean triples, specifically this point of the YT video. I have figured out with some online research how to produce the grid and the dots on the initial lattice (see below), and I know that the answer is somewhere in here. However, it is difficult to sort out exactly at which point the corresponding function is defined, and even more critically, there are so many incompatibilities in the syntax between different versions of Manim as to make the process cumbersome. So I wonder if someone can give me some idea as to how to go about it.

This is the code so far:

# From https://github.com/3b1b/videos/blob/master/_2017/triples.py
from manim import *
from manimlib import *

class PythagorasTriples(Scene):
    def construct(self):

        # Complex map
        n = 10
        c_grid = ComplexPlane(
            x_range=(- n, n),  # Adjust the x-axis range
            y_range=(- n, n),  # Adjust the y-axis range
        )

        # Add a lattice of points to the plane
        for x in range(- n, n):
            for y in range(-n, n):
                point = c_grid.number_to_point(x + 1j * y)
                dot = Dot(point, color=BLUE)
                self.add(dot)


        moving_c_grid = c_grid.copy()
        moving_c_grid.set_stroke(WHITE, 4)
        moving_c_grid.prepare_for_nonlinear_transform()
        c_grid.set_stroke(BLUE, 4)

        self.play(
            Write(c_grid, run_time=0),
            FadeIn(moving_c_grid),
        )
        self.wait()
        self.play(
            moving_c_grid.animate.apply_complex_function(lambda z: z**2),
            run_time=6,
        )
        
        self.wait(2)