From 162f6700506484f7a4661f9e816718b4011bd55a Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Mon, 26 Feb 2024 23:06:21 +0100
Subject: [PATCH] use affine root coords as geom if possible

Before 5034398d3 the geometry of a rectilinear mesh with uniform vertices was
an affine transformation of the root coordinates. In patch 5034398d3 this
special case was removed due to a pending change in the transform chains for
structured topologies: the replacement of the transformation that places an
element in the root space with an index, thereby giving every element unit
square root coordinates.

The effect of this change is a suboptimal gradient of the geometry to the root
coordinates. Since the geometry is now always the inner product of a linear
basis with vertices, the gradient of the geometry does not simplify to a
constant if the vertices are evenly spaced.

This patch restores the special cased geometry for rectilinear meshes defined
by numbers of elements by creating a reference geometry out of the (unit
square) root coordinates and the element index.
---
 nutils/mesh.py | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/nutils/mesh.py b/nutils/mesh.py
index b6237c992..06bbcf824 100644
--- a/nutils/mesh.py
+++ b/nutils/mesh.py
@@ -39,9 +39,18 @@ def rectilinear(richshape: Sequence[Union[int, Sequence[float]]], periodic: Sequ
     axes = [transformseq.DimAxis(i=0, j=n, mod=n if idim in periodic else 0, isperiodic=idim in periodic) for idim, n in enumerate(shape)]
     topo = topology.StructuredTopology(space, root, axes)
 
-    funcsp = topo.basis('spline', degree=1, periodic=())
-    coords = numeric.meshgrid(*verts).reshape(ndims, -1)
-    geom = (funcsp * coords).sum(-1)
+    if all(map(numeric.isint, richshape)):
+        flat_index = topo.f_index
+        rev_unraveled_index = []
+        for n in shape[:0:-1]:
+            rev_unraveled_index.append(flat_index % n)
+            flat_index //= n
+        rev_unraveled_index.append(flat_index)
+        geom = topo.f_coords + numpy.stack(rev_unraveled_index[::-1])
+    else:
+        funcsp = topo.basis('spline', degree=1, periodic=())
+        coords = numeric.meshgrid(*verts).reshape(ndims, -1)
+        geom = (funcsp * coords).sum(-1)
 
     return topo, geom
 
@@ -52,10 +61,12 @@ def rectilinear(richshape: Sequence[Union[int, Sequence[float]]], periodic: Sequ
 def line(nodes: Union[int, Sequence[float]], periodic: bool = False, bnames: Optional[Tuple[str, str]] = None, *, space: str = 'X', root: Optional[TransformItem] = None) -> Tuple[Topology, function.Array]:
     if root is None:
         root = transform.Index(1, 0)
+    nelems = nodes if isinstance(nodes, int) else len(nodes) - 1
+    domain = topology.StructuredLine(space, root, 0, nelems, periodic=periodic, bnames=bnames)
     if isinstance(nodes, int):
-        nodes = numpy.arange(nodes + 1)
-    domain = topology.StructuredLine(space, root, 0, len(nodes) - 1, periodic=periodic, bnames=bnames)
-    geom = domain.basis('std', degree=1, periodic=[]).dot(nodes)
+        geom = domain.f_coords[0] + domain.f_index
+    else:
+        geom = domain.basis('std', degree=1, periodic=[]).dot(nodes)
     return domain, geom