use affine root coords as geom if possible

Before 5034398 the geometry of a rectilinear mesh was an affine
transformation of the root coordinates. In patch 5034398 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 by creating a reference geometry
out of the (unit square) root coordinates and the element index and applying an
affine transformation.

nutils/mesh.py

@@ -39,9 +39,32 @@ 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 the vertices are uniform, create a geometry by applying an affine
+    # transformation to the root coordinates. Otherwise create a geometry by
+    # dotting the vertices with a linear basis.
+    flat_index = topo.f_index
+    rev_unraveled_index = []
+    rev_offset = []
+    rev_scale = []
+    for n, v, root_coord in zip(reversed(shape), reversed(richshape), reversed(topo.f_coords)):
+        rev_unraveled_index.append(flat_index % n)
+        flat_index //= n
+        if numeric.isint(v):
+            rev_offset.append(0)
+            rev_scale.append(1)
+        elif numpy.equal(v, numpy.linspace(v[0], v[-1], len(v))).all():
+            rev_offset.append(v[0])
+            rev_scale.append((v[-1] - v[0]) / n)
+        else:
+            # Vertices are not uniform.
+            funcsp = topo.basis('spline', degree=1, periodic=())
+            coords = numeric.meshgrid(*verts).reshape(ndims, -1)
+            geom = (funcsp * coords).sum(-1)
+            return topo, geom
+    unraveled_index = numpy.stack(list(reversed(rev_unraveled_index)))
+    offset = numpy.stack(list(reversed(rev_offset)))
+    scale = numpy.stack(list(reversed(rev_scale)))
+    geom = offset + scale * (topo.f_coords + unraveled_index)
 
     return topo, geom
 
@@ -52,10 +75,15 @@ 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)
+    root_coords = domain.f_coords[0] + domain.f_index
     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 = root_coords
+    elif numpy.equal(nodes, numpy.linspace(nodes[0], nodes[-1], nelems + 1)).all():
+        geom = nodes[0] + (nodes[-1] - nodes[0]) / nelems * root_coords
+    else:
+        geom = domain.basis('std', degree=1, periodic=[]).dot(nodes)
     return domain, geom