Water generating waves and reflection

Question

For the last few days I have been working on wave generation script. As there are a lot of calculations involved I was looking for possible optimizations to improve the performance of the code.

This specifically applies to the update function, as most of the calculations are done in here. As the full code is quite large, I made a pastebin of it which can be found over here.

void Update()
{
    //Get sun reflection dir from sun object
    if (Sun != null)
    {
        SunDirection = Sun.transform.forward;
        OceanMaterial.SetVector("_SunDir", SunDirection);
    }

    if (this.IsRenderingReflection)
        RenderObject();

    OceanMaterial.SetTextureOffset("_Foam", Time.time * new Vector2(33,33));

    var halfHeight = height / 2f;
    var halfWidth = width / 2f;
    var time = Time.time;

    for (var y = 0; y < height; y++)
    {
        for (var x = 0; x < width; x++)
        {
            var idx = width * y + x;
            var yc = y < halfHeight ? y : -height + y;
            var xc = x < halfWidth ? x : -width + x;
            Vector2 vec_k = new Vector2(2.0f * Mathf.PI * xc / ChunkSize.x, 2.0f * Mathf.PI * yc / ChunkSize.z);

            var sqrtMagnitude = (float)System.Math.Sqrt((vec_k.x * vec_k.x) + (vec_k.y * vec_k.y));
            var iwkt = (float)System.Math.Sqrt(9.81f * sqrtMagnitude) * time * WaveSpeed;
            ComplexF coeffA = new ComplexF((float)System.Math.Cos(iwkt), (float)System.Math.Sin(iwkt));
            ComplexF coeffB;
            coeffB.Re = coeffA.Re;
            coeffB.Im = -coeffA.Im;

            var ny = y > 0 ? height - y : 0;
            var nx = x > 0 ? width - x : 0;

            waterHeightMatrix[idx] = h0[idx] * coeffA + h0[width * ny + nx].GetConjugate() * coeffB;
            t_x[idx] = waterHeightMatrix[idx] * new ComplexF(0.0f, vec_k.x) - waterHeightMatrix[idx] * vec_k.y;

            // Choppy wave calculations
            if (x + y > 0)
                waterHeightMatrix[idx] += waterHeightMatrix[idx] * vec_k.x / sqrtMagnitude;
        }
    }

    Fourier.FFT2(waterHeightMatrix, width, height, FourierDirection.Backward);
    Fourier.FFT2(t_x, width, height, FourierDirection.Backward);

    // Get base values for vertices and uv coordinates.
    if (baseHeight == null)
    {
        baseHeight = baseMesh.vertices;
        vertices = new Vector3[baseHeight.Length];
        normals = new Vector3[baseHeight.Length];
        tangents = new Vector4[baseHeight.Length];
    }

    var wh = width * height;
    var scaleA = ChoppyScale / wh;
    var scaleB = WaveScale / wh;
    var scaleBinv = 1.0f / scaleB;

    for (var i = 0; i < wh; i++)
    {
        var iw = i + i / width;
        vertices[iw] = baseHeight[iw];
        vertices[iw].x += waterHeightMatrix[i].Im * scaleA;
        vertices[iw].y = waterHeightMatrix[i].Re * scaleB;

        normals[iw] = Vector3.Normalize(new Vector3(t_x[i].Re, scaleBinv, t_x[i].Im));

        if (((i + 1) % width) == 0)
        {
            var iwi = iw + 1;
            var iwidth = i + 1 - width;
            vertices[iwi] = baseHeight[iwi];
            vertices[iwi].x += waterHeightMatrix[iwidth].Im * scaleA;
            vertices[iwi].y = waterHeightMatrix[iwidth].Re * scaleB;

            normals[iwi] = Vector3.Normalize(new Vector3(t_x[iwidth].Re, scaleBinv, t_x[iwidth].Im));

        }
    }

    var offset = geometryWidth * (geometryHeight - 1);

    for (var i = 0; i < geometryWidth; i++)
    {
        var ioff = i + offset;
        int mod = i % width;
        vertices[ioff] = baseHeight[ioff];
        vertices[ioff].x += waterHeightMatrix[mod].Im * scaleA;
        vertices[ioff].y = waterHeightMatrix[mod].Re * scaleB;

        normals[ioff] = Vector3.Normalize(new Vector3(t_x[mod].Re, scaleBinv, t_x[mod].Im));
    }

    var geoWidthHeight = geometryWidth * geometryHeight - 1;
    for (var i = 0; i < geoWidthHeight; i++)
    {

        //Need to preserve w in refraction/reflection mode
        if (!reflectionRefractionEnabled)
        {
            if (((i + 1) % geometryWidth) == 0)
            {
                tangents[i] = Vector3.Normalize((vertices[i - width + 1] + new Vector3(ChunkSize.x, 0.0f, 0.0f) - vertices[i]));
            }
            else
            {
                tangents[i] = Vector3.Normalize((vertices[i + 1] - vertices[i]));
            }

            tangents[i].w = 1.0f;
        }
        else
        {
            Vector3 tmp; // = Vector3.zero;

            if (((i + 1) % geometryWidth) == 0)
            {
                tmp = Vector3.Normalize(vertices[i - width + 1] + new Vector3(ChunkSize.x, 0.0f, 0.0f) - vertices[i]);
            }
            else
            {
                tmp = Vector3.Normalize(vertices[i + 1] - vertices[i]);
            }

            tangents[i] = new Vector4(tmp.x, tmp.y, tmp.z, tangents[i].w);
        }
    }

    //In reflection mode, use tangent w for foam strength
    if (reflectionRefractionEnabled)
    {
        for (var y = 0; y < geometryHeight; y++)
        {
            for (var x = 0; x < geometryWidth; x++)
            {
                var item = x + geometryWidth * y;
                if (x + 1 >= geometryWidth)
                {
                    tangents[item].w = tangents[geometryWidth * y].w;

                    continue;
                }

                if (y + 1 >= geometryHeight)
                {
                    tangents[item].w = tangents[x].w;

                    continue;
                }

                var right = vertices[(x + 1) + geometryWidth * y].x - vertices[item].x;

                var foam = right / (ChunkSize.x / geometryWidth);


                if (foam < 0.0f)
                    tangents[item].w = 1f;
                else if (foam < 0.5f)
                    tangents[item].w += 3.0f * Time.deltaTime;
                else
                    tangents[item].w -= 0.4f * Time.deltaTime;

                if (Player != null)
                {
                    Vector3 player2Vertex = (Player.position - vertices[item] - transform.position);
                    // foam around boat
                    if (player2Vertex.x >= ChunkSize.x)
                        player2Vertex.x -= ChunkSize.x;

                    if (player2Vertex.x <= -ChunkSize.x)
                        player2Vertex.x += ChunkSize.x;

                    if (player2Vertex.z >= ChunkSize.z)
                        player2Vertex.z -= ChunkSize.z;

                    if (player2Vertex.z <= -ChunkSize.z)
                        player2Vertex.z += ChunkSize.z;
                    player2Vertex.y = 0;

                    if (player2Vertex.sqrMagnitude < WakeDistance * WakeDistance)
                        tangents[item].w += 3.0f * Time.deltaTime;
                }

                tangents[item].w = Mathf.Clamp(tangents[item].w, 0.0f, 2.0f);
            }
        }
    }

    tangents[geoWidthHeight] = Vector4.Normalize(vertices[geoWidthHeight] + new Vector3(ChunkSize.x, 0.0f, 0.0f) - vertices[1]);

    for (var L0D = 0; L0D < maxLOD; L0D++)
    {
        var den = (int)System.Math.Pow(2f, L0D);
        var itemcount = (int)((height / den + 1) * (width / den + 1));

        Vector4[] tangentsLOD = new Vector4[itemcount];
        Vector3[] verticesLOD = new Vector3[itemcount];
        Vector3[] normalsLOD = new Vector3[itemcount];

        var idx = 0;

        for (var y = 0; y < geometryHeight; y += den)
        {
            for (var x = 0; x < geometryWidth; x += den)
            {
                var idx2 = geometryWidth * y + x;
                verticesLOD[idx] = vertices[idx2];
                tangentsLOD[idx] = tangents[idx2];
                normalsLOD[idx++] = normals[idx2];
            }
        }
        for (var k = 0; k < tilesLOD[L0D].Count; k++)
        {
            Mesh meshLOD = tilesLOD[L0D][k];
            meshLOD.vertices = verticesLOD;
            meshLOD.normals = normalsLOD;
            meshLOD.tangents = tangentsLOD;
        }
    }
}

Answer 1

I'm not familiar with unity or simulation of water waves so just some generic remarks:

1. There are several variables used in your method which are not declared in your method but look like they are local variables such as width, height, maxLOD. If these are in fact class members then the two common C# conventions are to either access them with this. or to prefix the names with _ - this way it will be easy to see by the reader that these variables are in fact class members. This is especially important as your method is very long and introduces a fair amount of local variables. Increases readability significantly for someone unfamiliar with your code (or for yourself when you come back to this in 12 month time trying to fix some obscure bug).

2. As mentioned the method is quite long and it's not easy to at least get a basic understanding of what all the different blocks do. I think it would help to possibly break out the different blocks into individual methods with descriptive names. Alternatively leave a descriptive comment before each block denoting what the next step is doing.

3. There is some code duplication in this block:

   var geoWidthHeight = geometryWidth * geometryHeight - 1;
   for (var i = 0; i < geoWidthHeight; i++)
   {
       //Need to preserve w in refraction/reflection mode
       if (!reflectionRefractionEnabled)
       {
           if (((i + 1) % geometryWidth) == 0)
           {
               tangents[i] = Vector3.Normalize((vertices[i - width + 1] + new Vector3(ChunkSize.x, 0.0f, 0.0f) - vertices[i]));
           }
           else
           {
               tangents[i] = Vector3.Normalize((vertices[i + 1] - vertices[i]));
           }
           tangents[i].w = 1.0f;
       }
       else
       {
           Vector3 tmp; // = Vector3.zero;
           if (((i + 1) % geometryWidth) == 0)
           {
               tmp = Vector3.Normalize(vertices[i - width + 1] + new Vector3(ChunkSize.x, 0.0f, 0.0f) - vertices[i]);
           }
           else
           {
               tmp = Vector3.Normalize(vertices[i + 1] - vertices[i]);
           }
           tangents[i] = new Vector4(tmp.x, tmp.y, tmp.z, tangents[i].w);
       }
   }
   

   which can be shortened to:

   var geoWidthHeight = geometryWidth * geometryHeight - 1;
   for (var i = 0; i < geoWidthHeight; i++)
   {
       Vector3 tmp;
       if (((i + 1) % geometryWidth) == 0)
       {
           tmp = Vector3.Normalize((vertices[i - width + 1] + new Vector3(ChunkSize.x, 0.0f, 0.0f) - vertices[i]));
       }
       else
       {
           tmp = Vector3.Normalize((vertices[i + 1] - vertices[i]));
       }
   
       //Need to preserve w in refraction/reflection mode
       if (!reflectionRefractionEnabled)
       {
           tangents[i] = tmp;
           tangents[i].w = 1.0f;
       }
       else
       {
           tangents[i] = new Vector4(tmp.x, tmp.y, tmp.z, tangents[i].w);
       }
   }
   

   Also you create new Vector3(ChunkSize.x, 0.0f, 0.0f) on every loop iteration (which is width x height number of iterations). It's possible that either the compiler or the JIT can optimize it but I'd consider extracting it out of the loop and store it in a variable.

4. In the next loop here:

   //In reflection mode, use tangent w for foam strength
   if (reflectionRefractionEnabled)
   {
       for (var y = 0; y < geometryHeight; y++)
       {
           for (var x = 0; x < geometryWidth; x++)
           {
               var item = x + geometryWidth * y;
   ...
   

   you repeat the term geometryWidth * y several times. However this only changes with the outer loop so I'd consider extracting it into a stride variable like this:

   if (reflectionRefractionEnabled)
   {
       var stride = 0;
       for (var y = 0; y < geometryHeight; y++, stride += geometryWidth)
       {
           for (var x = 0; x < geometryWidth; x++)
           {
               var item = x + stride;
   ...
   

   and use stride in the inner loop wherever you used geometryWidth * y before. Again the JIT might perform this optimization but I wouldn't count on it.

5. This is not really a good way to calculate 2^n:

   var den = (int)System.Math.Pow(2f, L0D);
   

   you are converting the int L0D to float and then the result back to int. The classic way to calculate x * 2^n where n is int is to simply do x << n so in your case

   var den = 1 << L0D
   

   Conversely calculating x / 2^n where n is int is easily done via x >> n. So this

   var itemcount = (int)((height / den + 1) * (width / den + 1));
   

   can be turned into

   var itemCount = ((height >> L0D) + 1) * ((width >> L0D) + 1);