From 4cfa5c74c0b1f1ef3337879477b9eb76b5ba10d5 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Thu, 18 Feb 2021 15:39:51 +0100
Subject: [PATCH 01/26] fix normal tests in namespace, expression

Two unittests where the normal in (namespace) expressions is being tested
incorrectly uses the dimension of the volume instead of the boundary. This
patch fixes the problem by explicitly defining the dimension to lower for.
---
 tests/test_function.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/test_function.py b/tests/test_function.py
index ebf1437de..a2c1afebb 100644
--- a/tests/test_function.py
+++ b/tests/test_function.py
@@ -672,7 +672,7 @@ def test_d_arg(self):
   def test_n(self):
     ns = function.Namespace()
     topo, ns.x = mesh.rectilinear([1])
-    self.assertEqualLowered(ns.eval_i('n(x_i)'), function.normal(ns.x), ndims=topo.ndims)
+    self.assertEqualLowered(ns.eval_i('n(x_i)'), function.normal(ns.x), ndims=topo.ndims-1)
 
   def test_functions(self):
     def sqr(a):
@@ -738,8 +738,8 @@ def setUp(self):
     self.ns.a32 = numpy.array([[1,2],[3,4],[5,6]])
     self.x = function.Argument('x',())
 
-  def assertEqualLowered(self, s, f):
-    self.assertEqual((s @ self.ns).prepare_eval(ndims=2).simplified, f.prepare_eval(ndims=2).simplified)
+  def assertEqualLowered(self, s, f, *, ndims=2):
+    self.assertEqual((s @ self.ns).prepare_eval(ndims=ndims).simplified, f.prepare_eval(ndims=ndims).simplified)
 
   def test_group(self): self.assertEqualLowered('(a)', self.ns.a)
   def test_arg(self): self.assertEqualLowered('a2_i ?x_i', function.dot(self.ns.a2, function.Argument('x', [2]), axes=[0]))
@@ -748,7 +748,7 @@ def test_multisubstitute(self): self.assertEqualLowered('(a2_i + ?x_i + ?y_i)(x_
   def test_call(self): self.assertEqualLowered('sin(a)', function.sin(self.ns.a))
   def test_call2(self): self.assertEqual(self.ns.eval_ij('arctan2(a2_i, a3_j)').prepare_eval(ndims=2).simplified, function.arctan2(self.ns.a2[:,None], self.ns.a3[None,:]).prepare_eval(ndims=2).simplified)
   def test_eye(self): self.assertEqualLowered('δ_ij a2_i', function.dot(function.eye(2), self.ns.a2, axes=[0]))
-  def test_normal(self): self.assertEqualLowered('n_i', self.ns.x.normal())
+  def test_normal(self): self.assertEqualLowered('n_i', self.ns.x.normal(), ndims=1)
   def test_getitem(self): self.assertEqualLowered('a2_0', self.ns.a2[0])
   def test_trace(self): self.assertEqualLowered('a22_ii', function.trace(self.ns.a22, 0, 1))
   def test_sum(self): self.assertEqualLowered('a2_i a2_i', function.sum(self.ns.a2 * self.ns.a2, axis=0))

From 672facdff101830b09ce896ee5c1c5f28e08a323 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Wed, 7 Apr 2021 12:50:06 +0200
Subject: [PATCH 02/26] fix dtype of Normal.evalf for 1d argument

The `evaluable.Normal` class has a fast implementation of `evalf` for
one-dimensional geometries. This implementation, however, returns an array with
the same dtype as the argument of `evalf`, the local gradient of the geometry,
which could be different from the dtype announced by `evaluable.Normal`:
`float`. This patch fixes the problem by adding an explicit cast to a `float`.
---
 nutils/evaluable.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 47a7d21be..283a53521 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -1088,7 +1088,7 @@ def __init__(self, lgrad:asarray):
   def evalf(self, lgrad):
     n = lgrad[...,-1]
     if equalindex(n.shape[-1], 1): # geom is 1D
-      return numpy.sign(n)
+      return numpy.sign(n).astype(float)
     # orthonormalize n to G
     G = lgrad[...,:-1]
     GG = numeric.contract(G[...,:,_,:], G[...,:,:,_], axis=-3)

From 50a7a304ebf48fc61afd722b3935f49f50b097cd Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Wed, 14 Apr 2021 14:35:03 +0200
Subject: [PATCH 03/26] fix MKL tests in Github Actions

The most recent version of MKL on PyPI ships only `libmkl_rt.so.1`, not
`libmkl_rt.so`. In Nutils we load the latter. This causes two problems: in
Github Actions we cannot find the location of the library directory where MKL
is being installed, because we explicitly look for `libmkl_rt.so` and Nutils
can't load the library because we explicitly load the missing unversioned
library. This patch fixes both problems by searching for the versioned library
and adding a symlink to the unversioned library.
---
 .github/workflows/test.yaml   |  7 +++----
 devtools/gha/configure_mkl.py | 25 +++++++++++++++++++++++++
 devtools/gha/get_lib_dir.py   | 16 ----------------
 3 files changed, 28 insertions(+), 20 deletions(-)
 create mode 100644 devtools/gha/configure_mkl.py
 delete mode 100644 devtools/gha/get_lib_dir.py

diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 603953625..4a0d36531 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -100,13 +100,12 @@ jobs:
           # Install Nutils from `dist` dir created in job
           # `build-python-package`.
           python -um pip install --no-index --find-links ./dist nutils
-      - name: Get library directory
-        id: get-lib-dir
+      - name: Configure MKL
+        id: configure-mkl
         if: ${{ matrix.matrix-backend == 'mkl' }}
-        run: python -um devtools.gha.get_lib_dir
+        run: python -um devtools.gha.configure_mkl
       - name: Test
         env:
-          LD_LIBRARY_PATH: ${{ steps.get-lib-dir.outputs.libdir }}
           NUTILS_DEBUG: all
         run: |
           mkdir testenv
diff --git a/devtools/gha/configure_mkl.py b/devtools/gha/configure_mkl.py
new file mode 100644
index 000000000..9ecff944e
--- /dev/null
+++ b/devtools/gha/configure_mkl.py
@@ -0,0 +1,25 @@
+from .. import log
+import sys, site, os
+
+libsubdir = 'lib'
+prefixes = list(site.PREFIXES)
+if hasattr(site, 'getuserbase'):
+  prefixes.append(site.getuserbase())
+for prefix in prefixes:
+  libdir = os.path.join(prefix, libsubdir)
+  if not os.path.exists(os.path.join(libdir, 'libmkl_rt.so.1')):
+    continue
+  log.set_output('libdir', libdir)
+  break
+else:
+  log.error('cannot find {} in any of the following dirs: {}'.format('libmkl_rt.so.1', ' '.join(prefixes)))
+
+lib = os.path.join(libdir, 'libmkl_rt.so')
+if not os.path.exists(lib):
+  os.symlink('libmkl_rt.so.1', lib)
+
+github_env = os.environ.get('GITHUB_ENV')
+assert github_env
+ld_library_path = ':'.join(filter(None, (os.environ.get('LD_LIBRARY_PATH', ''), libdir)))
+with open(github_env, 'a') as f:
+  print('LD_LIBRARY_PATH={}'.format(ld_library_path), file=f)
diff --git a/devtools/gha/get_lib_dir.py b/devtools/gha/get_lib_dir.py
deleted file mode 100644
index 9e411ab92..000000000
--- a/devtools/gha/get_lib_dir.py
+++ /dev/null
@@ -1,16 +0,0 @@
-from .. import log
-import sys, site, os
-
-libname = 'libmkl_rt.so'
-libsubdir = 'lib'
-prefixes = list(site.PREFIXES)
-if hasattr(site, 'getuserbase'):
-  prefixes.append(site.getuserbase())
-for prefix in prefixes:
-  libdir = os.path.join(prefix, libsubdir)
-  if not os.path.exists(os.path.join(libdir, libname)):
-    continue
-  log.set_output('libdir', libdir)
-  break
-else:
-  log.error('cannot find {} in any of the following dirs: {}'.format(libname, ' '.join(prefixes)))

From 41db6c45b256e30e7ecd9db2f31854f407779638 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Wed, 14 Apr 2021 14:45:42 +0200
Subject: [PATCH 04/26] fix rootless Buildah mount prob. in Github Actions

Due to a recent change in Buildah and a configuration error in the Github
virtual environment, Buildah cannot mount layers in rootless mode. This patch
fixes the problem by explicitly specifying the correct mount program to use, as
suggested in [this bug
report](https://github.com/actions/virtual-environments/issues/3080).
---
 .github/workflows/test.yaml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 4a0d36531..361885608 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -134,6 +134,9 @@ jobs:
     name: Build container image
     needs: build-python-package
     runs-on: ubuntu-20.04
+    env:
+      # Fixes https://github.com/actions/virtual-environments/issues/3080
+      STORAGE_OPTS: overlay.mount_program=/usr/bin/fuse-overlayfs
     steps:
       - name: Checkout
         uses: actions/checkout@v2

From 00218c1b3f27924ab8dfc2747fc1e9e94cc63c30 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Thu, 15 Apr 2021 21:43:34 +0200
Subject: [PATCH 05/26] fix invalid length error with meshio.CellBlock

In commit [60d11b70] of meshio the length of `CellBlock`, a subclass of
`namedtuple`, is redefined to the length of the `data` attribute. This breaks
`namedtuple._replace`, where the length of the newly created `CellBlock` is
verified. This patch circumvents the problem by using the `CellBlock`
constructor directly instead of `namedtuple._replace`.

[60d11b70]: https://github.com/nschloe/meshio/commit/60d11b705c9b9074284feefdb7475db0410953e9
---
 nutils/mesh.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nutils/mesh.py b/nutils/mesh.py
index 11e86eaf5..64e8cbe54 100644
--- a/nutils/mesh.py
+++ b/nutils/mesh.py
@@ -357,7 +357,7 @@ def parsegmsh(mshdata):
       if keep.all():
         renum = numpy.arange(len(cells.data))
       else:
-        msh.cells[icell] = cells._replace(data=cells.data[numpy.hstack([True, keep])])
+        msh.cells[icell] = type(cells)(type=cells.type, data=cells.data[numpy.hstack([True, keep])])
         renum = numpy.hstack([0, keep.cumsum()])
       renums.append(renum)
     for name, (itag, nd) in msh.field_data.items():

From ea3a49330d849a078f524aace5a9ca39647e14a5 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Wed, 14 Apr 2021 14:40:44 +0200
Subject: [PATCH 06/26] include Python 3.9 in Github Actions

---
 .github/workflows/test.yaml | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 361885608..1d5b5605c 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -39,20 +39,21 @@ jobs:
       matrix:
         include:
           # base
-          - {os: ubuntu-latest , python-version: 3.8, matrix-backend: numpy, nprocs: 1, numpy-version: latest}
+          - {os: ubuntu-latest , python-version: 3.9, matrix-backend: numpy, nprocs: 1, numpy-version: latest}
           # os
-          - {os: windows-latest, python-version: 3.8, matrix-backend: numpy, nprocs: 1, numpy-version: latest}
-          - {os: macos-latest  , python-version: 3.8, matrix-backend: numpy, nprocs: 1, numpy-version: latest}
+          - {os: windows-latest, python-version: 3.9, matrix-backend: numpy, nprocs: 1, numpy-version: latest}
+          - {os: macos-latest  , python-version: 3.9, matrix-backend: numpy, nprocs: 1, numpy-version: latest}
           # python-version
           - {os: ubuntu-latest , python-version: 3.5, matrix-backend: numpy, nprocs: 1, numpy-version: latest}
           - {os: ubuntu-latest , python-version: 3.6, matrix-backend: numpy, nprocs: 1, numpy-version: latest}
           - {os: ubuntu-latest , python-version: 3.7, matrix-backend: numpy, nprocs: 1, numpy-version: latest}
+          - {os: ubuntu-latest , python-version: 3.8, matrix-backend: numpy, nprocs: 1, numpy-version: latest}
           # matrix-backend
-          - {os: ubuntu-latest , python-version: 3.8, matrix-backend: scipy, nprocs: 1, numpy-version: latest}
-          - {os: ubuntu-latest , python-version: 3.8, matrix-backend: mkl  , nprocs: 1, numpy-version: latest}
-          - {os: ubuntu-latest , python-version: 3.8, matrix-backend: mkl  , nprocs: 2, numpy-version: latest}
+          - {os: ubuntu-latest , python-version: 3.9, matrix-backend: scipy, nprocs: 1, numpy-version: latest}
+          - {os: ubuntu-latest , python-version: 3.9, matrix-backend: mkl  , nprocs: 1, numpy-version: latest}
+          - {os: ubuntu-latest , python-version: 3.9, matrix-backend: mkl  , nprocs: 2, numpy-version: latest}
           # nprocs
-          - {os: ubuntu-latest , python-version: 3.8, matrix-backend: numpy, nprocs: 2, numpy-version: latest}
+          - {os: ubuntu-latest , python-version: 3.9, matrix-backend: numpy, nprocs: 2, numpy-version: latest}
           # numpy-version
           - {os: ubuntu-latest , python-version: 3.8, matrix-backend: numpy, nprocs: 1, numpy-version: 1.16  }
       fail-fast: false

From b8edba005844755d38da9b0e7fd895e0e75279d2 Mon Sep 17 00:00:00 2001
From: Gertjan van Zwieten <gertjanvanzwieten@evalf.nl>
Date: Wed, 24 Mar 2021 21:03:31 +0100
Subject: [PATCH 07/26] pull add through loopsum iff either arg is const

---
 nutils/evaluable.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 283a53521..687775746 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -1834,7 +1834,8 @@ def _unravel(self, axis, shape):
     return Add([unravel(func, axis, shape) for func in self.funcs])
 
   def _loopsum(self, index, length):
-    return Add([LoopSum(func, index, length) for func in self.funcs])
+    if any(index not in func.arguments for func in self.funcs):
+      return Add([LoopSum(func, index, length) for func in self.funcs])
 
   def _desparsify(self, axis):
     assert isinstance(self._axes[axis], Sparse)

From 1be2d6f1baa994c23cd8f7c3954d64c4b0e9846f Mon Sep 17 00:00:00 2001
From: Gertjan van Zwieten <gertjanvanzwieten@evalf.nl>
Date: Thu, 8 Apr 2021 13:55:09 +0200
Subject: [PATCH 08/26] implement LoopSum._sum, ._add

---
 nutils/evaluable.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 687775746..efc800fcf 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -3391,6 +3391,13 @@ def _take(self, index, axis):
   def _unravel(self, axis, shape):
     return LoopSum(unravel(self.func, axis, shape), self.index, self.length)
 
+  def _sum(self, axis):
+    return LoopSum(sum(self.func, axis), self.index, self.length)
+
+  def _add(self, other):
+    if isinstance(other, LoopSum) and other.length == self.length and other.index == self.index:
+      return LoopSum(self.func + other.func, self.index, self.length)
+
   @property
   def _assparse(self):
     chunks = []

From b40bee17248b55751326e4b67b083ea06b5bbd1f Mon Sep 17 00:00:00 2001
From: Gertjan van Zwieten <gertjanvanzwieten@evalf.nl>
Date: Thu, 8 Apr 2021 13:36:01 +0200
Subject: [PATCH 09/26] simplify sum to take for length=1 axes

---
 nutils/evaluable.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index efc800fcf..024a8d35f 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -1918,6 +1918,8 @@ def __init__(self, func:asarray):
     super().__init__(args=[func], shape=shape, dtype=int if func.dtype == bool else func.dtype)
 
   def _simplified(self):
+    if equalindex(self.func.shape[-1], 1):
+      return Take(self.func, 0)
     return self.func._sum(self.ndim)
 
   def evalf(self, arr):

From 53522d9e0487c5bf3eb7ba518324a69dbdb02785 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Wed, 7 Apr 2021 12:57:20 +0200
Subject: [PATCH 10/26] simplify Normal of 1D geometry to Sign

The normal at the edge of a line is the sign of the local gradient of the
geometry. The `evaluable.Normal` has a fast implementation of `evalf` for this
situation. This patch replaces the fast evalf with a simplification of the same
kind in case the dimension of the geometry is constant one and adds unittests.
---
 nutils/evaluable.py     | 6 ++++--
 tests/test_evaluable.py | 3 +++
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 024a8d35f..5204947f0 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -1085,10 +1085,12 @@ def __init__(self, lgrad:asarray):
     self.lgrad = lgrad
     super().__init__(args=[lgrad], shape=lgrad.shape[:-1], dtype=float)
 
+  def _simplified(self):
+    if equalindex(self.shape[-1], 1):
+      return Sign(Take(self.lgrad, 0))
+
   def evalf(self, lgrad):
     n = lgrad[...,-1]
-    if equalindex(n.shape[-1], 1): # geom is 1D
-      return numpy.sign(n).astype(float)
     # orthonormalize n to G
     G = lgrad[...,:-1]
     GG = numeric.contract(G[...,:,_,:], G[...,:,:,_], axis=-3)
diff --git a/tests/test_evaluable.py b/tests/test_evaluable.py
index d0bb745ce..89f190453 100644
--- a/tests/test_evaluable.py
+++ b/tests/test_evaluable.py
@@ -439,6 +439,9 @@ def _check(name, op, n_op, shapes, hasgrad=True, zerograd=False, ndim=2, low=-1,
 _check('take-duplicate', lambda f: evaluable.Take(f, [0,3,0]), lambda a: a[:,[0,3,0]], [(2,4)])
 _check('choose', lambda a, b, c: evaluable.Choose(evaluable.appendaxes(evaluable.Int(a)%2, (3,3)), [b,c]), lambda a, b, c: numpy.choose(a[_,_].astype(int)%2, [b,c]), [(), (3,3), (3,3)])
 _check('slice', lambda a: evaluable.asarray(a)[::2], lambda a: a[::2], [(5,3)])
+_check('normal1d', lambda a: evaluable.Normal(a), lambda a: numpy.sign(a[...,0]), [(3,1,1)])
+_check('normal2d', lambda a: evaluable.Normal(a), lambda a: numpy.stack([Q[:,-1]*numpy.sign(R[-1,-1]) for ai in a for Q, R in [numpy.linalg.qr(ai, mode='complete')]], axis=0), [(1,2,2)])
+_check('normal3d', lambda a: evaluable.Normal(a), lambda a: numpy.stack([Q[:,-1]*numpy.sign(R[-1,-1]) for ai in a for Q, R in [numpy.linalg.qr(ai, mode='complete')]], axis=0), [(2,3,3)])
 _check('loopsum1', lambda: evaluable.LoopSum(evaluable.Argument('index', (), int), evaluable.Argument('index', (), int), 3), lambda: numpy.array(3), [])
 _check('loopsum2', lambda a: evaluable.LoopSum(a, evaluable.Argument('index', (), int), 2), lambda a: 2*a, [(3,4,2,4)])
 _check('loopsum3', lambda a: evaluable.LoopSum(evaluable.get(a, 0, evaluable.Argument('index', (), int)), evaluable.Argument('index', (), int), 3), lambda a: numpy.sum(a, 0), [(3,4,2,4)])

From 84cdb3d49060218fc434a37a1c49684dc8f7e092 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Thu, 8 Apr 2021 22:51:05 +0200
Subject: [PATCH 11/26] simplify repeated scalar inflate to diagonalize

This patch simplifies nested `Inflate` objects with the same scalar dofmap and
length to a `Diagonalize` of the inner `Inflate`. This simplification enables
two additional simplifications: a sum of diagonal (nested) inflates becomes a
diagonal of sum of inflates and the inverse of this is the reciprocal.
---
 nutils/evaluable.py     | 2 ++
 tests/test_evaluable.py | 1 +
 2 files changed, 3 insertions(+)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 5204947f0..690917201 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -2556,6 +2556,8 @@ def _desparsify(self, axis):
   def _inflate(self, dofmap, length, axis):
     if dofmap.ndim == 1 and axis == self.ndim-1:
       return Inflate(self.func, Take(dofmap, self.dofmap), length)
+    if dofmap.ndim == 0 and dofmap == self.dofmap and length == self.length:
+      return diagonalize(self, -1, axis)
 
   def _derivative(self, var, seen):
     return _inflate(derivative(self.func, var, seen), self.dofmap, self.length, self.ndim-1)
diff --git a/tests/test_evaluable.py b/tests/test_evaluable.py
index 89f190453..d3dc7b701 100644
--- a/tests/test_evaluable.py
+++ b/tests/test_evaluable.py
@@ -435,6 +435,7 @@ def _check(name, op, n_op, shapes, hasgrad=True, zerograd=False, ndim=2, low=-1,
 _check('inflate-duplicate', lambda f: evaluable.Inflate(f,dofmap=[0,1,0,3],length=4), lambda a: numpy.stack([a[:,0]+a[:,2], a[:,1], numpy.zeros_like(a[:,0]), a[:,3]], axis=1), [(2,4)])
 _check('inflate-block', lambda f: evaluable.Inflate(f,dofmap=[[5,4,3],[2,1,0]],length=6), lambda a: a.ravel()[::-1], [(2,3)])
 _check('inflate-scalar', lambda f: evaluable.Inflate(f,dofmap=1,length=3), lambda a: numpy.array([0,a,0]), [()])
+_check('inflate-diagonal', lambda f: evaluable.Inflate(evaluable.Inflate(f,1,3),1,3), lambda a: numpy.diag(numpy.array([0,a,0])), [()])
 _check('take', lambda f: evaluable.Take(f, [0,3,2]), lambda a: a[:,[0,3,2]], [(2,4)])
 _check('take-duplicate', lambda f: evaluable.Take(f, [0,3,0]), lambda a: a[:,[0,3,0]], [(2,4)])
 _check('choose', lambda a, b, c: evaluable.Choose(evaluable.appendaxes(evaluable.Int(a)%2, (3,3)), [b,c]), lambda a, b, c: numpy.choose(a[_,_].astype(int)%2, [b,c]), [(), (3,3), (3,3)])

From 9d19d1b3a175d5b794cd2f121fe247e1cef3f151 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Fri, 9 Apr 2021 10:39:09 +0200
Subject: [PATCH 12/26] simplify Inflate of length one to InsertAxis

This patch simplifies an `Inflate` of length one with scalar dofmap to an
`InsertAxis`.
---
 nutils/evaluable.py     | 2 ++
 tests/test_evaluable.py | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 690917201..20236abd6 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -2536,6 +2536,8 @@ def _simplified(self):
       if isinstance(self.func._axes[self.func.ndim-self.dofmap.ndim+axis], Sparse):
         items = [Inflate(f, _take(self.dofmap, ind, axis), self.shape[-1]) for ind, f in self.func._desparsify(self.func.ndim-self.dofmap.ndim+axis)]
         return util.sum(items) if items else zeros_like(self)
+    if self.dofmap.ndim == 0 and equalindex(self.dofmap, 0) and equalindex(self.length, 1):
+      return InsertAxis(self.func, 1)
     return self.func._inflate(self.dofmap, self.length, self.ndim-1)
 
   def evalf(self, array, indices, length):
diff --git a/tests/test_evaluable.py b/tests/test_evaluable.py
index d3dc7b701..cabe3db3c 100644
--- a/tests/test_evaluable.py
+++ b/tests/test_evaluable.py
@@ -436,6 +436,8 @@ def _check(name, op, n_op, shapes, hasgrad=True, zerograd=False, ndim=2, low=-1,
 _check('inflate-block', lambda f: evaluable.Inflate(f,dofmap=[[5,4,3],[2,1,0]],length=6), lambda a: a.ravel()[::-1], [(2,3)])
 _check('inflate-scalar', lambda f: evaluable.Inflate(f,dofmap=1,length=3), lambda a: numpy.array([0,a,0]), [()])
 _check('inflate-diagonal', lambda f: evaluable.Inflate(evaluable.Inflate(f,1,3),1,3), lambda a: numpy.diag(numpy.array([0,a,0])), [()])
+_check('inflate-one', lambda f: evaluable.Inflate(f,0,1), lambda a: numpy.array([a]), [()])
+_check('inflate-range', lambda f: evaluable.Inflate(f,evaluable.Range(3),3), lambda a: a, [(3,)])
 _check('take', lambda f: evaluable.Take(f, [0,3,2]), lambda a: a[:,[0,3,2]], [(2,4)])
 _check('take-duplicate', lambda f: evaluable.Take(f, [0,3,0]), lambda a: a[:,[0,3,0]], [(2,4)])
 _check('choose', lambda a, b, c: evaluable.Choose(evaluable.appendaxes(evaluable.Int(a)%2, (3,3)), [b,c]), lambda a, b, c: numpy.choose(a[_,_].astype(int)%2, [b,c]), [(), (3,3), (3,3)])

From f7c0bf18f144da0158cbf45fb922a750359274ed Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Wed, 7 Apr 2021 14:23:21 +0200
Subject: [PATCH 13/26] drop offset from evaluable.InRange

The `InRange` evaluable takes `length`, `offset` and `index` as arguments. In
`evalf` the `index` is asserted to be within 0 and `length` after which the
`offset` is added to `index` and returned. This patch removes the `offset`
argument in favor of an explicit `Add(InRange(...), offset)` and reorders the
remaining to `index` and `length`.
---
 nutils/evaluable.py | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 20236abd6..e6be6f63c 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -3094,7 +3094,7 @@ def _take(self, index, axis):
     if index.ndim == 1 and isinstance(index, Range) and self.isconstant and index.isconstant:
       assert (index.offset + index.length).eval() <= self.length.eval()
       return Range(index.length, self.offset + index.offset)
-    return InRange(self.length, self.offset, index)
+    return InRange(index, self.length) + self.offset
 
   def _add(self, offset):
     if isinstance(self._axes[0], Inserted):
@@ -3105,18 +3105,17 @@ def evalf(self, length, offset):
 
 class InRange(Array):
 
-  __slots__ = 'length', 'offset', 'index'
+  __slots__ = 'index', 'length'
 
   @types.apply_annotations
-  def __init__(self, length:asarray, offset:asarray, index:asarray):
-    self.length = length
-    self.offset = offset
+  def __init__(self, index:asarray, length:asarray):
     self.index = index
-    super().__init__(args=[length, offset, index], shape=index.shape, dtype=int)
+    self.length = length
+    super().__init__(args=[index, length], shape=index.shape, dtype=int)
 
-  def evalf(self, length, offset, index):
+  def evalf(self, index, length):
     assert index.size == 0 or 0 <= index.min() and index.max() < length
-    return index + offset
+    return index
 
 class Polyval(Array):
   '''

From facd4745c9fa0933c34dcf360039bba039dd53b9 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Tue, 20 Apr 2021 11:15:00 +0200
Subject: [PATCH 14/26] drop offset from evaluable.Range

The `Range` evaluable currently takes `length` and `offset` as arguments. To
simplify the evaluable, this patch drops the `offset` argument in favor of an
explicit `Add(Range(...), offset)`.
---
 nutils/evaluable.py     | 24 ++++++++----------------
 nutils/function.py      |  6 +++---
 nutils/sample.py        |  2 +-
 tests/test_evaluable.py |  2 +-
 4 files changed, 13 insertions(+), 21 deletions(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index e6be6f63c..1ecbd19c9 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -2380,7 +2380,7 @@ def Elemwise(data:types.tuple[types.arraydata], index:asarray, dtype:asdtype):
   for i in reversed(range(ndim-1)):
     cumprod[i] *= cumprod[i+1] # work backwards so that the shape check matches in Unravel
   offsets = numpy.cumsum([0, *map(len, raveled[:-1])])
-  elemwise = Take(concat, Range(cumprod[0], Take(offsets, index)))
+  elemwise = Take(concat, Range(cumprod[0]) + Take(offsets, index))
   for i in range(ndim-1):
     elemwise = Unravel(elemwise, shape[i], cumprod[i+1])
   return elemwise
@@ -3082,26 +3082,18 @@ def _assparse(self):
 
 class Range(Array):
 
-  __slots__ = 'length', 'offset'
+  __slots__ = 'length'
 
   @types.apply_annotations
-  def __init__(self, length:asindex, offset:asindex=Zeros((), int)):
+  def __init__(self, length:asindex):
     self.length = length
-    self.offset = offset
-    super().__init__(args=[length, offset], shape=[length], dtype=int)
+    super().__init__(args=[length], shape=[length], dtype=int)
 
   def _take(self, index, axis):
-    if index.ndim == 1 and isinstance(index, Range) and self.isconstant and index.isconstant:
-      assert (index.offset + index.length).eval() <= self.length.eval()
-      return Range(index.length, self.offset + index.offset)
-    return InRange(index, self.length) + self.offset
-
-  def _add(self, offset):
-    if isinstance(self._axes[0], Inserted):
-      return Range(self.length, self.offset + offset._uninsert(0))
+    return InRange(index, self.length)
 
-  def evalf(self, length, offset):
-    return numpy.arange(offset, offset+length)
+  def evalf(self, length):
+    return numpy.arange(length)
 
 class InRange(Array):
 
@@ -3899,7 +3891,7 @@ def _takeslice(arg:asarray, s:types.strict[slice], axis:types.strictint):
     stop = n if s.stop is None else s.stop if s.stop >= 0 else s.stop + n
     if start == 0 and stop == n:
       return arg
-    index = Range(stop-start, start)
+    index = Range(stop-start) + start
   elif n.isconstant:
     index = Constant(numpy.arange(*s.indices(arg.shape[axis])))
   else:
diff --git a/nutils/function.py b/nutils/function.py
index bcf1a457d..a6a4e96da 100644
--- a/nutils/function.py
+++ b/nutils/function.py
@@ -1896,7 +1896,7 @@ def get(__array: IntoArray, __axis: int, __index: IntoArray) -> Array:
 def _range(__length: IntoArray, __offset: IntoArray) -> Array:
   length = Array.cast(__length, dtype=int, ndim=0)
   offset = Array.cast(__offset, dtype=int, ndim=0)
-  return _Wrapper(evaluable.Range, _WithoutPoints(length), _WithoutPoints(offset), shape=(length,), dtype=int)
+  return _Wrapper(lambda l, o: evaluable.Range(l) + o, _WithoutPoints(length), _WithoutPoints(offset), shape=(length,), dtype=int)
 
 def _takeslice(__array: IntoArray, __s: slice, __axis: int) -> Array:
   array = Array.cast(__array)
@@ -2772,7 +2772,7 @@ def get_support(self, dof: Union[int, numpy.ndarray]) -> numpy.ndarray:
 
   def f_dofs_coeffs(self, index: evaluable.Array) -> Tuple[evaluable.Array,evaluable.Array]:
     coeffs = evaluable.Elemwise(self._coeffs, index, dtype=float)
-    dofs = evaluable.Range(coeffs.shape[0], evaluable.get(self._offsets, 0, index))
+    dofs = evaluable.Range(coeffs.shape[0]) + evaluable.get(self._offsets, 0, index)
     return dofs, coeffs
 
 class MaskedBasis(Basis):
@@ -2875,7 +2875,7 @@ def f_dofs_coeffs(self, index: evaluable.Array) -> Tuple[evaluable.Array,evaluab
     for lengths_i, offsets_i, ndofs_i, index_i in zip(self._ndofs, self._start_dofs, self._dofs_shape, indices):
       length = evaluable.get(lengths_i, 0, index_i)
       offset = evaluable.get(offsets_i, 0, index_i)
-      dofs_i = evaluable.Range(length, offset) % ndofs_i
+      dofs_i = (evaluable.Range(length) + offset) % ndofs_i
       dofs = dofs_i if dofs is None else evaluable.ravel(evaluable.insertaxis(dofs * ndofs_i, 1, length) + dofs_i, axis=0)
     coeffs = functools.reduce(evaluable.PolyOuterProduct,
       [evaluable.Elemwise(coeffs_i, index_i, float) for coeffs_i, index_i in zip(self._coeffs, indices)])
diff --git a/nutils/sample.py b/nutils/sample.py
index 81081d319..369bf51fd 100644
--- a/nutils/sample.py
+++ b/nutils/sample.py
@@ -338,7 +338,7 @@ def hull(self):
   def get_evaluable_indices(self, ielem):
     npoints = self.points.get_evaluable_coords(ielem).shape[0]
     offset = evaluable.get(_offsets(self.points), 0, ielem)
-    return evaluable.Range(npoints, offset)
+    return evaluable.Range(npoints) + offset
 
 class _CustomIndex(Sample):
 
diff --git a/tests/test_evaluable.py b/tests/test_evaluable.py
index cabe3db3c..1ca38fa7f 100644
--- a/tests/test_evaluable.py
+++ b/tests/test_evaluable.py
@@ -369,7 +369,7 @@ def _check(name, op, n_op, shapes, hasgrad=True, zerograd=False, ndim=2, low=-1,
 _check('const', lambda f: evaluable.asarray(numpy.arange(16, dtype=float).reshape(2,4,2)), lambda a: numpy.arange(16, dtype=float).reshape(2,4,2), [(2,4,2)])
 _check('zeros', lambda f: evaluable.zeros([1,4,3,4]), lambda a: numpy.zeros([1,4,3,4]), [(4,3,4)])
 _check('ones', lambda f: evaluable.ones([1,4,3,4]), lambda a: numpy.ones([1,4,3,4]), [(4,3,4)])
-_check('range', lambda f: evaluable.Range(4, offset=2)[numpy.newaxis], lambda a: numpy.arange(2,6)[numpy.newaxis], [(4,)])
+_check('range', lambda f: evaluable.Range(4) + 2, lambda a: numpy.arange(2,6), [(4,)])
 _check('sin', evaluable.sin, numpy.sin, [(4,)])
 _check('cos', evaluable.cos, numpy.cos, [(4,)])
 _check('tan', evaluable.tan, numpy.tan, [(4,)])

From e24afe4ae829f363786e8d114e9e60d883151990 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Wed, 17 Mar 2021 23:53:23 +0100
Subject: [PATCH 15/26] replace all calls to `prepare_eval` with `lower`

This patch replaces all calls to the old `prepare_eval` with `lower`, awaiting
deprecation when `lower` is finalized.
---
 nutils/topology.py      |  21 ++++---
 tests/test_evaluable.py |  15 +++--
 tests/test_function.py  | 124 +++++++++++++++++++++++-----------------
 tests/test_sample.py    |   2 +-
 tests/test_topology.py  |   5 +-
 5 files changed, 95 insertions(+), 72 deletions(-)

diff --git a/nutils/topology.py b/nutils/topology.py
index 6264131e1..a37db65d5 100644
--- a/nutils/topology.py
+++ b/nutils/topology.py
@@ -308,11 +308,12 @@ def project(self, fun, onto, geometry, ischeme='gauss', degree=None, droptol=1e-
       F = numpy.zeros(onto.shape[0])
       W = numpy.zeros(onto.shape[0])
       I = numpy.zeros(onto.shape[0], dtype=bool)
-      fun = function.asarray(fun).prepare_eval(ndims=self.ndims)
-      data = evaluable.Tuple(evaluable.Tuple([fun, onto_f.simplified, evaluable.Tuple(onto_ind)]) for onto_ind, onto_f in evaluable.blocks(onto.prepare_eval(ndims=self.ndims)))
-      for ref, trans, opp in zip(self.references, self.transforms, self.opposites):
-        ipoints = ref.getpoints('bezier2')
-        for fun_, onto_f_, onto_ind_ in data.eval(_transforms=(trans, opp), _points=ipoints, **arguments or {}):
+      ielem_arg = evaluable.Argument('_project_index', (), dtype=int)
+      lower_args = dict(transform_chains=(evaluable.TransformChainFromSequence(self.transforms, ielem_arg), evaluable.TransformChainFromSequence(self.opposites, ielem_arg)), coordinates=(self.references.getpoints('bezier', 2).get_evaluable_coords(ielem_arg),)*2)
+      fun = function.lower(**lower_args)
+      data = evaluable.Tuple(evaluable.Tuple([fun, onto_f.simplified, evaluable.Tuple(onto_ind)]) for onto_ind, onto_f in evaluable.blocks(onto.lower(**lower_args))).optimized_for_numpy
+      for ielem in range(len(self)):
+        for fun_, onto_f_, onto_ind_ in data.eval(_project_index=ielem, **arguments or {}):
           onto_f_ = onto_f_.swapaxes(0,1) # -> dof axis, point axis, ...
           indfun_ = fun_[(slice(None),)+numpy.ix_(*onto_ind_[1:])]
           assert onto_f_.shape[0] == len(onto_ind_[0])
@@ -361,15 +362,17 @@ def trim(self, levelset, maxrefine, ndivisions=8, name='trimmed', leveltopo=None
     if arguments is None:
       arguments = {}
 
-    levelset = levelset.prepare_eval(ndims=self.ndims).optimized_for_numpy
     refs = []
     if leveltopo is None:
-      with log.iter.percentage('trimming', self.references, self.transforms, self.opposites) as items:
-        for ref, trans, opp in items:
-          levels = levelset.eval(_transforms=(trans, opp), _points=ref.getpoints('vertex', maxrefine), **arguments)
+      ielem_arg = evaluable.Argument('_trim_index', (), dtype=int)
+      levelset = levelset.lower(transform_chains=(evaluable.TransformChainFromSequence(self.transforms, ielem_arg), evaluable.TransformChainFromSequence(self.opposites, ielem_arg)), coordinates=(self.references.getpoints('vertex', maxrefine).get_evaluable_coords(ielem_arg),)*2).optimized_for_numpy
+      with log.iter.percentage('trimming', range(len(self)), self.references) as items:
+        for ielem, ref in items:
+          levels = levelset.eval(_trim_index=ielem, **arguments)
           refs.append(ref.trim(levels, maxrefine=maxrefine, ndivisions=ndivisions))
     else:
       log.info('collecting leveltopo elements')
+      levelset = levelset.lower(transform_chains=(evaluable.SelectChain(0), evaluable.SelectChain(1)), coordinates=(evaluable.Points(evaluable.NPoints(), self.ndims),)*2).optimized_for_numpy
       bins = [set() for ielem in range(len(self))]
       for trans in leveltopo.transforms:
         ielem, tail = self.transforms.index_with_tail(trans)
diff --git a/tests/test_evaluable.py b/tests/test_evaluable.py
index 1ca38fa7f..99da0d4d3 100644
--- a/tests/test_evaluable.py
+++ b/tests/test_evaluable.py
@@ -576,8 +576,7 @@ class elemwise(TestCase):
 
   def setUp(self):
     super().setUp()
-    self.domain, geom = mesh.rectilinear([7])
-    self.index = self.domain.f_index.prepare_eval(ndims=self.domain.ndims, npoints=None)
+    self.index = evaluable.Argument('index', (), int)
     self.data = tuple(map(types.frozenarray, (
       numpy.arange(1, dtype=float).reshape(1,1),
       numpy.arange(2, dtype=float).reshape(1,2),
@@ -590,20 +589,20 @@ def setUp(self):
     self.func = evaluable.Elemwise(self.data, self.index, float)
 
   def test_evalf(self):
-    for i, (trans, points) in enumerate(zip(self.domain.transforms, self.domain.sample('gauss', 1).points)):
+    for i in range(7):
       with self.subTest(i=i):
-        numpy.testing.assert_array_almost_equal(self.func.eval(_transforms=(trans,), _points=points), self.data[i])
+        numpy.testing.assert_array_almost_equal(self.func.eval(index=i), self.data[i])
 
   def test_shape(self):
-    for i, (trans, points) in enumerate(zip(self.domain.transforms, self.domain.sample('gauss', 1).points)):
+    for i in range(7):
       with self.subTest(i=i):
-        self.assertEqual(self.func.size.eval(_transforms=(trans,), _points=points), self.data[i].size)
+        self.assertEqual(self.func.size.eval(index=i), self.data[i].size)
 
   def test_derivative(self):
-    self.assertTrue(evaluable.iszero(evaluable.localgradient(self.func, self.domain.ndims).simplified))
+    self.assertTrue(evaluable.iszero(evaluable.localgradient(self.func, 2).simplified))
 
   def test_shape_derivative(self):
-    self.assertEqual(evaluable.localgradient(self.func, self.domain.ndims).shape, self.func.shape+(evaluable.Constant(self.domain.ndims),))
+    self.assertEqual(evaluable.localgradient(self.func, 2).shape, self.func.shape+(evaluable.Constant(2),))
 
 class jacobian(TestCase):
 
diff --git a/tests/test_function.py b/tests/test_function.py
index a2c1afebb..4dc5b5179 100644
--- a/tests/test_function.py
+++ b/tests/test_function.py
@@ -29,7 +29,7 @@ def test_shape(self):
     with self.subTest('const Array'):
       self.assertEqual(l2, 3)
     with self.subTest('unknown'):
-      self.assertEqual(l3.prepare_eval(npoints=None).simplified, n.prepare_eval(npoints=None).simplified)
+      self.assertEqual(l3.as_evaluable_array().simplified, n.as_evaluable_array().simplified)
 
   def test_size_known(self):
     self.assertEqual(function.Argument('a', (2,3)).size, 6)
@@ -41,7 +41,7 @@ def test_size_unknown(self):
     n = function.Argument('n', (), dtype=int)
     size = function.Argument('a', (2,n,3)).size
     self.assertIsInstance(size, function.Array)
-    self.assertEqual(size.prepare_eval(npoints=None).simplified, (2*n*3).prepare_eval(npoints=None).simplified)
+    self.assertEqual(size.as_evaluable_array().simplified, (2*n*3).as_evaluable_array().simplified)
 
   def test_len_0d(self):
     with self.assertRaisesRegex(Exception, '^len\\(\\) of unsized object$'):
@@ -60,8 +60,8 @@ def test_iter_0d(self):
 
   def test_iter_known(self):
     a, b = function.Array.cast([1,2])
-    self.assertEqual(a.prepare_eval(npoints=None).eval(), 1)
-    self.assertEqual(b.prepare_eval(npoints=None).eval(), 2)
+    self.assertEqual(a.as_evaluable_array().eval(), 1)
+    self.assertEqual(b.as_evaluable_array().eval(), 2)
 
   def test_iter_unknown(self):
     with self.assertRaisesRegex(Exception, '^iteration over array with unknown length$'):
@@ -79,6 +79,21 @@ def test_deprecated_simplified(self):
     with self.assertWarns(warnings.NutilsDeprecationWarning):
       function.Array.cast([1,2]).simplified
 
+  def test_prepare_eval(self):
+    topo, geom = mesh.rectilinear([2])
+    f = topo.basis('discont', 0).dot([1, 2])
+    lowered = f.prepare_eval(ndims=topo.ndims, opposite=True)
+    smpl = topo.sample('bezier', 2)
+    with _builtin_warnings.catch_warnings():
+      _builtin_warnings.simplefilter('ignore', category=evaluable.ExpensiveEvaluationWarning)
+      self.assertAllAlmostEqual(lowered.eval(_transforms=(smpl.transforms[0][0], smpl.transforms[0][0]), _points=smpl.points[0]), numpy.array([1, 1]))
+
+  def test_prepare_eval_without_points(self):
+    f = function.ones((2,), int)
+    lowered = f.prepare_eval(ndims=None, npoints=None)
+    with _builtin_warnings.catch_warnings():
+      _builtin_warnings.simplefilter('ignore', category=evaluable.ExpensiveEvaluationWarning)
+      self.assertAllEqual(lowered.eval(), numpy.array([1, 1]))
 
 class integral_compatibility(TestCase):
 
@@ -172,7 +187,7 @@ def assertArrayAlmostEqual(self, actual, desired, decimal):
 
   def test_lower_eval(self):
     args = tuple((numpy.random.randint if self.dtype == int else numpy.random.uniform)(size=shape, low=self.low, high=self.high) for shape in self.shapes)
-    actual = self.op(*args).prepare_eval(npoints=None).eval()
+    actual = self.op(*args).as_evaluable_array().eval()
     desired = self.n_op(*args)
     self.assertArrayAlmostEqual(actual, desired, decimal=15)
 
@@ -340,10 +355,10 @@ def test_runtime_check(self):
     b = function.Argument('b', (m,), dtype=int)
     (a_, b_), shape, dtype = function._broadcast(a, b)
     with self.subTest('match'):
-      self.assertEqual(shape[0].prepare_eval(npoints=None).eval(n=numpy.array(2), m=numpy.array(2)), 2)
+      self.assertEqual(shape[0].as_evaluable_array().eval(n=numpy.array(2), m=numpy.array(2)), 2)
     with self.subTest('mismatch'):
       with self.assertRaises(evaluable.EvaluationError):
-        shape[0].prepare_eval(npoints=None).eval(n=numpy.array(2), m=numpy.array(3))
+        shape[0].as_evaluable_array().eval(n=numpy.array(2), m=numpy.array(3))
 
 
 @parametrize
@@ -402,8 +417,7 @@ class elemwise(TestCase):
 
   def setUp(self):
     super().setUp()
-    self.domain, geom = mesh.rectilinear([5])
-    self.index = self.domain.f_index
+    self.index = function.Argument('index', (), dtype=int)
     self.data = tuple(map(types.frozenarray, (
       numpy.arange(1, dtype=float).reshape(1,1),
       numpy.arange(2, dtype=float).reshape(1,2),
@@ -414,20 +428,14 @@ def setUp(self):
     self.func = function.Elemwise(self.data, self.index, float)
 
   def test_evalf(self):
-    for i, trans in enumerate(self.domain.transforms):
+    for i in range(5):
       with self.subTest(i=i):
-        numpy.testing.assert_array_almost_equal(self.func.prepare_eval(ndims=self.domain.ndims).eval(_transforms=(trans,), _points=points.SimplexGaussPoints(self.domain.ndims, 1)), self.data[i][_])
+        numpy.testing.assert_array_almost_equal(self.func.as_evaluable_array().eval(index=i), self.data[i])
 
   def test_shape(self):
-    for i, trans in enumerate(self.domain.transforms):
+    for i in range(5):
       with self.subTest(i=i):
-        self.assertEqual(self.func.size.prepare_eval(ndims=self.domain.ndims, npoints=None).eval(_transforms=(trans,)), self.data[i].size)
-
-  def test_derivative(self):
-    self.assertTrue(evaluable.iszero(function.localgradient(self.func, self.domain.ndims).prepare_eval(ndims=self.domain.ndims)))
-
-  def test_shape_derivative(self):
-    self.assertEqual(function.localgradient(self.func, self.domain.ndims).shape, self.func.shape+(self.domain.ndims,))
+        self.assertEqual(self.func.size.as_evaluable_array().eval(index=i), self.data[i].size)
 
 
 class replace_arguments(TestCase):
@@ -435,38 +443,38 @@ class replace_arguments(TestCase):
   def test_array(self):
     a = function.Argument('a', (2,))
     b = function.Array.cast([1,2])
-    self.assertEqual(function.replace_arguments(a, dict(a=b)).prepare_eval(npoints=None), b.prepare_eval(npoints=None))
+    self.assertEqual(function.replace_arguments(a, dict(a=b)).as_evaluable_array(), b.as_evaluable_array())
 
   def test_argument(self):
     a = function.Argument('a', (2,))
     b = function.Argument('b', (2,))
-    self.assertEqual(function.replace_arguments(a, dict(a=b)).prepare_eval(npoints=None), b.prepare_eval(npoints=None))
+    self.assertEqual(function.replace_arguments(a, dict(a=b)).as_evaluable_array(), b.as_evaluable_array())
 
   def test_argument_array(self):
     a = function.Argument('a', (2,))
     b = function.Argument('b', (2,))
     c = function.Array.cast([1,2])
-    self.assertEqual(function.replace_arguments(function.replace_arguments(a, dict(a=b)), dict(b=c)).prepare_eval(npoints=None), c.prepare_eval(npoints=None))
+    self.assertEqual(function.replace_arguments(function.replace_arguments(a, dict(a=b)), dict(b=c)).as_evaluable_array(), c.as_evaluable_array())
 
   def test_swap(self):
     a = function.Argument('a', (2,))
     b = function.Argument('b', (2,))
-    self.assertEqual(function.replace_arguments(2*a+3*b, dict(a=b, b=a)).prepare_eval(npoints=None), (2*b+3*a).prepare_eval(npoints=None))
+    self.assertEqual(function.replace_arguments(2*a+3*b, dict(a=b, b=a)).as_evaluable_array(), (2*b+3*a).as_evaluable_array())
 
   def test_ignore_replaced(self):
     a = function.Argument('a', (2,))
     b = function.Array.cast([1,2])
     c = function.Array.cast([2,3])
-    self.assertEqual(function.replace_arguments(function.replace_arguments(a, dict(a=b)), dict(a=c)).prepare_eval(npoints=None), b.prepare_eval(npoints=None))
+    self.assertEqual(function.replace_arguments(function.replace_arguments(a, dict(a=b)), dict(a=c)).as_evaluable_array(), b.as_evaluable_array())
 
   def test_ignore_recursion(self):
     a = function.Argument('a', (2,))
-    self.assertEqual(function.replace_arguments(a, dict(a=2*a)).prepare_eval(npoints=None), (2*a).prepare_eval(npoints=None))
+    self.assertEqual(function.replace_arguments(a, dict(a=2*a)).as_evaluable_array(), (2*a).as_evaluable_array())
 
   def test_replace_derivative(self):
     a = function.Argument('a', ())
     b = function.Argument('b', ())
-    self.assertEqual(function.replace_arguments(function.derivative(a, a), dict(a=b)).prepare_eval(npoints=None).simplified, evaluable.ones(()).simplified)
+    self.assertEqual(function.replace_arguments(function.derivative(a, a), dict(a=b)).as_evaluable_array().simplified, evaluable.ones(()).simplified)
 
 
 class namespace(TestCase):
@@ -526,16 +534,21 @@ def test_invalid_default_geometry_no_variable(self):
     with self.assertRaises(ValueError):
       function.Namespace(default_geometry_name='foo_bar')
 
-  def assertEqualLowered(self, actual, desired, **lowerargs):
-    return self.assertEqual(actual.prepare_eval(**lowerargs), desired.prepare_eval(**lowerargs))
+  def assertEqualLowered(self, actual, desired, *, topo=None):
+    if topo:
+      smpl = topo.sample('gauss', 2)
+      lower = lambda f: evaluable.asarray(f @ smpl)
+    else:
+      lower = evaluable.asarray
+    return self.assertEqual(lower(actual), lower(desired))
 
   def test_default_geometry_property(self):
     ns = function.Namespace()
     ns.x = 1
-    self.assertEqualLowered(ns.default_geometry, ns.x, ndims=0)
+    self.assertEqualLowered(ns.default_geometry, ns.x)
     ns = function.Namespace(default_geometry_name='y')
     ns.y = 2
-    self.assertEqualLowered(ns.default_geometry, ns.y, ndims=0)
+    self.assertEqualLowered(ns.default_geometry, ns.y)
 
   def test_copy(self):
     ns = function.Namespace()
@@ -549,7 +562,7 @@ def test_copy_change_geom(self):
     ns1.basis = domain.basis('spline', degree=2)
     ns2 = ns1.copy_(default_geometry_name='y')
     self.assertEqual(ns2.default_geometry_name, 'y')
-    self.assertEqualLowered(ns2.eval_ni('basis_n,i'), ns2.basis.grad(ns2.y), ndims=domain.ndims)
+    self.assertEqualLowered(ns2.eval_ni('basis_n,i'), ns2.basis.grad(ns2.y), topo=domain)
 
   def test_copy_preserve_geom(self):
     ns1 = function.Namespace(default_geometry_name='y')
@@ -557,7 +570,7 @@ def test_copy_preserve_geom(self):
     ns1.basis = domain.basis('spline', degree=2)
     ns2 = ns1.copy_()
     self.assertEqual(ns2.default_geometry_name, 'y')
-    self.assertEqualLowered(ns2.eval_ni('basis_n,i'), ns2.basis.grad(ns2.y), ndims=domain.ndims)
+    self.assertEqualLowered(ns2.eval_ni('basis_n,i'), ns2.basis.grad(ns2.y), topo=domain)
 
   def test_copy_fixed_lengths(self):
     ns = function.Namespace(length_i=2)
@@ -590,12 +603,12 @@ def test_eval_fallback_length(self):
   def test_matmul_0d(self):
     ns = function.Namespace()
     ns.foo = 2
-    self.assertEqualLowered('foo' @ ns, ns.foo, npoints=None)
+    self.assertEqualLowered('foo' @ ns, ns.foo)
 
   def test_matmul_1d(self):
     ns = function.Namespace()
     ns.foo = function.zeros([2])
-    self.assertEqualLowered('foo_i' @ ns, ns.foo, npoints=None)
+    self.assertEqualLowered('foo_i' @ ns, ns.foo)
 
   def test_matmul_2d(self):
     ns = function.Namespace()
@@ -621,7 +634,7 @@ def test_replace(self):
     ns.foo = function.Argument('arg', [2,3])
     ns.bar_ij = 'sin(foo_ij) + cos(2 foo_ij)'
     ns = ns(arg=function.zeros([2,3]))
-    self.assertEqualLowered(ns.foo, function.zeros([2,3]), npoints=None)
+    self.assertEqualLowered(ns.foo, function.zeros([2,3]))
     self.assertEqual(ns.default_geometry_name, 'y')
 
   def test_pickle(self):
@@ -633,7 +646,7 @@ def test_pickle(self):
     orig.f = 'cosh(x_0)'
     pickled = pickle.loads(pickle.dumps(orig))
     for attr in ('x', 'v', 'u', 'f'):
-      self.assertEqualLowered(getattr(pickled, attr), getattr(orig, attr), ndims=domain.ndims)
+      self.assertEqualLowered(getattr(pickled, attr), getattr(orig, attr), topo=domain)
     self.assertEqual(pickled.arg_shapes['lhs'], orig.arg_shapes['lhs'])
 
   def test_pickle_default_geometry_name(self):
@@ -662,17 +675,17 @@ def test_unexpected_keyword_argument(self):
   def test_d_geom(self):
     ns = function.Namespace()
     topo, ns.x = mesh.rectilinear([1])
-    self.assertEqualLowered(ns.eval_ij('d(x_i, x_j)'), function.grad(ns.x, ns.x), ndims=topo.ndims)
+    self.assertEqualLowered(ns.eval_ij('d(x_i, x_j)'), function.grad(ns.x, ns.x), topo=topo)
 
   def test_d_arg(self):
     ns = function.Namespace()
     ns.a = '?a'
-    self.assertEqual(ns.eval_('d(2 ?a + 1, ?a)').prepare_eval(npoints=None).simplified, function.asarray(2).prepare_eval(npoints=None).simplified)
+    self.assertEqual(ns.eval_('d(2 ?a + 1, ?a)').as_evaluable_array().simplified, function.asarray(2).as_evaluable_array().simplified)
 
   def test_n(self):
     ns = function.Namespace()
     topo, ns.x = mesh.rectilinear([1])
-    self.assertEqualLowered(ns.eval_i('n(x_i)'), function.normal(ns.x), ndims=topo.ndims-1)
+    self.assertEqualLowered(ns.eval_i('n(x_i)'), function.normal(ns.x), topo=topo.boundary)
 
   def test_functions(self):
     def sqr(a):
@@ -686,30 +699,29 @@ def mul(*args):
     ns.a = numpy.array([1, 2, 3])
     ns.b = numpy.array([4, 5])
     ns.A = numpy.array([[6, 7, 8], [9, 10, 11]])
-    l = lambda f: f.prepare_eval(npoints=None).simplified
+    l = lambda f: f.as_evaluable_array().simplified
     self.assertEqual(l(ns.eval_i('sqr(a_i)')), l(sqr(ns.a)))
     self.assertEqual(l(ns.eval_ij('mul(a_i, b_j)')), l(ns.eval_ij('a_i b_j')))
     self.assertEqual(l(ns.eval_('mul(b_i, A_ij, a_j)')), l(ns.eval_('b_i A_ij a_j')))
 
   def test_builtin_functions(self):
     ns = function.Namespace()
-    domain, ns.x = mesh.rectilinear([1]*2)
     ns.a = numpy.array([1, 2, 3])
     ns.A = numpy.array([[6, 7, 8], [9, 10, 11]])
-    l = lambda f: f.prepare_eval(npoints=4, ndims=2).simplified
+    l = lambda f: f.as_evaluable_array().simplified
     self.assertEqual(l(ns.eval_('norm2(a)')), l(function.norm2(ns.a)))
     self.assertEqual(l(ns.eval_i('sum:j(A_ij)')), l(function.sum(ns.A, 1)))
 
   def test_builtin_jacobian_vector(self):
     ns = function.Namespace()
     domain, ns.x = mesh.rectilinear([1]*2)
-    l = lambda f: f.prepare_eval(npoints=4, ndims=2).simplified
+    l = lambda f: evaluable.asarray(f @ domain.sample('gauss', 2)).simplified
     self.assertEqual(l(ns.eval_('J(x)')), l(function.jacobian(ns.x)))
 
   def test_builtin_jacobian_scalar(self):
     ns = function.Namespace()
     domain, (ns.t,) = mesh.rectilinear([1])
-    l = lambda f: f.prepare_eval(npoints=4, ndims=1).simplified
+    l = lambda f: evaluable.asarray(f @ domain.sample('gauss', 2)).simplified
     self.assertEqual(l(ns.eval_('J(t)')), l(function.jacobian(ns.t[None])))
 
   def test_builtin_jacobian_matrix(self):
@@ -726,11 +738,11 @@ class eval_ast(TestCase):
 
   def setUp(self):
     super().setUp()
-    domain, x = mesh.rectilinear([2,2])
+    self.domain, x = mesh.rectilinear([2,2])
     self.ns = function.Namespace()
     self.ns.x = x
     self.ns.altgeom = function.concatenate([self.ns.x, [0]], 0)
-    self.ns.basis = domain.basis('spline', degree=2)
+    self.ns.basis = self.domain.basis('spline', degree=2)
     self.ns.a = 2
     self.ns.a2 = numpy.array([1,2])
     self.ns.a3 = numpy.array([1,2,3])
@@ -738,17 +750,25 @@ def setUp(self):
     self.ns.a32 = numpy.array([[1,2],[3,4],[5,6]])
     self.x = function.Argument('x',())
 
-  def assertEqualLowered(self, s, f, *, ndims=2):
-    self.assertEqual((s @ self.ns).prepare_eval(ndims=ndims).simplified, f.prepare_eval(ndims=ndims).simplified)
+  def assertEqualLowered(self, s, f, *, topo=None, indices=None):
+    if topo is None:
+      topo = self.domain
+    smpl = topo.sample('gauss', 2)
+    lower = lambda g: evaluable.asarray(g @ smpl).simplified
+    if indices:
+      evaluated = getattr(self.ns, 'eval_'+indices)(s)
+    else:
+      evaluated = s @ self.ns
+    self.assertEqual(lower(evaluated), lower(f))
 
   def test_group(self): self.assertEqualLowered('(a)', self.ns.a)
   def test_arg(self): self.assertEqualLowered('a2_i ?x_i', function.dot(self.ns.a2, function.Argument('x', [2]), axes=[0]))
   def test_substitute(self): self.assertEqualLowered('(?x_i^2)(x_i=a2_i)', self.ns.a2**2)
   def test_multisubstitute(self): self.assertEqualLowered('(a2_i + ?x_i + ?y_i)(x_i=?y_i, y_i=?x_i)', self.ns.a2 + function.Argument('y', [2]) + function.Argument('x', [2]))
   def test_call(self): self.assertEqualLowered('sin(a)', function.sin(self.ns.a))
-  def test_call2(self): self.assertEqual(self.ns.eval_ij('arctan2(a2_i, a3_j)').prepare_eval(ndims=2).simplified, function.arctan2(self.ns.a2[:,None], self.ns.a3[None,:]).prepare_eval(ndims=2).simplified)
+  def test_call2(self): self.assertEqualLowered('arctan2(a2_i, a3_j)', function.arctan2(self.ns.a2[:,None], self.ns.a3[None,:]), indices='ij')
   def test_eye(self): self.assertEqualLowered('δ_ij a2_i', function.dot(function.eye(2), self.ns.a2, axes=[0]))
-  def test_normal(self): self.assertEqualLowered('n_i', self.ns.x.normal(), ndims=1)
+  def test_normal(self): self.assertEqualLowered('n_i', self.ns.x.normal(), topo=self.domain.boundary)
   def test_getitem(self): self.assertEqualLowered('a2_0', self.ns.a2[0])
   def test_trace(self): self.assertEqualLowered('a22_ii', function.trace(self.ns.a22, 0, 1))
   def test_sum(self): self.assertEqualLowered('a2_i a2_i', function.sum(self.ns.a2 * self.ns.a2, axis=0))
@@ -1098,11 +1118,11 @@ def checkeval(self, ielem, points):
   def test_lower(self):
     ref = element.PointReference() if self.basis.coords.shape[0] == 0 else element.LineReference()**self.basis.coords.shape[0]
     points = ref.getpoints('bezier', 4)
-    lowered = self.basis.prepare_eval(ndims=points.ndims)
+    lowered = self.basis.lower(transform_chains=(evaluable.TransformChainFromSequence(self.checktransforms, evaluable.Argument('ielem', (), int)),), coordinates=(evaluable.Constant(points.coords),))
     with _builtin_warnings.catch_warnings():
       _builtin_warnings.simplefilter('ignore', category=evaluable.ExpensiveEvaluationWarning)
       for ielem in range(self.checknelems):
-        value = lowered.eval(_transforms=(self.checktransforms[ielem],), _points=points)
+        value = lowered.eval(ielem=ielem)
         if value.shape[0] == 1:
           value = numpy.tile(value, (points.npoints, 1))
         self.assertEqual(value.tolist(), self.checkeval(ielem, points))
diff --git a/tests/test_sample.py b/tests/test_sample.py
index 897fa2059..e73473629 100644
--- a/tests/test_sample.py
+++ b/tests/test_sample.py
@@ -46,7 +46,7 @@ def test_asfunction(self):
       self.bezier2.eval(sampled)
     self.assertAllEqual(self.gauss2.eval(sampled), values)
     arg = function.Argument('dofs', [2,3])
-    self.assertTrue(evaluable.iszero(function.derivative(sampled, arg).prepare_eval(ndims=self.domain.ndims)))
+    self.assertTrue(evaluable.iszero(evaluable.asarray(function.derivative(sampled, arg) @ self.gauss2)))
 
 class integral(TestCase):
 
diff --git a/tests/test_topology.py b/tests/test_topology.py
index 162c65bda..025d6bc6e 100644
--- a/tests/test_topology.py
+++ b/tests/test_topology.py
@@ -10,6 +10,7 @@ def assertConnectivity(self, domain, geom):
     interfaces = domain.interfaces
     bmask = numpy.zeros(len(boundary), dtype=int)
     imask = numpy.zeros(len(interfaces), dtype=int)
+    lowered_geom = geom.lower(transform_chains=(evaluable.SelectChain(0),), coordinates=(evaluable.Points(evaluable.NPoints(), boundary.ndims),))
     for ielem, ioppelems in enumerate(domain.connectivity):
       for iedge, ioppelem in enumerate(ioppelems):
         etrans, eref = domain.references[ielem].edges[iedge]
@@ -31,8 +32,8 @@ def assertConnectivity(self, domain, geom):
           imask[index] += 1
           self.assertEqual(eref, opperef)
           points = eref.getpoints('gauss', 2)
-          a0 = geom.prepare_eval(ndims=eref.ndims).eval(_transforms=[trans], _points=points)
-          a1 = geom.prepare_eval(ndims=eref.ndims).eval(_transforms=[opptrans], _points=points)
+          a0 = lowered_geom.eval(_transforms=[trans], _points=points)
+          a1 = lowered_geom.eval(_transforms=[opptrans], _points=points)
           numpy.testing.assert_array_almost_equal(a0, a1)
     self.assertTrue(numpy.equal(bmask, 1).all())
     self.assertTrue(numpy.equal(imask, 2).all())

From 520e1885a09506e32816968d09644cd81d7192fa Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Fri, 9 Apr 2021 12:23:58 +0200
Subject: [PATCH 16/26] preserve as much shape as possible in Elemwise

---
 nutils/evaluable.py     | 29 +++++++++++++++++++----------
 tests/test_evaluable.py | 39 ++++++++++++++-------------------------
 2 files changed, 33 insertions(+), 35 deletions(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 1ecbd19c9..2e17b17ec 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -2369,21 +2369,30 @@ def Elemwise(data:types.tuple[types.arraydata], index:asarray, dtype:asdtype):
   if len(unique) == 1:
     return Constant(unique[0])
   shape = [Take(s, index).simplified for s in numpy.array([d.shape for d in data]).T] # use original index to avoid potential inconsistencies with other arrays
-  ndim = len(shape)
   if len(unique) < len(data):
     index = Take(indices, index)
-  raveled = list(map(numpy.ravel, unique))
-  concat = numpy.concatenate(raveled)
-  if ndim == 0:
+  # Move all axes with constant shape to the left and ravel the remainder.
+  is_constant = numpy.array([n.isconstant for n in shape], dtype=bool)
+  nconstant = is_constant.sum()
+  reorder = numpy.argsort(~is_constant)
+  shape = [shape[i] for i in reorder]
+  var_shape = shape[nconstant:]
+  reshape = [n.__index__() for n in shape[:nconstant]] + [-1]
+  raveled = [numpy.transpose(d, reorder).reshape(reshape) for d in unique]
+  # Concatenate the raveled axis, take slices, unravel and reorder the axes to
+  # the original position.
+  concat = numpy.concatenate(raveled, axis=-1)
+  if len(var_shape) == 0:
+    assert tuple(reorder) == tuple(range(nconstant))
     return Take(concat, index)
-  cumprod = list(shape)
-  for i in reversed(range(ndim-1)):
+  cumprod = list(var_shape)
+  for i in reversed(range(len(var_shape)-1)):
     cumprod[i] *= cumprod[i+1] # work backwards so that the shape check matches in Unravel
-  offsets = numpy.cumsum([0, *map(len, raveled[:-1])])
+  offsets = _SizesToOffsets(asarray([d.shape[-1] for d in raveled]))
   elemwise = Take(concat, Range(cumprod[0]) + Take(offsets, index))
-  for i in range(ndim-1):
-    elemwise = Unravel(elemwise, shape[i], cumprod[i+1])
-  return elemwise
+  for i in range(len(var_shape)-1):
+    elemwise = Unravel(elemwise, var_shape[i], cumprod[i+1])
+  return Transpose(elemwise, tuple(numpy.argsort(reorder)))
 
 class Eig(Evaluable):
 
diff --git a/tests/test_evaluable.py b/tests/test_evaluable.py
index 99da0d4d3..cef5c9c9e 100644
--- a/tests/test_evaluable.py
+++ b/tests/test_evaluable.py
@@ -574,35 +574,24 @@ def test_deriv(self):
 
 class elemwise(TestCase):
 
-  def setUp(self):
-    super().setUp()
-    self.index = evaluable.Argument('index', (), int)
-    self.data = tuple(map(types.frozenarray, (
-      numpy.arange(1, dtype=float).reshape(1,1),
-      numpy.arange(2, dtype=float).reshape(1,2),
-      numpy.arange(3, dtype=float).reshape(3,1),
-      numpy.arange(4, dtype=float).reshape(2,2),
-      numpy.arange(6, dtype=float).reshape(3,2),
-      numpy.arange(4, dtype=float).reshape(2,2),
-      numpy.arange(3, dtype=float).reshape(3,1),
-    )))
-    self.func = evaluable.Elemwise(self.data, self.index, float)
+  def assertElemwise(self, items):
+    items = tuple(map(numpy.array, items))
+    index = evaluable.Argument('index', (), int)
+    elemwise = evaluable.Elemwise(items, index, int)
+    for i, item in enumerate(items):
+      self.assertEqual(elemwise.eval(index=i).tolist(), item.tolist())
 
-  def test_evalf(self):
-    for i in range(7):
-      with self.subTest(i=i):
-        numpy.testing.assert_array_almost_equal(self.func.eval(index=i), self.data[i])
+  def test_const_values(self):
+    self.assertElemwise((numpy.arange(2*3*4).reshape(2,3,4),)*3)
 
-  def test_shape(self):
-    for i in range(7):
-      with self.subTest(i=i):
-        self.assertEqual(self.func.size.eval(index=i), self.data[i].size)
+  def test_const_shape(self):
+    self.assertElemwise(numpy.arange(4*2*3*4).reshape(4,2,3,4))
 
-  def test_derivative(self):
-    self.assertTrue(evaluable.iszero(evaluable.localgradient(self.func, 2).simplified))
+  def test_mixed_shape(self):
+    self.assertElemwise(numpy.arange(4*i*j*3).reshape(4,i,j,3) for i, j in ((1,2),(2,4)))
 
-  def test_shape_derivative(self):
-    self.assertEqual(evaluable.localgradient(self.func, 2).shape, self.func.shape+(evaluable.Constant(2),))
+  def test_var_shape(self):
+    self.assertElemwise(numpy.arange(i*j).reshape(i,j) for i, j in ((1,2),(2,4)))
 
 class jacobian(TestCase):
 

From 9f3a1efb885c31c32760c406901efe791ee7cbc3 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Sat, 10 Apr 2021 14:42:37 +0200
Subject: [PATCH 17/26] cache function.Array.as_evaluable_array

This patch changes `as_evaluable_array` into a property and caches the result
of `function.Array.as_evaluable_array`. If a `function.Array`, e.g. a
`Sample.integral`, is passed to `solver.newton` inside a time loop, this
prevents repeated lowering of the `function.Array` to a `evaluable.Array`.
---
 nutils/evaluable.py    | 10 +++++++---
 nutils/function.py     |  7 ++++---
 nutils/sample.py       |  2 +-
 nutils/solver.py       |  4 ++--
 tests/test_function.py | 38 +++++++++++++++++++-------------------
 5 files changed, 33 insertions(+), 28 deletions(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 2e17b17ec..c159901d6 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -66,9 +66,11 @@ def simplified(value):
 asdtype = lambda arg: arg if any(arg is dtype for dtype in (bool, int, float, complex)) else {'f': float, 'i': int, 'b': bool, 'c': complex}[numpy.dtype(arg).kind]
 
 def asarray(arg):
-  if hasattr(arg, 'as_evaluable_array'):
-    return arg.as_evaluable_array()
-  elif _containsarray(arg):
+  try:
+    return arg.as_evaluable_array
+  except AttributeError:
+    pass
+  if _containsarray(arg):
     return stack(arg, axis=0)
   else:
     return Constant(arg)
@@ -824,6 +826,7 @@ def __new__(mcls, name, bases, namespace):
 class AsEvaluableArray(Protocol):
   'Protocol for conversion into an :class:`Array`.'
 
+  @property
   def as_evaluable_array(self) -> 'Array':
     'Lower this object to a :class:`nutils.evaluable.Array`.'
 
@@ -1035,6 +1038,7 @@ def _derivative(self, var, seen):
       return Zeros(self.shape + var.shape, dtype=self.dtype)
     raise NotImplementedError('derivative not defined for {}'.format(self.__class__.__name__))
 
+  @property
   def as_evaluable_array(self):
     'return self'
 
diff --git a/nutils/function.py b/nutils/function.py
index a6a4e96da..59c353fec 100644
--- a/nutils/function.py
+++ b/nutils/function.py
@@ -124,13 +124,14 @@ def __init__(self, shape: Shape, dtype: DType) -> None:
         # Try to convert `length` to an `int` by lowering to an `Evaluable` and
         # evaluating.
         try:
-          length = int(length.lower().eval())
+          length = int(length.as_evaluable_array.eval())
         except:
           pass
       shape_.append(length)
     self.shape = tuple(shape_)
     self.dtype = dtype
 
+  @util.cached_property
   def as_evaluable_array(self) -> evaluable.Array:
     return self.lower()
 
@@ -401,7 +402,7 @@ def derivative(self, __var: Union[str, 'Argument']) -> 'Array':
     'Differentiate this function to `var`.'
 
     if isinstance(__var, str):
-      for arg in self.as_evaluable_array().arguments:
+      for arg in self.as_evaluable_array.arguments:
         if isinstance(arg, evaluable.Argument) and arg._name == __var:
           if not all(n.isconstant for n in arg.shape):
             raise ValueError('arguments with variable shapes are not supported')
@@ -424,7 +425,7 @@ def contains(self, __name: str) -> bool:
   @property
   def argshapes(self) -> Mapping[str, Tuple[int, ...]]:
     shapes = {} # type: Dict[str, Tuple[int, ...]]
-    for arg in self.as_evaluable_array().arguments:
+    for arg in self.as_evaluable_array.arguments:
       if isinstance(arg, evaluable.Argument):
         if arg._name in shapes:
           if shapes.get(arg._name, arg.shape) != arg.shape:
diff --git a/nutils/sample.py b/nutils/sample.py
index 369bf51fd..968923331 100644
--- a/nutils/sample.py
+++ b/nutils/sample.py
@@ -408,7 +408,7 @@ def eval_integrals_sparse(*integrals, **arguments):
   results : :class:`tuple` of arrays and/or :class:`nutils.matrix.Matrix` objects.
   '''
 
-  integrals = tuple(integral.as_evaluable_array().assparse for integral in integrals)
+  integrals = tuple(integral.as_evaluable_array.assparse for integral in integrals)
   with evaluable.Tuple(tuple(integrals)).optimized_for_numpy.session(graphviz=graphviz) as eval:
     return eval(**arguments)
 
diff --git a/nutils/solver.py b/nutils/solver.py
index fd14d2756..3a31a3aa5 100644
--- a/nutils/solver.py
+++ b/nutils/solver.py
@@ -59,10 +59,10 @@
 argdict = types.frozendict[types.strictstr,types.arraydata]
 
 def integraltuple(arg):
-  return tuple(a.as_evaluable_array() for a in arg)
+  return tuple(a.as_evaluable_array for a in arg)
 
 def optionalintegraltuple(arg):
-  return tuple(None if a is None else a.as_evaluable_array() for a in arg)
+  return tuple(None if a is None else a.as_evaluable_array for a in arg)
 
 def arrayordict(arg):
   return types.arraydata(arg) if numeric.isarray(arg) else argdict(arg)
diff --git a/tests/test_function.py b/tests/test_function.py
index 4dc5b5179..89e1ccae5 100644
--- a/tests/test_function.py
+++ b/tests/test_function.py
@@ -29,7 +29,7 @@ def test_shape(self):
     with self.subTest('const Array'):
       self.assertEqual(l2, 3)
     with self.subTest('unknown'):
-      self.assertEqual(l3.as_evaluable_array().simplified, n.as_evaluable_array().simplified)
+      self.assertEqual(l3.as_evaluable_array.simplified, n.as_evaluable_array.simplified)
 
   def test_size_known(self):
     self.assertEqual(function.Argument('a', (2,3)).size, 6)
@@ -41,7 +41,7 @@ def test_size_unknown(self):
     n = function.Argument('n', (), dtype=int)
     size = function.Argument('a', (2,n,3)).size
     self.assertIsInstance(size, function.Array)
-    self.assertEqual(size.as_evaluable_array().simplified, (2*n*3).as_evaluable_array().simplified)
+    self.assertEqual(size.as_evaluable_array.simplified, (2*n*3).as_evaluable_array.simplified)
 
   def test_len_0d(self):
     with self.assertRaisesRegex(Exception, '^len\\(\\) of unsized object$'):
@@ -60,8 +60,8 @@ def test_iter_0d(self):
 
   def test_iter_known(self):
     a, b = function.Array.cast([1,2])
-    self.assertEqual(a.as_evaluable_array().eval(), 1)
-    self.assertEqual(b.as_evaluable_array().eval(), 2)
+    self.assertEqual(a.as_evaluable_array.eval(), 1)
+    self.assertEqual(b.as_evaluable_array.eval(), 2)
 
   def test_iter_unknown(self):
     with self.assertRaisesRegex(Exception, '^iteration over array with unknown length$'):
@@ -187,7 +187,7 @@ def assertArrayAlmostEqual(self, actual, desired, decimal):
 
   def test_lower_eval(self):
     args = tuple((numpy.random.randint if self.dtype == int else numpy.random.uniform)(size=shape, low=self.low, high=self.high) for shape in self.shapes)
-    actual = self.op(*args).as_evaluable_array().eval()
+    actual = self.op(*args).as_evaluable_array.eval()
     desired = self.n_op(*args)
     self.assertArrayAlmostEqual(actual, desired, decimal=15)
 
@@ -355,10 +355,10 @@ def test_runtime_check(self):
     b = function.Argument('b', (m,), dtype=int)
     (a_, b_), shape, dtype = function._broadcast(a, b)
     with self.subTest('match'):
-      self.assertEqual(shape[0].as_evaluable_array().eval(n=numpy.array(2), m=numpy.array(2)), 2)
+      self.assertEqual(shape[0].as_evaluable_array.eval(n=numpy.array(2), m=numpy.array(2)), 2)
     with self.subTest('mismatch'):
       with self.assertRaises(evaluable.EvaluationError):
-        shape[0].as_evaluable_array().eval(n=numpy.array(2), m=numpy.array(3))
+        shape[0].as_evaluable_array.eval(n=numpy.array(2), m=numpy.array(3))
 
 
 @parametrize
@@ -430,12 +430,12 @@ def setUp(self):
   def test_evalf(self):
     for i in range(5):
       with self.subTest(i=i):
-        numpy.testing.assert_array_almost_equal(self.func.as_evaluable_array().eval(index=i), self.data[i])
+        numpy.testing.assert_array_almost_equal(self.func.as_evaluable_array.eval(index=i), self.data[i])
 
   def test_shape(self):
     for i in range(5):
       with self.subTest(i=i):
-        self.assertEqual(self.func.size.as_evaluable_array().eval(index=i), self.data[i].size)
+        self.assertEqual(self.func.size.as_evaluable_array.eval(index=i), self.data[i].size)
 
 
 class replace_arguments(TestCase):
@@ -443,38 +443,38 @@ class replace_arguments(TestCase):
   def test_array(self):
     a = function.Argument('a', (2,))
     b = function.Array.cast([1,2])
-    self.assertEqual(function.replace_arguments(a, dict(a=b)).as_evaluable_array(), b.as_evaluable_array())
+    self.assertEqual(function.replace_arguments(a, dict(a=b)).as_evaluable_array, b.as_evaluable_array)
 
   def test_argument(self):
     a = function.Argument('a', (2,))
     b = function.Argument('b', (2,))
-    self.assertEqual(function.replace_arguments(a, dict(a=b)).as_evaluable_array(), b.as_evaluable_array())
+    self.assertEqual(function.replace_arguments(a, dict(a=b)).as_evaluable_array, b.as_evaluable_array)
 
   def test_argument_array(self):
     a = function.Argument('a', (2,))
     b = function.Argument('b', (2,))
     c = function.Array.cast([1,2])
-    self.assertEqual(function.replace_arguments(function.replace_arguments(a, dict(a=b)), dict(b=c)).as_evaluable_array(), c.as_evaluable_array())
+    self.assertEqual(function.replace_arguments(function.replace_arguments(a, dict(a=b)), dict(b=c)).as_evaluable_array, c.as_evaluable_array)
 
   def test_swap(self):
     a = function.Argument('a', (2,))
     b = function.Argument('b', (2,))
-    self.assertEqual(function.replace_arguments(2*a+3*b, dict(a=b, b=a)).as_evaluable_array(), (2*b+3*a).as_evaluable_array())
+    self.assertEqual(function.replace_arguments(2*a+3*b, dict(a=b, b=a)).as_evaluable_array, (2*b+3*a).as_evaluable_array)
 
   def test_ignore_replaced(self):
     a = function.Argument('a', (2,))
     b = function.Array.cast([1,2])
     c = function.Array.cast([2,3])
-    self.assertEqual(function.replace_arguments(function.replace_arguments(a, dict(a=b)), dict(a=c)).as_evaluable_array(), b.as_evaluable_array())
+    self.assertEqual(function.replace_arguments(function.replace_arguments(a, dict(a=b)), dict(a=c)).as_evaluable_array, b.as_evaluable_array)
 
   def test_ignore_recursion(self):
     a = function.Argument('a', (2,))
-    self.assertEqual(function.replace_arguments(a, dict(a=2*a)).as_evaluable_array(), (2*a).as_evaluable_array())
+    self.assertEqual(function.replace_arguments(a, dict(a=2*a)).as_evaluable_array, (2*a).as_evaluable_array)
 
   def test_replace_derivative(self):
     a = function.Argument('a', ())
     b = function.Argument('b', ())
-    self.assertEqual(function.replace_arguments(function.derivative(a, a), dict(a=b)).as_evaluable_array().simplified, evaluable.ones(()).simplified)
+    self.assertEqual(function.replace_arguments(function.derivative(a, a), dict(a=b)).as_evaluable_array.simplified, evaluable.ones(()).simplified)
 
 
 class namespace(TestCase):
@@ -680,7 +680,7 @@ def test_d_geom(self):
   def test_d_arg(self):
     ns = function.Namespace()
     ns.a = '?a'
-    self.assertEqual(ns.eval_('d(2 ?a + 1, ?a)').as_evaluable_array().simplified, function.asarray(2).as_evaluable_array().simplified)
+    self.assertEqual(ns.eval_('d(2 ?a + 1, ?a)').as_evaluable_array.simplified, function.asarray(2).as_evaluable_array.simplified)
 
   def test_n(self):
     ns = function.Namespace()
@@ -699,7 +699,7 @@ def mul(*args):
     ns.a = numpy.array([1, 2, 3])
     ns.b = numpy.array([4, 5])
     ns.A = numpy.array([[6, 7, 8], [9, 10, 11]])
-    l = lambda f: f.as_evaluable_array().simplified
+    l = lambda f: f.as_evaluable_array.simplified
     self.assertEqual(l(ns.eval_i('sqr(a_i)')), l(sqr(ns.a)))
     self.assertEqual(l(ns.eval_ij('mul(a_i, b_j)')), l(ns.eval_ij('a_i b_j')))
     self.assertEqual(l(ns.eval_('mul(b_i, A_ij, a_j)')), l(ns.eval_('b_i A_ij a_j')))
@@ -708,7 +708,7 @@ def test_builtin_functions(self):
     ns = function.Namespace()
     ns.a = numpy.array([1, 2, 3])
     ns.A = numpy.array([[6, 7, 8], [9, 10, 11]])
-    l = lambda f: f.as_evaluable_array().simplified
+    l = lambda f: f.as_evaluable_array.simplified
     self.assertEqual(l(ns.eval_('norm2(a)')), l(function.norm2(ns.a)))
     self.assertEqual(l(ns.eval_i('sum:j(A_ij)')), l(function.sum(ns.A, 1)))
 

From 3610ef152d32cf77c21fc1a7a8c1fc5ed24461c2 Mon Sep 17 00:00:00 2001
From: Gertjan van Zwieten <gertjanvanzwieten@evalf.nl>
Date: Wed, 24 Mar 2021 21:03:04 +0100
Subject: [PATCH 18/26] add Array._dot

---
 nutils/evaluable.py | 50 +++++++++++++++++++++++++++++++++++++++------
 1 file changed, 44 insertions(+), 6 deletions(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index c159901d6..5f9a4e23d 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -1014,6 +1014,7 @@ def _node(self, cache, subgraph, times):
   _unravel = lambda self, axis, shape: None
   _ravel = lambda self, axis: None
   _loopsum = lambda self, index, length: None
+  _dot = lambda self, other, axis: None
 
   def _uninsert(self, axis):
     assert isinstance(self._axes[axis], Inserted)
@@ -1253,6 +1254,14 @@ def _multiply(self, other):
       if isinstance(self._axes[axis], Inserted) and isinstance(other._axes[axis], Inserted):
         return insertaxis(Multiply([self._uninsert(axis), other._uninsert(axis)]), axis, self.shape[axis])
 
+  def _dot(self, other, sumaxis):
+    # Since `_dot` cannot look beyond this `InsertAxis` and `InsertAxis`
+    # instances clot together, we check here whether `sumaxis` is inserted, not
+    # whether `sumaxis` is the last axis (for which we obviously know it is
+    # inserted).
+    if isinstance(self._axes[sumaxis], Inserted):
+      return multiply(self._uninsert(sumaxis), sum(other, sumaxis))
+
   def _insertaxis(self, axis, length):
     if axis == self.ndim - 1:
       return InsertAxis(InsertAxis(self.func, length), self.length)
@@ -1377,6 +1386,12 @@ def _multiply(self, other):
     if trymultiply is not None:
       return Transpose(trymultiply, self.axes)
 
+  def _dot(self, other, axis):
+    invaxis = self.axes[axis]
+    trydot = self.func._dot(Transpose(other, self._invaxes), invaxis)
+    if trydot is not None:
+      return Transpose(trydot, [ax-(ax>invaxis) for ax in self.axes if ax != invaxis])
+
   def _add(self, other):
     if isinstance(other, Transpose) and self.axes == other.axes:
       return Transpose(Add([self.func, other.func]), self.axes)
@@ -1685,12 +1700,12 @@ def evalf(self, arr1, arr2):
 
   def _sum(self, axis):
     func1, func2 = self.funcs
-    if equalindex(self.shape[axis], 1):
-      return multiply(get(func1, axis, 0), get(func2, axis, 0))
-    if isinstance(func1._axes[axis], Inserted):
-      return multiply(func1._uninsert(axis), func2.sum(axis))
-    if isinstance(func2._axes[axis], Inserted):
-      return multiply(func1.sum(axis), func2._uninsert(axis))
+    trydot1 = func1._dot(func2, axis)
+    if trydot1 is not None:
+      return trydot1
+    trydot2 = func2._dot(func1, axis)
+    if trydot2 is not None:
+      return trydot2
 
   def _add(self, other):
     func1, func2 = self.funcs
@@ -1843,6 +1858,15 @@ def _loopsum(self, index, length):
     if any(index not in func.arguments for func in self.funcs):
       return Add([LoopSum(func, index, length) for func in self.funcs])
 
+  def _dot(self, other, axis):
+    func1, func2 = self.funcs
+    trydot1 = func1._dot(other, axis)
+    trydot2 = func2._dot(other, axis)
+    if trydot1 is not None or trydot2 is not None:
+      dot1 = sum(func1 * other, axis) if trydot1 is None else trydot1
+      dot2 = sum(func2 * other, axis) if trydot2 is None else trydot2
+      return dot1 + dot2
+
   def _desparsify(self, axis):
     assert isinstance(self._axes[axis], Sparse)
     return _gatherblocks(block for func in self.funcs for block in func._desparsify(axis))
@@ -2583,6 +2607,14 @@ def _insertaxis(self, axis, length):
   def _multiply(self, other):
     return Inflate(Multiply([self.func, Take(other, self.dofmap)]), self.dofmap, self.length)
 
+  def _dot(self, other, axis):
+    if axis == self.ndim - 1:
+      return self._multiply(other)._sum(axis)
+    else:
+      trydot = self.func._dot(Take(other, self.dofmap), axis)
+      if trydot is not None:
+        return Inflate(trydot, self.dofmap, self.length)
+
   def _add(self, other):
     if isinstance(other, Inflate) and self.dofmap == other.dofmap:
       return Inflate(Add([self.func, other.func]), self.dofmap, self.length)
@@ -3410,6 +3442,12 @@ def _add(self, other):
     if isinstance(other, LoopSum) and other.length == self.length and other.index == self.index:
       return LoopSum(self.func + other.func, self.index, self.length)
 
+  def _dot(self, other, axis):
+    if self.index not in other.arguments:
+      trydot = self.func._dot(other, axis)
+      if trydot is not None:
+        return LoopSum(trydot, self.index, self.length)
+
   @property
   def _assparse(self):
     chunks = []

From ea94fb479f24b2af0914faf409343cb49071b33e Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Thu, 8 Apr 2021 21:27:38 +0200
Subject: [PATCH 19/26] add evaluable.Array._intbounds

This patch adds the `_intbounds` property to `evaluable.Array`, which defines
arrays the lower and upper (inclusive) bounds for integer valued arrays. This
property will be used in simplifications to be introduced in follow-up commits.
---
 nutils/evaluable.py     | 148 ++++++++++++++++++++++++++++++++--
 nutils/pointsseq.py     |   2 +-
 tests/test_evaluable.py | 171 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 315 insertions(+), 6 deletions(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 5f9a4e23d..c2be81112 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -667,7 +667,7 @@ def __len__(self):
 
   @property
   def index(self):
-    return ArrayFromTuple(self, index=0, shape=(), dtype=int)
+    return ArrayFromTuple(self, index=0, shape=(), dtype=int, _lower=0, _upper=max(0, len(self._transforms)-1))
 
   @property
   def tail(self):
@@ -846,7 +846,7 @@ class Array(Evaluable, metaclass=_ArrayMeta):
   '''
 
   __slots__ = '_axes', 'dtype', '__index'
-  __cache__ = 'blocks', 'assparse', '_assparse'
+  __cache__ = 'blocks', 'assparse', '_assparse', '_intbounds'
 
   __array_priority__ = 1. # http://stackoverflow.com/questions/7042496/numpy-coercion-problem-for-left-sided-binary-operator/7057530#7057530
 
@@ -1045,6 +1045,22 @@ def as_evaluable_array(self):
 
     return self
 
+  @property
+  def _intbounds(self):
+    # inclusive lower and upper bounds
+    if self.ndim == 0 and self.dtype == int and self.isconstant:
+      value = self.__index__()
+      return value, value
+    else:
+      lower, upper = self._intbounds_impl()
+      assert isinstance(lower, int) or lower == float('-inf') or lower == float('inf')
+      assert isinstance(upper, int) or upper == float('-inf') or upper == float('inf')
+      assert lower <= upper
+      return lower, upper
+
+  def _intbounds_impl(self):
+    return float('-inf'), float('inf')
+
 class NPoints(Array):
   'The length of the points axis.'
 
@@ -1057,6 +1073,9 @@ def evalf(self, evalargs):
     points = evalargs['_points'].coords
     return types.frozenarray(points.shape[0])
 
+  def _intbounds_impl(self):
+    return 0, float('inf')
+
 class Points(Array):
 
   __slots__ = ()
@@ -1207,6 +1226,12 @@ def _determinant(self, axis1, axis2):
     value = numpy.transpose(self.value, tuple(i for i in range(self.ndim) if i != axis1 and i != axis2) + (axis1, axis2))
     return Constant(numpy.linalg.det(value))
 
+  def _intbounds_impl(self):
+    if self.dtype == int and self.value.size:
+      return int(self.value.min()), int(self.value.max())
+    else:
+      return super()._intbounds_impl()
+
 class InsertAxis(Array):
 
   __slots__ = 'func', 'length'
@@ -1307,6 +1332,9 @@ def _loopsum(self, index, length):
   def _assparse(self):
     return tuple((*(InsertAxis(idx, self.length) for idx in indices), prependaxes(Range(self.length), values.shape), InsertAxis(values, self.length)) for *indices, values in self.func._assparse)
 
+  def _intbounds_impl(self):
+    return self.func._intbounds
+
 class Transpose(Array):
 
   __slots__ = 'func', 'axes'
@@ -1473,6 +1501,9 @@ def _desparsify(self, axis):
   def _assparse(self):
     return tuple((*(indices[i] for i in self.axes), values) for *indices, values in self.func._assparse)
 
+  def _intbounds_impl(self):
+    return self.func._intbounds
+
 class Product(Array):
 
   __slots__ = 'func',
@@ -1797,6 +1828,13 @@ def _assparse(self):
     else:
       return super()._assparse
 
+  def _intbounds_impl(self):
+    func1, func2 = self.funcs
+    lower1, upper1 = func1._intbounds
+    lower2, upper2 = func2._intbounds
+    extrema = lower1 * lower2, lower1 * upper2, upper1 * lower2, upper1 * upper2
+    return min(extrema), max(extrema)
+
 class Add(Array):
 
   __slots__ = 'funcs',
@@ -1876,6 +1914,12 @@ def _assparse(self):
     func1, func2 = self.funcs
     return _gathersparsechunks(itertools.chain(func1._assparse, func2._assparse))
 
+  def _intbounds_impl(self):
+    func1, func2 = self.funcs
+    lower1, upper1 = func1._intbounds
+    lower2, upper2 = func2._intbounds
+    return lower1 + lower2, upper1 + upper2
+
 class Einsum(Array):
 
   __slots__ = 'args', 'out_idx', 'args_idx', '_einsumfmt', '_has_summed_axes'
@@ -1987,6 +2031,16 @@ def _assparse(self):
       chunks.append((*indices, values))
     return _gathersparsechunks(chunks)
 
+  def _intbounds_impl(self):
+    lower_func, upper_func = self.func._intbounds
+    lower_length, upper_length = self.func.shape[-1]._intbounds
+    if upper_length <= 0:
+      return 0, 0
+    elif lower_length <= 0:
+      return min(0, lower_func * upper_length), max(0, upper_func * upper_length)
+    else:
+      return min(lower_func * lower_length, lower_func * upper_length), max(upper_func * lower_length, upper_func * upper_length)
+
 class TakeDiag(Array):
 
   __slots__ = 'func'
@@ -2043,6 +2097,9 @@ def _assparse(self):
         chunks.append(tuple(take(arr, mask, 0) for arr in (*indices[:-1], values)))
     return _gathersparsechunks(chunks)
 
+  def _intbounds_impl(self):
+    return self.func._intbounds
+
 class Take(Array):
 
   __slots__ = 'func', 'indices'
@@ -2100,6 +2157,9 @@ def _desparsify(self, axis):
     assert axis < self.func.ndim-1 and isinstance(self._axes[axis], Sparse)
     return [(ind, Take(f, self.indices)) for ind, f in self.func._desparsify(axis)]
 
+  def _intbounds_impl(self):
+    return self.func._intbounds
+
 class Power(Array):
 
   __slots__ = 'func', 'power'
@@ -2234,6 +2294,10 @@ class Negative(Pointwise):
   __slots__ = ()
   evalf = numpy.negative
 
+  def _intbounds_impl(self):
+    lower, upper = self.args[0]._intbounds
+    return -upper, -lower
+
 class Square(Pointwise):
   __slots__ = ()
   evalf = numpy.square
@@ -2242,6 +2306,14 @@ def _sum(self, axis):
     idx = tuple(range(func.ndim))
     return Einsum((func, func), (idx, idx), idx)._sum(axis)
 
+  def _intbounds_impl(self):
+    lower, upper = self.args[0]._intbounds
+    extrema = lower**2, upper**2
+    if lower <= 0 and upper >= 0:
+      return 0, max(extrema)
+    else:
+      return min(extrema), max(extrema)
+
 class FloorDivide(Pointwise):
   __slots__ = ()
   evalf = numpy.floor_divide
@@ -2250,6 +2322,14 @@ class Absolute(Pointwise):
   __slots__ = ()
   evalf = numpy.absolute
 
+  def _intbounds_impl(self):
+    lower, upper = self.args[0]._intbounds
+    extrema = builtins.abs(lower), builtins.abs(upper)
+    if lower <= 0 and upper >= 0:
+      return 0, max(extrema)
+    else:
+      return min(extrema), max(extrema)
+
 class Cos(Pointwise):
   'Cosine, element-wise.'
   __slots__ = ()
@@ -2300,6 +2380,18 @@ class Mod(Pointwise):
   __slots__ = ()
   evalf = numpy.mod
 
+  def _intbounds_impl(self):
+    dividend, divisor = self.args
+    lower_divisor, upper_divisor = divisor._intbounds
+    if lower_divisor > 0:
+      lower_dividend, upper_dividend = dividend._intbounds
+      if 0 <= lower_dividend and upper_dividend < lower_divisor:
+        return lower_dividend, upper_dividend
+      else:
+        return 0, upper_divisor - 1
+    else:
+      return super()._intbounds_impl()
+
 class ArcTan2(Pointwise):
   __slots__ = ()
   evalf = numpy.arctan2
@@ -2335,6 +2427,12 @@ class Int(Pointwise):
   evalf = staticmethod(lambda a: a.astype(int))
   deriv = lambda a: Zeros(a.shape, int),
 
+  def _intbounds_impl(self):
+    if self.args[0].dtype == bool:
+      return 0, 1
+    else:
+      return self.args[0]._intbounds
+
 class Sign(Array):
 
   __slots__ = 'func',
@@ -2365,6 +2463,10 @@ def _unravel(self, axis, shape):
   def _derivative(self, var, seen):
     return Zeros(self.shape + var.shape, dtype=self.dtype)
 
+  def _intbounds_impl(self):
+    lower, upper = self.func._intbounds
+    return int(numpy.sign(lower)), int(numpy.sign(upper))
+
 class Sampled(Array):
   '''Basis-like identity operator.
 
@@ -2453,12 +2555,14 @@ def evalf(self, arr):
 
 class ArrayFromTuple(Array):
 
-  __slots__ = 'arrays', 'index'
+  __slots__ = 'arrays', 'index', '_lower', '_upper'
 
   @types.apply_annotations
-  def __init__(self, arrays:strictevaluable, index:types.strictint, shape:asshape, dtype:asdtype):
+  def __init__(self, arrays:strictevaluable, index:types.strictint, shape:asshape, dtype:asdtype, *, _lower=float('-inf'), _upper=float('inf')):
     self.arrays = arrays
     self.index = index
+    self._lower = _lower
+    self._upper = _upper
     super().__init__(args=[arrays], shape=shape, dtype=dtype)
 
   def evalf(self, arrays):
@@ -2474,6 +2578,9 @@ def _node(self, cache, subgraph, times):
     else:
       return super()._node(cache, subgraph, times)
 
+  def _intbounds_impl(self):
+    return self._lower, self._upper
+
 class Zeros(Array):
   'zero'
 
@@ -2547,6 +2654,9 @@ def _determinant(self, axis1, axis2):
   def _assparse(self):
     return ()
 
+  def _intbounds_impl(self):
+    return 0, 0
+
 class Inflate(Array):
 
   __slots__ = 'func', 'dofmap', 'length', 'warn'
@@ -2682,6 +2792,10 @@ def _assparse(self):
       chunks.append((*indices[:keep_dim], inflate_indices, values))
     return tuple(chunks)
 
+  def _intbounds_impl(self):
+    lower, upper = self.func._intbounds
+    return min(lower, 0), max(upper, 0)
+
 class SwapInflateTake(Evaluable):
 
   def __init__(self, inflateidx, takeidx):
@@ -2690,7 +2804,7 @@ def __init__(self, inflateidx, takeidx):
     super().__init__(args=[inflateidx, takeidx])
 
   def __iter__(self):
-    shape = ArrayFromTuple(self, index=2, shape=(), dtype=int),
+    shape = ArrayFromTuple(self, index=2, shape=(), dtype=int, _lower=0),
     return (ArrayFromTuple(self, index=index, shape=shape, dtype=int) for index in range(2))
 
   def evalf(self, inflateidx, takeidx):
@@ -3140,6 +3254,14 @@ def _take(self, index, axis):
   def evalf(self, length):
     return numpy.arange(length)
 
+  def _intbounds_impl(self):
+    lower, upper = self.length._intbounds
+    # If the upper bound of the length `upper` is zero, or worse: negative, we
+    # define the bounds of this range as `[0,0]` and hope nobody actually cares
+    # about the bounds of an empty array. Note that `numpy.arange` returns an
+    # empty array when presented with a negative length.
+    return 0, max(0, upper - 1)
+
 class InRange(Array):
 
   __slots__ = 'index', 'length'
@@ -3154,6 +3276,12 @@ def evalf(self, index, length):
     assert index.size == 0 or 0 <= index.min() and index.max() < length
     return index
 
+  def _intbounds_impl(self):
+    lower_index, upper_index = self.index._intbounds
+    lower_length, upper_length = self.length._intbounds
+    upper = min(upper_index, max(0, upper_length - 1))
+    return max(0, min(lower_index, upper)), upper
+
 class Polyval(Array):
   '''
   Computes the :math:`k`-dimensional array
@@ -3347,6 +3475,16 @@ def _simplified(self):
     if self.length.isconstant and self.index.isconstant:
       return Constant(self.eval())
 
+  def _intbounds_impl(self):
+    lower_length, upper_length = self.length._intbounds
+    lower_index, upper_index = self.index._intbounds
+    if lower_index >= 0:
+      return min(lower_index, upper_length - 1), min(upper_index, upper_length - 1)
+    elif upper_index < 0 and isinstance(lower_length, int) and lower_length == upper_length:
+      return max(lower_index + lower_length, 0), max(upper_index + lower_length, 0)
+    else:
+      return 0, upper_length - 1
+
 class LoopSum(Array):
 
   __cache__ = '_serialized'
diff --git a/nutils/pointsseq.py b/nutils/pointsseq.py
index cee6fc8a2..8cc68d10d 100644
--- a/nutils/pointsseq.py
+++ b/nutils/pointsseq.py
@@ -675,6 +675,6 @@ def weights(self) -> evaluable.Array:
 
   @property
   def npoints(self) -> evaluable.Array:
-    return evaluable.ArrayFromTuple(self, index=2, shape=(), dtype=int)
+    return evaluable.ArrayFromTuple(self, index=2, shape=(), dtype=int, _lower=0)
 
 # vim:sw=2:sts=2:et
diff --git a/tests/test_evaluable.py b/tests/test_evaluable.py
index cef5c9c9e..aad814c68 100644
--- a/tests/test_evaluable.py
+++ b/tests/test_evaluable.py
@@ -463,6 +463,177 @@ def _check(name, op, n_op, shapes, hasgrad=True, zerograd=False, ndim=2, low=-1,
 _check('polyval_2d_p2', lambda c, x: evaluable.Polyval(c*_polyval_mask(c.shape,x.shape[1]), x), _polyval_desired, [(3,3),(4,2)], ndim=2)
 _check('polyval_2d_p1_23', lambda c, x: evaluable.Polyval(c*_polyval_mask(c.shape,x.shape[1]), x), _polyval_desired, [(2,3,2,2),(4,2)], ndim=2)
 
+class intbounds(TestCase):
+
+  @staticmethod
+  def R(start, shape):
+    # A range of numbers starting at `start` with the given `shape`.
+    if isinstance(shape, int):
+      size = shape
+      shape = shape,
+    else:
+      size = util.product(shape)
+    return evaluable.Constant(numpy.arange(start, start+size).reshape(*shape))
+
+  class S(evaluable.Array):
+    # An evaluable scalar argument with given bounds.
+    def __init__(self, argname, lower, upper):
+      self._argname = argname
+      self._lower = lower
+      self._upper = upper
+      super().__init__(args=(evaluable.EVALARGS,), shape=(), dtype=int)
+    def evalf(self, evalargs):
+      value = numpy.array(evalargs[self._argname])
+      assert self._lower <= value <= self._upper
+      return numpy.array(value)
+    @property
+    def _intbounds(self):
+      return self._lower, self._upper
+
+  def assertBounds(self, func, *, tight_lower=True, tight_upper=True, **evalargs):
+    lower, upper = func._intbounds
+    value = func.eval(**evalargs)
+    (self.assertEqual if tight_lower else self.assertLessEqual)(lower, value.min())
+    (self.assertEqual if tight_upper else self.assertGreaterEqual)(upper, value.max())
+
+  def test_default(self):
+    class Test(evaluable.Array):
+      def __init__(self):
+        super().__init__(args=(evaluable.Argument('dummy', (), int),), shape=(), dtype=int)
+      def evalf(self):
+        raise NotImplementedError
+    self.assertEqual(Test()._intbounds, (float('-inf'), float('inf')))
+
+  def test_constant(self):
+    self.assertEqual(self.R(-4,[2,3,4])._intbounds, (-4, 19))
+
+  def test_constant_empty(self):
+    self.assertEqual(self.R(0,[0])._intbounds, (float('-inf'), float('inf')))
+
+  def test_insertaxis(self):
+    arg = self.R(-4,[2,3,4])
+    self.assertEqual(evaluable.InsertAxis(arg, 2)._intbounds, arg._intbounds)
+
+  def test_transpose(self):
+    arg = self.R(-4,[2,3,4])
+    self.assertEqual(evaluable.Transpose(arg, (2,0,1))._intbounds, arg._intbounds)
+
+  def test_multiply(self):
+    args = tuple(self.R(low, [high+1-low]) for low, high in ((-13, -5), (-2, 7), (3, 11)))
+    for arg1 in args:
+      for arg2 in args:
+        self.assertBounds(evaluable.Multiply((evaluable.insertaxis(arg1, 1, arg2.shape[0]), evaluable.insertaxis(arg2, 0, arg1.shape[0]))))
+
+  def test_add(self):
+    self.assertBounds(evaluable.Add((evaluable.insertaxis(self.R(-5,[8]), 1, 5), evaluable.insertaxis(self.R(2,[5]), 0, 8))))
+
+  def test_sum_zero_axis(self):
+    self.assertEqual(evaluable.Sum(self.R(0,[0]))._intbounds, (0, 0))
+
+  def test_sum_variable_axis_including_zero(self):
+    self.assertEqual(evaluable.Sum(evaluable.Argument('test', (self.S('n', 0, 4),), int))._intbounds, (float('-inf'), float('inf')))
+
+  def test_sum_zero_size(self):
+    self.assertEqual(evaluable.Sum(self.R(0,[2,3,0]))._intbounds, (0, 0))
+
+  def test_sum_nonzero(self):
+    self.assertBounds(evaluable.Sum(self.R(-3,[9,1])))
+
+  def test_sum_unknown(self):
+    func = lambda l, h: evaluable.Sum(evaluable.InsertAxis(self.R(l,[h+1-l]), self.S('n',2,5)))
+    self.assertBounds(func(-3, 5), n=5)
+    self.assertBounds(func(-3, 5), n=5, tight_lower=False, tight_upper=False)
+    self.assertBounds(func(3, 5), n=5, tight_lower=False)
+    self.assertBounds(func(3, 5), n=2, tight_upper=False)
+    self.assertBounds(func(-3, -2), n=5, tight_upper=False)
+    self.assertBounds(func(-3, -2), n=2, tight_lower=False)
+
+  def test_takediag(self):
+    arg = self.R(-4,[2,3,3])
+    self.assertEqual(evaluable.TakeDiag(arg)._intbounds, arg._intbounds)
+
+  def test_take(self):
+    arg = self.R(-4,[2,3,4])
+    idx = self.R(0,[1])
+    self.assertEqual(evaluable.Take(arg, idx)._intbounds, arg._intbounds)
+
+  def test_negative(self):
+    self.assertBounds(evaluable.Negative(self.R(-4,[2,3,4])))
+
+  def test_square_negative(self):
+    self.assertBounds(evaluable.Square(self.R(-4,[4])))
+
+  def test_square_positive(self):
+    self.assertBounds(evaluable.Square(self.R(1,[4])))
+
+  def test_square_full(self):
+    self.assertBounds(evaluable.Square(self.R(-3,[7])))
+
+  def test_absolute_negative(self):
+    self.assertBounds(evaluable.Absolute(self.R(-4,[3])))
+
+  def test_absolute_positive(self):
+    self.assertBounds(evaluable.Absolute(self.R(1,[3])))
+
+  def test_absolute_full(self):
+    self.assertBounds(evaluable.Absolute(self.R(-3,[7])))
+
+  def test_mod_nowrap(self):
+    self.assertBounds(evaluable.Mod(evaluable.insertaxis(self.R(1,[4]), 1, 3), evaluable.insertaxis(self.R(5,[3]), 0, 4)))
+
+  def test_mod_wrap_negative(self):
+    self.assertBounds(evaluable.Mod(evaluable.insertaxis(self.R(-3,[7]), 1, 3), evaluable.insertaxis(self.R(5,[3]), 0, 7)))
+
+  def test_mod_wrap_positive(self):
+    self.assertBounds(evaluable.Mod(evaluable.insertaxis(self.R(3,[7]), 1, 3), evaluable.insertaxis(self.R(5,[3]), 0, 7)))
+
+  def test_mod_negative_divisor(self):
+    self.assertEqual(evaluable.Mod(evaluable.Argument('d', (2,), int), self.R(-3,[2]))._intbounds, (float('-inf'), float('inf')))
+
+  def test_sign(self):
+    for i in range(-2, 3):
+      for j in range(i, 3):
+        self.assertBounds(evaluable.Sign(self.R(i,[j-i+1])))
+
+  def test_zeros(self):
+    self.assertEqual(evaluable.Zeros((2,3), int)._intbounds, (0, 0))
+
+  def test_range(self):
+    self.assertEqual(evaluable.Range(self.S('n', 0, 0))._intbounds, (0, 0))
+    self.assertBounds(evaluable.Range(self.S('n', 1, 3)), n=3)
+
+  def test_inrange_loose(self):
+    self.assertEqual(evaluable.InRange(self.S('n', 3, 5), evaluable.Constant(6))._intbounds, (3, 5))
+
+  def test_inrange_strict(self):
+    self.assertEqual(evaluable.InRange(self.S('n', float('-inf'), float('inf')), self.S('m', 2, 4))._intbounds, (0, 3))
+
+  def test_inrange_empty(self):
+    self.assertEqual(evaluable.InRange(self.S('n', float('-inf'), float('inf')), evaluable.Constant(0))._intbounds, (0, 0))
+
+  def test_npoints(self):
+    self.assertEqual(evaluable.NPoints()._intbounds, (0, float('inf')))
+
+  def test_int_bool(self):
+    self.assertEqual(evaluable.Int(evaluable.Constant(numpy.array([False, True], dtype=bool)))._intbounds, (0, 1))
+
+  def test_int_int(self):
+    self.assertEqual(evaluable.Int(self.S('n', 3, 5))._intbounds, (3, 5))
+
+  def test_array_from_tuple(self):
+    self.assertEqual(evaluable.ArrayFromTuple(evaluable.Tuple((evaluable.Argument('n', (3,), int),)), 0, (3,), int, _lower=-2, _upper=3)._intbounds, (-2, 3))
+
+  def test_inflate(self):
+    self.assertEqual(evaluable.Inflate(self.R(4, (2,3)), evaluable.Constant(numpy.arange(6).reshape(2,3)), 7)._intbounds, (0, 9))
+
+  def test_normdim_positive(self):
+    self.assertEqual(evaluable.NormDim(self.S('l', 2, 4), self.S('i', 1, 3))._intbounds, (1, 3))
+
+  def test_normdim_negative(self):
+    self.assertEqual(evaluable.NormDim(self.S('l', 4, 4), self.S('i', -3, -1))._intbounds, (1, 3))
+
+  def test_normdim_mixed(self):
+    self.assertEqual(evaluable.NormDim(self.S('l', 4, 5), self.S('i', -3, 2))._intbounds, (0, 4))
 
 class blocks(TestCase):
 

From cf14be5e90fd709c0c7337f65c9532e1c3025cb2 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Thu, 8 Apr 2021 21:28:01 +0200
Subject: [PATCH 20/26] simplify InRange based on intbounds

This patch simplifies `InRange` to its argument if the bounds of the argument
imply that `InRange` is a noop.
---
 nutils/evaluable.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index c2be81112..fe1841e4c 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -3276,6 +3276,12 @@ def evalf(self, index, length):
     assert index.size == 0 or 0 <= index.min() and index.max() < length
     return index
 
+  def _simplified(self):
+    lower_length, upper_length = self.length._intbounds
+    lower_index, upper_index = self.index._intbounds
+    if 0 <= lower_index <= upper_index < lower_length:
+      return self.index
+
   def _intbounds_impl(self):
     lower_index, upper_index = self.index._intbounds
     lower_length, upper_length = self.length._intbounds

From ea1d6f06130e9fa6cde3f5ef0efe80e9711101ee Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Fri, 23 Apr 2021 14:52:44 +0200
Subject: [PATCH 21/26] simplify NormDim based on intbounds

This patch simplifies `NormDim` to its argument if the bounds of the argument
imply that `NormDim` is a noop.
---
 nutils/evaluable.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index fe1841e4c..016595f99 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -3478,6 +3478,12 @@ def evalf(self, length, index):
     return result
 
   def _simplified(self):
+    lower_length, upper_length = self.length._intbounds
+    lower_index, upper_index = self.index._intbounds
+    if 0 <= lower_index and upper_index < lower_length:
+      return self.index
+    if isinstance(lower_length, int) and lower_length == upper_length and -lower_length <= lower_index and upper_index < 0:
+      return self.index + lower_length
     if self.length.isconstant and self.index.isconstant:
       return Constant(self.eval())
 

From 34c16f84357b3b09c55c42cea69e02af4d448e69 Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Fri, 23 Apr 2021 14:54:29 +0200
Subject: [PATCH 22/26] assert bounds in NormDim

---
 nutils/evaluable.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 016595f99..c89259d96 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -3464,6 +3464,13 @@ class NormDim(Array):
   def __init__(self, length: asarray, index: asarray):
     assert length.dtype == int
     assert index.dtype == int
+    assert equalshape(length.shape, index.shape)
+    # The following corner cases makes the assertion fail, hence we can only
+    # assert the bounds if the arrays are guaranteed to be unempty:
+    #
+    #     Take(func, NormDim(func.shape[-1], Range(0) + func.shape[-1]))
+    if all(n._intbounds[0] > 0 for n in index.shape):
+      assert -length._intbounds[1] <= index._intbounds[0] and index._intbounds[1] <= length._intbounds[1] - 1
     self.length = length
     self.index = index
     super().__init__(args=[length, index], shape=index.shape, dtype=index.dtype)

From 57ea527b176aec56840607a34dda6aa3a2dca48d Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Fri, 9 Apr 2021 10:45:26 +0200
Subject: [PATCH 23/26] cast boolean to int in evaluable.asindex

The `asindex` function casts its input to an array and asserts that the
dimension is zero and the dtype as `bool` or `int`. Callers of this function
expect the return value to have dtype `int` unconditionally. This patch casts
the value to an `int` if necessary.
---
 nutils/evaluable.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index c89259d96..a94b7411e 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -81,6 +81,8 @@ def asindex(arg):
   arg = asarray(arg)
   if arg.ndim or arg.dtype not in (int, bool): # NOTE: bool to be removed after introduction of Cast
     raise ValueError('argument is not an index: {}'.format(arg))
+  if arg.dtype == bool:
+    arg = Int(arg)
   return arg
 
 @types.apply_annotations

From c368f7069c6dc8bfb2e3ea59db1ba9cd60190bcf Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Fri, 23 Apr 2021 14:39:50 +0200
Subject: [PATCH 24/26] disallow Sum of boolean array

This patch disallows a boolean array as input. The main motivation is that a
boolean array does not have `_intbounds`, and this causes the `_intbounds` of
`Sum` to be infinite with the current implementation. The cast function `Int`
already defines `_intbounds` for boolean input, hence `Sum(Int(mask))` has
proper `_intbounds`.
---
 nutils/evaluable.py     | 6 ++++--
 tests/test_evaluable.py | 3 ++-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index a94b7411e..0774be461 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -1989,9 +1989,11 @@ class Sum(Array):
   def __init__(self, func:asarray):
     if func.ndim == 0:
       raise Exception('cannot sum a scalar function')
+    if func.dtype == bool:
+      raise Exception('cannot sum a boolean function')
     self.func = func
     shape = func._axes[:-1]
-    super().__init__(args=[func], shape=shape, dtype=int if func.dtype == bool else func.dtype)
+    super().__init__(args=[func], shape=shape, dtype=func.dtype)
 
   def _simplified(self):
     if equalindex(self.func.shape[-1], 1):
@@ -3000,7 +3002,7 @@ class Find(Array):
   def __init__(self, where:asarray):
     assert isarray(where) and where.ndim == 1 and where.dtype == bool
     self.where = where
-    super().__init__(args=[where], shape=[where.sum()], dtype=int)
+    super().__init__(args=[where], shape=[Int(where).sum()], dtype=int)
 
   def evalf(self, where):
     return where.nonzero()[0]
diff --git a/tests/test_evaluable.py b/tests/test_evaluable.py
index aad814c68..d8022869a 100644
--- a/tests/test_evaluable.py
+++ b/tests/test_evaluable.py
@@ -214,10 +214,11 @@ def test_getslice(self):
         actual=self.op_args[s])
 
   def test_sumaxis(self):
+    op_args = evaluable.Int(self.op_args) if self.op_args.dtype == bool else self.op_args
     for idim in range(self.op_args.ndim):
       self.assertFunctionAlmostEqual(decimal=14,
         desired=self.n_op_argsfun.sum(idim),
-        actual=self.op_args.sum(idim))
+        actual=op_args.sum(idim))
 
   def test_add(self):
     self.assertFunctionAlmostEqual(decimal=15,

From ef03d34265f1b99d296693db205e2e55553d9b2b Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Thu, 22 Apr 2021 17:00:38 +0200
Subject: [PATCH 25/26] require nonnegative argument in evaluable.asindex

The result of `evaluable.asindex` should be nonnegative. This patch asserts
that its argument satisfies this condition. In addition the shape of every
`evaluable.Array` is now nonnegative as well.
---
 nutils/evaluable.py    | 11 +++++------
 tests/test_function.py |  2 +-
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 0774be461..162e80275 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -83,6 +83,8 @@ def asindex(arg):
     raise ValueError('argument is not an index: {}'.format(arg))
   if arg.dtype == bool:
     arg = Int(arg)
+  elif arg._intbounds[0] < 0:
+    raise ValueError('index must be non-negative')
   return arg
 
 @types.apply_annotations
@@ -2038,9 +2040,9 @@ def _assparse(self):
   def _intbounds_impl(self):
     lower_func, upper_func = self.func._intbounds
     lower_length, upper_length = self.func.shape[-1]._intbounds
-    if upper_length <= 0:
+    if upper_length == 0:
       return 0, 0
-    elif lower_length <= 0:
+    elif lower_length == 0:
       return min(0, lower_func * upper_length), max(0, upper_func * upper_length)
     else:
       return min(lower_func * lower_length, lower_func * upper_length), max(upper_func * lower_length, upper_func * upper_length)
@@ -3260,10 +3262,7 @@ def evalf(self, length):
 
   def _intbounds_impl(self):
     lower, upper = self.length._intbounds
-    # If the upper bound of the length `upper` is zero, or worse: negative, we
-    # define the bounds of this range as `[0,0]` and hope nobody actually cares
-    # about the bounds of an empty array. Note that `numpy.arange` returns an
-    # empty array when presented with a negative length.
+    assert lower >= 0
     return 0, max(0, upper - 1)
 
 class InRange(Array):
diff --git a/tests/test_function.py b/tests/test_function.py
index 89e1ccae5..3485c9975 100644
--- a/tests/test_function.py
+++ b/tests/test_function.py
@@ -417,7 +417,7 @@ class elemwise(TestCase):
 
   def setUp(self):
     super().setUp()
-    self.index = function.Argument('index', (), dtype=int)
+    self.index = function._Wrapper(lambda: evaluable.InRange(evaluable.Argument('index', (), int), 5), shape=(), dtype=int)
     self.data = tuple(map(types.frozenarray, (
       numpy.arange(1, dtype=float).reshape(1,1),
       numpy.arange(2, dtype=float).reshape(1,2),

From b5a57f66768f679e74be1be57cbb8a576ab82ada Mon Sep 17 00:00:00 2001
From: Joost van Zwieten <joostvanzwieten@evalf.com>
Date: Wed, 7 Apr 2021 14:03:57 +0200
Subject: [PATCH 26/26] add loop_index for loop_sum, loop_concatenate

The `evaluable.InRange` evaluable checks the bounds of its argument at runtime.
If `InRange` is applied to the index of a `LoopSum` or `LoopConcatenate`, we
could simplify the `InRange` away, provided that we can get hold of the length
of the loop. This patch introduces `evaluable.loop_index` with an explicit
reference to the length of the loop and mandates that the loop index of
`evaluable.LoopSum` and `evaluable.LoopConcatenate` is of this class.
---
 nutils/evaluable.py     | 170 ++++++++++++++++++++++------------------
 nutils/sample.py        |  10 +--
 tests/test_evaluable.py |  89 ++++++++++-----------
 3 files changed, 139 insertions(+), 130 deletions(-)

diff --git a/nutils/evaluable.py b/nutils/evaluable.py
index 162e80275..b2fb8d5b8 100644
--- a/nutils/evaluable.py
+++ b/nutils/evaluable.py
@@ -453,14 +453,14 @@ def _combine_loop_concatenates(self, outer_exclude):
       # Collect all top-level `LoopConcatenate` instances in `combine` and all
       # their dependent `LoopConcatenate` instances in `exclude`.
       for lc in self._loop_concatenate_deps:
-        lcs = combine.setdefault((lc.index, lc.length), [])
+        lcs = combine.setdefault(lc.index, [])
         if lc not in lcs:
           lcs.append(lc)
           exclude.update(set(lc._loop_concatenate_deps) - {lc})
       # Combine top-level `LoopConcatenate` instances excluding those in
       # `exclude`.
       replacements = {}
-      for (index, length), lcs in combine.items():
+      for index, lcs in combine.items():
         lcs = [lc for lc in lcs if lc not in exclude]
         if not lcs:
           continue
@@ -477,7 +477,7 @@ def _combine_loop_concatenates(self, outer_exclude):
         # (the remainder of `exclude`). The latter consists of loops that are
         # invariant w.r.t. the current loop `index`.
         data = data._combine_loop_concatenates(exclude)
-        combined = LoopConcatenateCombined(data, index, length)
+        combined = LoopConcatenateCombined(data, index._name, index.length)
         for i, lc in enumerate(lcs):
           replacements[lc] = ArrayFromTuple(combined, i, lc.shape, lc.dtype)
       if replacements:
@@ -1017,7 +1017,7 @@ def _node(self, cache, subgraph, times):
   _inflate = lambda self, dofmap, length, axis: None
   _unravel = lambda self, axis, shape: None
   _ravel = lambda self, axis: None
-  _loopsum = lambda self, index, length: None
+  _loopsum = lambda self, loop_index: None # NOTE: type of `loop_index` is `_LoopIndex`
   _dot = lambda self, other, axis: None
 
   def _uninsert(self, axis):
@@ -1329,8 +1329,8 @@ def _inverse(self, axis1, axis2):
     if axis1 < self.ndim-1 and axis2 < self.ndim-1:
       return InsertAxis(inverse(self.func, (axis1, axis2)), self.length)
 
-  def _loopsum(self, index, length):
-    return InsertAxis(LoopSum(self.func, index, length), self.length)
+  def _loopsum(self, index):
+    return InsertAxis(loop_sum(self.func, index), self.length)
 
   @property
   def _assparse(self):
@@ -1494,8 +1494,8 @@ def _diagonalize(self, axis):
   def _insertaxis(self, axis, length):
     return Transpose(InsertAxis(self.func, length), self.axes[:axis] + (self.ndim,) + self.axes[axis:])
 
-  def _loopsum(self, index, length):
-    return Transpose(LoopSum(self.func, index, length), self.axes)
+  def _loopsum(self, index):
+    return Transpose(loop_sum(self.func, index), self.axes)
 
   def _desparsify(self, axis):
     assert isinstance(self._axes[axis], Sparse)
@@ -1802,12 +1802,12 @@ def _inverse(self, axis1, axis2):
     if all(isinstance(func2._axes[axis], Inserted) for axis in (axis1, axis2)):
       return divide(inverse(func1, (axis1, axis2)), func2)
 
-  def _loopsum(self, index, length):
+  def _loopsum(self, index):
     func1, func2 = self.funcs
     if func1 != index and index not in func1.dependencies:
-      return Multiply([func1, LoopSum(func2, index, length)])
+      return Multiply([func1, loop_sum(func2, index)])
     if func2 != index and index not in func2.dependencies:
-      return Multiply([LoopSum(func1, index, length), func2])
+      return Multiply([loop_sum(func1, index), func2])
 
   def _desparsify(self, axis):
     assert isinstance(self._axes[axis], Sparse)
@@ -1896,9 +1896,9 @@ def _add(self, other):
   def _unravel(self, axis, shape):
     return Add([unravel(func, axis, shape) for func in self.funcs])
 
-  def _loopsum(self, index, length):
+  def _loopsum(self, index):
     if any(index not in func.arguments for func in self.funcs):
-      return Add([LoopSum(func, index, length) for func in self.funcs])
+      return Add([loop_sum(func, index) for func in self.funcs])
 
   def _dot(self, other, axis):
     func1, func2 = self.funcs
@@ -2922,8 +2922,8 @@ def _product(self):
     if numeric.isint(self.shape[-1]) and self.shape[-1] > 1:
       return Zeros(self.shape[:-1], dtype=self.dtype)
 
-  def _loopsum(self, index, length):
-    return Diagonalize(LoopSum(self.func, index, length))
+  def _loopsum(self, index):
+    return Diagonalize(loop_sum(self.func, index))
 
   def _desparsify(self, axis):
     assert isinstance(self._axes[axis], Sparse)
@@ -3180,8 +3180,8 @@ def _sign(self):
   def _product(self):
     return Product(Product(self.func))
 
-  def _loopsum(self, index, length):
-    return Ravel(LoopSum(self.func, index, length))
+  def _loopsum(self, index):
+    return Ravel(loop_sum(self.func, index))
 
   def _desparsify(self, axis):
     assert isinstance(self._axes[axis], Sparse)
@@ -3507,6 +3507,32 @@ def _intbounds_impl(self):
     else:
       return 0, upper_length - 1
 
+class _LoopIndex(Argument):
+
+  __slots__ = 'length'
+
+  @types.apply_annotations
+  def __init__(self, name: types.strictstr, length: asindex):
+    self.length = length
+    super().__init__(name, (), int)
+
+  def __str__(self):
+    try:
+      length = self.length.__index__()
+    except EvaluationError:
+      length = '?'
+    return 'LoopIndex({}, length={})'.format(self._name, length)
+
+  def _node(self, cache, subgraph, times):
+    if self in cache:
+      return cache[self]
+    cache[self] = node = RegularNode('LoopIndex', (), dict(length=self.length._node(cache, subgraph, times)), (type(self).__name__, _Stats()), subgraph)
+    return node
+
+  def _intbounds_impl(self):
+    lower_length, upper_length = self.length._intbounds
+    return 0, max(0, upper_length - 1)
+
 class LoopSum(Array):
 
   __cache__ = '_serialized'
@@ -3521,18 +3547,13 @@ def prepare_funcdata(arg):
     return (arg, *arg.shape)
 
   @types.apply_annotations
-  def __init__(self, funcdata:prepare_funcdata, index:types.strict[Argument], length:asindex):
+  def __init__(self, funcdata:prepare_funcdata, index_name:types.strictstr, length:asindex):
     shape = Tuple(funcdata[1:])
-    if index.dtype != int or index.ndim != 0:
-      raise ValueError('expected an index with dtype int and dimension zero but got {}'.format(index))
-    if index in shape.arguments:
+    self.index = loop_index(index_name, length)
+    if self.index in shape.arguments:
       raise ValueError('the shape of the function must not depend on the index')
-    if index in length.arguments:
-      raise ValueError('the length of the loop must not depend on the index')
     self.func = funcdata[0]
-    self.index = index
-    self.length = length
-    self._invariants, self._dependencies = _dependencies_sans_invariants(self.func, index)
+    self._invariants, self._dependencies = _dependencies_sans_invariants(self.func, self.index)
     axes = tuple(Axis(axis.length) if isinstance(axis, Sparse) else axis for axis in self.func._axes)
     super().__init__(args=(shape, length, *self._invariants), shape=axes, dtype=self.func.dtype)
 
@@ -3563,7 +3584,7 @@ def evalf_withtimes(self, times, shape, length, *args):
     return result
 
   def _derivative(self, var, seen):
-    return LoopSum(derivative(self.func, var, seen), self.index, self.length)
+    return loop_sum(derivative(self.func, var, seen), self.index)
 
   def _node(self, cache, subgraph, times):
     if self in cache:
@@ -3572,7 +3593,6 @@ def _node(self, cache, subgraph, times):
     for arg in self._Evaluable__args:
       subcache[arg] = arg._node(cache, subgraph, times)
     loopgraph = Subgraph('Loop', subgraph)
-    subcache[self.index] = RegularNode('LoopIndex', (), dict(length=self.length._node(cache, subgraph, times)), (type(self).__name__, _Stats()), loopgraph)
     subtimes = times.get(self, collections.defaultdict(_Stats))
     sum_kwargs = {'shape[{}]'.format(i): n._node(cache, subgraph, times) for i, n in enumerate(self.shape)}
     sum_kwargs['func'] = self.func._node(subcache, loopgraph, subtimes)
@@ -3583,37 +3603,37 @@ def _simplified(self):
     if iszero(self.func):
       return zeros_like(self)
     elif self.index not in self.func.arguments:
-      return self.func * self.length
-    return self.func._loopsum(self.index, self.length)
+      return self.func * self.index.length
+    return self.func._loopsum(self.index)
 
   def _takediag(self, axis1, axis2):
-    return LoopSum(_takediag(self.func, axis1, axis2), self.index, self.length)
+    return loop_sum(_takediag(self.func, axis1, axis2), self.index)
 
   def _take(self, index, axis):
-    return LoopSum(_take(self.func, index, axis), self.index, self.length)
+    return loop_sum(_take(self.func, index, axis), self.index)
 
   def _unravel(self, axis, shape):
-    return LoopSum(unravel(self.func, axis, shape), self.index, self.length)
+    return loop_sum(unravel(self.func, axis, shape), self.index)
 
   def _sum(self, axis):
-    return LoopSum(sum(self.func, axis), self.index, self.length)
+    return loop_sum(sum(self.func, axis), self.index)
 
   def _add(self, other):
-    if isinstance(other, LoopSum) and other.length == self.length and other.index == self.index:
-      return LoopSum(self.func + other.func, self.index, self.length)
+    if isinstance(other, LoopSum) and other.index == self.index:
+      return loop_sum(self.func + other.func, self.index)
 
   def _dot(self, other, axis):
     if self.index not in other.arguments:
       trydot = self.func._dot(other, axis)
       if trydot is not None:
-        return LoopSum(trydot, self.index, self.length)
+        return loop_sum(trydot, self.index)
 
   @property
   def _assparse(self):
     chunks = []
     for *elem_indices, elem_values in self.func._assparse:
       if self.ndim == 0:
-        values = loop_concatenate(InsertAxis(elem_values, 1), self.index, self.length)
+        values = loop_concatenate(InsertAxis(elem_values, 1), self.index)
         while values.ndim:
           values = Sum(values)
         chunks.append((values,))
@@ -3627,7 +3647,7 @@ def _assparse(self):
           for i in variable[:-1]:
             *elem_indices, elem_values = map(Ravel, (*elem_indices, elem_values))
           assert all(n.isconstant for n in elem_values.shape[:-1])
-        chunks.append(tuple(loop_concatenate(arr, self.index, self.length) for arr in (*elem_indices, elem_values)))
+        chunks.append(tuple(loop_concatenate(arr, self.index) for arr in (*elem_indices, elem_values)))
     return tuple(chunks)
 
 class _SizesToOffsets(Array):
@@ -3648,18 +3668,15 @@ def _simplified(self):
 class LoopConcatenate(Array):
 
   @types.apply_annotations
-  def __init__(self, funcdata:asarrays, index:types.strict[Argument], length:asindex):
+  def __init__(self, funcdata:asarrays, index_name:types.strictstr, length:asindex):
     self.funcdata = funcdata
     self.func, self.start, stop, *shape = funcdata
+    self.index = loop_index(index_name, length)
     if not self.func.ndim:
       raise ValueError('expected an array with at least one axis')
-    if any(index in n.arguments for n in shape):
+    if any(self.index in n.arguments for n in shape):
       raise ValueError('the shape of the function must not depend on the index')
-    if index in length.arguments:
-      raise ValueError('the length of the loop must not depend on the index')
-    self.index = index
-    self.length = length
-    self._lcc = LoopConcatenateCombined((self.funcdata,), index, length)
+    self._lcc = LoopConcatenateCombined((self.funcdata,), index_name, length)
     axes = (*(Axis(axis.length) if isinstance(axis, Sparse) else axis for axis in self.func._axes[:-1]), Axis(shape[-1]))
     super().__init__(args=[self._lcc], shape=axes, dtype=self.func.dtype)
 
@@ -3671,7 +3688,7 @@ def evalf_withtimes(self, times, arg):
       return arg[0]
 
   def _derivative(self, var, seen):
-    return Transpose.from_end(loop_concatenate(Transpose.to_end(derivative(self.func, var, seen), self.ndim-1), self.index, self.length), self.ndim-1)
+    return Transpose.from_end(loop_concatenate(Transpose.to_end(derivative(self.func, var, seen), self.ndim-1), self.index), self.ndim-1)
 
   def _node(self, cache, subgraph, times):
     if self in cache:
@@ -3684,31 +3701,31 @@ def _simplified(self):
     if iszero(self.func):
       return zeros_like(self)
     elif self.index not in self.func.arguments:
-      return Ravel(Transpose.from_end(InsertAxis(self.func, self.length), -2))
+      return Ravel(Transpose.from_end(InsertAxis(self.func, self.index.length), -2))
     for iaxis, axis in enumerate(self.func._axes[:-1]):
       if isinstance(axis, Inserted) and self.index not in axis.length.arguments:
-        return insertaxis(loop_concatenate(self.func._uninsert(iaxis), self.index, self.length), iaxis, axis.length)
+        return insertaxis(loop_concatenate(self.func._uninsert(iaxis), self.index), iaxis, axis.length)
     if isinstance(self.func._axes[-1], Inserted) and self.index not in self.func.shape[-1].arguments and not equalindex(self.func.shape[-1], 1):
-      return Ravel(InsertAxis(loop_concatenate(InsertAxis(self.func._uninsert(self.func.ndim-1), 1), self.index, self.length), self.func.shape[-1]))
+      return Ravel(InsertAxis(loop_concatenate(InsertAxis(self.func._uninsert(self.func.ndim-1), 1), self.index), self.func.shape[-1]))
 
   def _takediag(self, axis1, axis2):
     if axis1 < self.ndim-1 and axis2 < self.ndim-1:
-      return Transpose.from_end(loop_concatenate(Transpose.to_end(_takediag(self.func, axis1, axis2), -2), self.index, self.length), -2)
+      return Transpose.from_end(loop_concatenate(Transpose.to_end(_takediag(self.func, axis1, axis2), -2), self.index), -2)
 
   def _take(self, index, axis):
     if axis < self.ndim-1:
-      return loop_concatenate(_take(self.func, index, axis), self.index, self.length)
+      return loop_concatenate(_take(self.func, index, axis), self.index)
 
   def _unravel(self, axis, shape):
     if axis < self.ndim-1:
-      return loop_concatenate(unravel(self.func, axis, shape), self.index, self.length)
+      return loop_concatenate(unravel(self.func, axis, shape), self.index)
 
   @property
   def _assparse(self):
     chunks = []
     for *indices, last_index, values in self.func._assparse:
       last_index = last_index + prependaxes(self.start, last_index.shape)
-      chunks.append(tuple(loop_concatenate(_flat(arr), self.index, self.length) for arr in (*indices, last_index, values)))
+      chunks.append(tuple(loop_concatenate(_flat(arr), self.index) for arr in (*indices, last_index, values)))
     return tuple(chunks)
 
   @property
@@ -3720,22 +3737,17 @@ class LoopConcatenateCombined(Evaluable):
   __cache__ = '_serialized'
 
   @types.apply_annotations
-  def __init__(self, funcdatas:types.tuple[asarrays], index:types.strict[Argument], length:asindex):
+  def __init__(self, funcdatas:types.tuple[asarrays], index_name:types.strictstr, length:asindex):
     self._funcdatas = funcdatas
-    if index.dtype != int or index.ndim != 0:
-      raise ValueError('expected an index with dtype int and dimension zero but got {}'.format(index))
     self._funcs = tuple(func for func, start, stop, *shape in funcdatas)
+    self._index = loop_index(index_name, length)
     if any(not func.ndim for func in self._funcs):
       raise ValueError('expected an array with at least one axis')
     shapes = [Tuple(shape) for func, start, stop, *shape in funcdatas]
-    if any(index in shape.arguments for shape in shapes):
+    if any(self._index in shape.arguments for shape in shapes):
       raise ValueError('the shape of the function must not depend on the index')
-    if index in length.arguments:
-      raise ValueError('the length of the loop must not depend on the index')
-    self._index = index
-    self._length = length
     self._invariants, self._dependencies = _dependencies_sans_invariants(
-      Tuple([Tuple([start, stop, func]) for func, start, stop, *shape in funcdatas]), index)
+      Tuple([Tuple([start, stop, func]) for func, start, stop, *shape in funcdatas]), self._index)
     super().__init__(args=(Tuple(shapes), length, *self._invariants))
 
   @property
@@ -3775,7 +3787,6 @@ def _node_tuple(self, cache, subgraph, times):
     for arg in self._invariants:
       subcache[arg] = arg._node(cache, subgraph, times)
     loopgraph = Subgraph('Loop', subgraph)
-    subcache[self._index] = RegularNode('LoopIndex', (), dict(length=self._length._node(cache, subgraph, times)), (type(self).__name__, _Stats()), loopgraph)
     subtimes = times.get(self, collections.defaultdict(_Stats))
     concats = []
     for func, start, stop, *shape in self._funcdatas:
@@ -4181,29 +4192,36 @@ def appendaxes(func, shape):
     func = insertaxis(func, func.ndim, n)
   return func
 
-def _loop_concatenate_data(func, index, length):
+def loop_index(name, length):
+  return _LoopIndex(name, length)
+
+def loop_sum(func, index):
+  func = asarray(func)
+  index = types.strict[_LoopIndex](index)
+  return LoopSum(func, index._name, index.length)
+
+def _loop_concatenate_data(func, index):
   func = asarray(func)
+  index = types.strict[_LoopIndex](index)
   chunk_size = func.shape[-1]
   if chunk_size.isconstant:
-    chunk_sizes = InsertAxis(chunk_size, length)
+    chunk_sizes = InsertAxis(chunk_size, index.length)
   else:
-    chunk_sizes = loop_concatenate(InsertAxis(func.shape[-1], 1), index, length)
+    chunk_sizes = loop_concatenate(InsertAxis(func.shape[-1], 1), index)
   offsets = _SizesToOffsets(chunk_sizes)
   start = Take(offsets, index)
   stop = Take(offsets, index+1)
-  return (func, start, stop, *func.shape[:-1], Take(offsets, length))
+  return (func, start, stop, *func.shape[:-1], Take(offsets, index.length))
 
-def loop_concatenate(func, index, length):
-  length = asindex(length)
-  funcdata = _loop_concatenate_data(func, index, length)
-  return LoopConcatenate(funcdata, index, length)
+def loop_concatenate(func, index):
+  funcdata = _loop_concatenate_data(func, index)
+  return LoopConcatenate(funcdata, index._name, index.length)
 
-def loop_concatenate_combined(funcs, index, length):
-  length = asindex(length)
+def loop_concatenate_combined(funcs, index):
   unique_funcs = []
   unique_funcs.extend(func for func in funcs if func not in unique_funcs)
-  unique_func_data = tuple(_loop_concatenate_data(func, index, length) for func in unique_funcs)
-  loop = LoopConcatenateCombined(unique_func_data, index, length)
+  unique_func_data = tuple(_loop_concatenate_data(func, index) for func in unique_funcs)
+  loop = LoopConcatenateCombined(unique_func_data, index._name, index.length)
   return tuple(ArrayFromTuple(loop, unique_funcs.index(func), shape, func.dtype) for func, start, stop, *shape in unique_func_data)
 
 @replace
diff --git a/nutils/sample.py b/nutils/sample.py
index 968923331..d8273dd9d 100644
--- a/nutils/sample.py
+++ b/nutils/sample.py
@@ -150,7 +150,7 @@ def get_evaluable_indices(self, ielem):
   def _lower_for_loop(self, func, **kwargs):
     if kwargs.pop('transform_chains', None) or kwargs.pop('coordinates', None):
       raise ValueError('nested integrals or samples are not yet supported')
-    ielem = evaluable.Argument('_ielem', (), dtype=int)
+    ielem = evaluable.loop_index('_ielem', self.nelems)
     return ielem, func.lower(**kwargs,
       transform_chains=tuple(evaluable.TransformChainFromSequence(t, ielem) for t in self.transforms),
       coordinates=(self.points.get_evaluable_coords(ielem),) * len(self.transforms))
@@ -436,7 +436,7 @@ def __init__(self, integrand: function.Array, sample: Sample) -> None:
   def lower(self, **kwargs) -> evaluable.Array:
     ielem, integrand = self._sample._lower_for_loop(self._integrand, **kwargs)
     contracted = evaluable.dot(evaluable.appendaxes(self._sample.points.get_evaluable_weights(ielem), integrand.shape[1:]), integrand, 0)
-    return evaluable.LoopSum(contracted, ielem, self._sample.nelems)
+    return evaluable.loop_sum(contracted, ielem)
 
 class _AtSample(function.Array):
 
@@ -449,7 +449,7 @@ def lower(self, **kwargs) -> evaluable.Array:
     ielem, func = self._sample._lower_for_loop(self._func, **kwargs)
     indices = self._sample.get_evaluable_indices(ielem)
     inflated = evaluable.Transpose.from_end(evaluable.Inflate(evaluable.Transpose.to_end(func, 0), indices, self._sample.npoints), 0)
-    return evaluable.LoopSum(inflated, ielem, self._sample.nelems)
+    return evaluable.loop_sum(inflated, ielem)
 
 class _Basis(function.Array):
 
@@ -467,8 +467,8 @@ def lower(self, *, transform_chains=(), coordinates=(), **kwargs):
     return evaluable.Inflate(sampled, dofmap=indices, length=self._sample.npoints)
 
 def _offsets(pointsseq):
-  ielem = evaluable.Argument('_ielem', shape=(), dtype=int)
+  ielem = evaluable.loop_index('_ielem', len(pointsseq))
   npoints, ndims = pointsseq.get_evaluable_coords(ielem).shape
-  return evaluable._SizesToOffsets(evaluable.loop_concatenate(evaluable.InsertAxis(npoints, 1), ielem, len(pointsseq)))
+  return evaluable._SizesToOffsets(evaluable.loop_concatenate(evaluable.InsertAxis(npoints, 1), ielem))
 
 # vim:sw=2:sts=2:et
diff --git a/tests/test_evaluable.py b/tests/test_evaluable.py
index d8022869a..e89bed8ba 100644
--- a/tests/test_evaluable.py
+++ b/tests/test_evaluable.py
@@ -285,25 +285,25 @@ def test_unravel(self):
         actual=evaluable.unravel(self.op_args, axis=idim, shape=unravelshape))
 
   def test_loopsum(self):
-    index = evaluable.Argument('_testindex', (), int)
     length = 3
+    index = evaluable.loop_index('_testindex', length)
     for iarg, shape in enumerate(self.shapes):
       testvalue = numpy.random.uniform(size=(length, *shape), low=self.low, high=self.high)
       n_op_argsfun = functools.reduce(operator.add, (self.n_op(*self.arg_values[:iarg], v, *self.arg_values[iarg+1:]) for v in testvalue))
       args = (*self.args[:iarg], evaluable.Guard(evaluable.get(evaluable.asarray(testvalue), 0, index)), *self.args[iarg+1:])
       self.assertFunctionAlmostEqual(decimal=15,
-        actual=evaluable.LoopSum(self.op(*args), index, length),
+        actual=evaluable.loop_sum(self.op(*args), index),
         desired=n_op_argsfun)
 
   def test_loopconcatenate(self):
-    index = evaluable.Argument('_testindex', (), int)
     length = 3
+    index = evaluable.loop_index('_testindex', length)
     for iarg, shape in enumerate(self.shapes):
       testvalue = numpy.random.uniform(size=(length, *shape), low=self.low, high=self.high)
       n_op_argsfun = numpy.concatenate([self.n_op(*self.arg_values[:iarg], v, *self.arg_values[iarg+1:]) for v in testvalue], axis=-1)
       args = (*self.args[:iarg], evaluable.Guard(evaluable.get(evaluable.asarray(testvalue), 0, index)), *self.args[iarg+1:])
       self.assertFunctionAlmostEqual(decimal=15,
-        actual=evaluable.loop_concatenate(self.op(*args), index, length),
+        actual=evaluable.loop_concatenate(self.op(*args), index),
         desired=n_op_argsfun)
 
   def test_desparsify(self):
@@ -446,13 +446,14 @@ def _check(name, op, n_op, shapes, hasgrad=True, zerograd=False, ndim=2, low=-1,
 _check('normal1d', lambda a: evaluable.Normal(a), lambda a: numpy.sign(a[...,0]), [(3,1,1)])
 _check('normal2d', lambda a: evaluable.Normal(a), lambda a: numpy.stack([Q[:,-1]*numpy.sign(R[-1,-1]) for ai in a for Q, R in [numpy.linalg.qr(ai, mode='complete')]], axis=0), [(1,2,2)])
 _check('normal3d', lambda a: evaluable.Normal(a), lambda a: numpy.stack([Q[:,-1]*numpy.sign(R[-1,-1]) for ai in a for Q, R in [numpy.linalg.qr(ai, mode='complete')]], axis=0), [(2,3,3)])
-_check('loopsum1', lambda: evaluable.LoopSum(evaluable.Argument('index', (), int), evaluable.Argument('index', (), int), 3), lambda: numpy.array(3), [])
-_check('loopsum2', lambda a: evaluable.LoopSum(a, evaluable.Argument('index', (), int), 2), lambda a: 2*a, [(3,4,2,4)])
-_check('loopsum3', lambda a: evaluable.LoopSum(evaluable.get(a, 0, evaluable.Argument('index', (), int)), evaluable.Argument('index', (), int), 3), lambda a: numpy.sum(a, 0), [(3,4,2,4)])
-_check('loopsum4', lambda: evaluable.LoopSum(evaluable.Inflate(evaluable.Argument('index', (), int), 0, 2), evaluable.Argument('index', (), int), 3), lambda: numpy.array([3, 0]), [])
-_check('loopconcatenate1', lambda a: evaluable.loop_concatenate(a+evaluable.prependaxes(evaluable.Argument('index', (), int), a.shape), evaluable.Argument('index', (), int), 3), lambda a: a+numpy.arange(3)[None], [(2,1)])
-_check('loopconcatenate2', lambda: evaluable.loop_concatenate(evaluable.Elemwise([numpy.arange(48).reshape(4,4,3)[:,:,a:b] for a, b in util.pairwise([0,2,3])], evaluable.Argument('index', (), int), int), evaluable.Argument('index', (), int), 2), lambda: numpy.arange(48).reshape(4,4,3), [])
-_check('loopconcatenatecombined', lambda a: evaluable.loop_concatenate_combined([a+evaluable.prependaxes(evaluable.Argument('index', (), int), a.shape)], evaluable.Argument('index', (), int), 3)[0], lambda a: a+numpy.arange(3)[None], [(2,1)], hasgrad=False)
+_check('loopsum1', lambda: evaluable.loop_sum(evaluable.loop_index('index', 3), evaluable.loop_index('index', 3)), lambda: numpy.array(3), [])
+_check('loopsum2', lambda a: evaluable.loop_sum(a, evaluable.loop_index('index', 2)), lambda a: 2*a, [(3,4,2,4)])
+_check('loopsum3', lambda a: evaluable.loop_sum(evaluable.get(a, 0, evaluable.loop_index('index', 3)), evaluable.loop_index('index', 3)), lambda a: numpy.sum(a, 0), [(3,4,2,4)])
+_check('loopsum4', lambda: evaluable.loop_sum(evaluable.Inflate(evaluable.loop_index('index', 3), 0, 2), evaluable.loop_index('index', 3)), lambda: numpy.array([3, 0]), [])
+_check('loopsum5', lambda: evaluable.loop_sum(evaluable.loop_index('index', 1), evaluable.loop_index('index', 1)), lambda: numpy.array(0), [])
+_check('loopconcatenate1', lambda a: evaluable.loop_concatenate(a+evaluable.prependaxes(evaluable.loop_index('index', 3), a.shape), evaluable.loop_index('index', 3)), lambda a: a+numpy.arange(3)[None], [(2,1)])
+_check('loopconcatenate2', lambda: evaluable.loop_concatenate(evaluable.Elemwise([numpy.arange(48).reshape(4,4,3)[:,:,a:b] for a, b in util.pairwise([0,2,3])], evaluable.loop_index('index', 2), int), evaluable.loop_index('index', 2)), lambda: numpy.arange(48).reshape(4,4,3), [])
+_check('loopconcatenatecombined', lambda a: evaluable.loop_concatenate_combined([a+evaluable.prependaxes(evaluable.loop_index('index', 3), a.shape)], evaluable.loop_index('index', 3))[0], lambda a: a+numpy.arange(3)[None], [(2,1)], hasgrad=False)
 
 _polyval_mask = lambda shape, ndim: 1 if ndim == 0 else numpy.array([sum(i[-ndim:]) < int(shape[-1]) for i in numpy.ndindex(shape)], dtype=int).reshape(shape)
 _polyval_desired = lambda c, x: sum(c[(...,*i)]*(x[(slice(None),*[None]*(c.ndim-x.shape[1]))]**i).prod(-1) for i in itertools.product(*[range(c.shape[-1])]*x.shape[1]) if sum(i) < c.shape[-1])
@@ -871,8 +872,8 @@ def test_asciitree(self):
 
   @unittest.skipIf(sys.version_info < (3, 6), 'test requires dicts maintaining insertion order')
   def test_loop_sum(self):
-    i = evaluable.Argument('i', (), int)
-    f = evaluable.LoopSum(i, i, evaluable.Constant(2))
+    i = evaluable.loop_index('i', 2)
+    f = evaluable.loop_sum(i, i)
     self.assertEqual(f.asciitree(richoutput=True),
                      'SUBGRAPHS\n'
                      'A\n'
@@ -884,8 +885,8 @@ def test_loop_sum(self):
 
   @unittest.skipIf(sys.version_info < (3, 6), 'test requires dicts maintaining insertion order')
   def test_loop_concatenate(self):
-    i = evaluable.Argument('i', (), int)
-    f = evaluable.loop_concatenate(evaluable.InsertAxis(i, 1), i, evaluable.Constant(2))
+    i = evaluable.loop_index('i', 2)
+    f = evaluable.loop_concatenate(evaluable.InsertAxis(i, 1), i)
     self.assertEqual(f.asciitree(richoutput=True),
                      'SUBGRAPHS\n'
                      'A\n'
@@ -913,8 +914,8 @@ def test_loop_concatenate(self):
 
   @unittest.skipIf(sys.version_info < (3, 6), 'test requires dicts maintaining insertion order')
   def test_loop_concatenatecombined(self):
-    i = evaluable.Argument('i', (), int)
-    f, = evaluable.loop_concatenate_combined([evaluable.InsertAxis(i, 1)], i, evaluable.Constant(2))
+    i = evaluable.loop_index('i', 2)
+    f, = evaluable.loop_concatenate_combined([evaluable.InsertAxis(i, 1)], i)
     self.assertEqual(f.asciitree(richoutput=True),
                      'SUBGRAPHS\n'
                      'A\n'
@@ -987,57 +988,47 @@ def _simplified(self):
 
 class combine_loop_concatenates(TestCase):
 
-  def test_same_index_same_length(self):
-    i = evaluable.Argument('i', (), int)
-    A = evaluable.LoopConcatenate((evaluable.InsertAxis(i, 1), i, i+1, 3,), i, 3)
-    B = evaluable.LoopConcatenate((evaluable.InsertAxis(i, 2), i*2, i*2+2, 6,), i, 3)
+  def test_same_index(self):
+    i = evaluable.loop_index('i', 3)
+    A = evaluable.LoopConcatenate((evaluable.InsertAxis(i, 1), i, i+1, 3,), i._name, i.length)
+    B = evaluable.LoopConcatenate((evaluable.InsertAxis(i, 2), i*2, i*2+2, 6,), i._name, i.length)
     actual = evaluable.Tuple((A, B))._combine_loop_concatenates(set())
-    L = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(i, 1), i, i+1, 3), (evaluable.InsertAxis(i, 2), i*2, i*2+2, 6)), i, 3)
+    L = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(i, 1), i, i+1, 3), (evaluable.InsertAxis(i, 2), i*2, i*2+2, 6)), i._name, i.length)
     desired = evaluable.Tuple((evaluable.ArrayFromTuple(L, 0, (3,), int), evaluable.ArrayFromTuple(L, 1, (6,), int)))
     self.assertEqual(actual, desired)
 
-  def test_same_index_different_length(self):
-    i = evaluable.Argument('i', (), int)
-    A = evaluable.LoopConcatenate((evaluable.InsertAxis(i, 1), i, i+1, 3,), i, 3)
-    B = evaluable.LoopConcatenate((evaluable.InsertAxis(i, 1), i, i+1, 4,), i, 4)
-    actual = evaluable.Tuple((A, B))._combine_loop_concatenates(set())
-    L1 = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(i, 1), i, i+1, 3),), i, 3)
-    L2 = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(i, 1), i, i+1, 4),), i, 4)
-    desired = evaluable.Tuple((evaluable.ArrayFromTuple(L1, 0, (3,), int), evaluable.ArrayFromTuple(L2, 0, (4,), int)))
-    self.assertEqual(actual, desired)
-
   def test_different_index(self):
-    i = evaluable.Argument('i', (), int)
-    j = evaluable.Argument('j', (), int)
-    A = evaluable.LoopConcatenate((evaluable.InsertAxis(i, 1), i, i+1, 3,), i, 3)
-    B = evaluable.LoopConcatenate((evaluable.InsertAxis(j, 1), j, j+1, 3,), j, 3)
+    i = evaluable.loop_index('i', 3)
+    j = evaluable.loop_index('j', 3)
+    A = evaluable.LoopConcatenate((evaluable.InsertAxis(i, 1), i, i+1, 3,), i._name, i.length)
+    B = evaluable.LoopConcatenate((evaluable.InsertAxis(j, 1), j, j+1, 3,), j._name, j.length)
     actual = evaluable.Tuple((A, B))._combine_loop_concatenates(set())
-    L1 = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(i, 1), i, i+1, 3),), i, 3)
-    L2 = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(j, 1), j, j+1, 3),), j, 3)
+    L1 = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(i, 1), i, i+1, 3),), i._name, i.length)
+    L2 = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(j, 1), j, j+1, 3),), j._name, j.length)
     desired = evaluable.Tuple((evaluable.ArrayFromTuple(L1, 0, (3,), int), evaluable.ArrayFromTuple(L2, 0, (3,), int)))
     self.assertEqual(actual, desired)
 
   def test_nested_invariant(self):
-    i = evaluable.Argument('i', (), int)
-    A = evaluable.LoopConcatenate((evaluable.InsertAxis(i, 1), i, i+1, 3,), i, 3)
-    B = evaluable.LoopConcatenate((A, i*3, i*3+3, 9,), i, 3)
+    i = evaluable.loop_index('i', 3)
+    A = evaluable.LoopConcatenate((evaluable.InsertAxis(i, 1), i, i+1, 3,), i._name, i.length)
+    B = evaluable.LoopConcatenate((A, i*3, i*3+3, 9,), i._name, i.length)
     actual = evaluable.Tuple((A, B))._combine_loop_concatenates(set())
-    L1 = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(i, 1), i, i+1, 3),), i, 3)
+    L1 = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(i, 1), i, i+1, 3),), i._name, i.length)
     A_ = evaluable.ArrayFromTuple(L1, 0, (3,), int)
-    L2 = evaluable.LoopConcatenateCombined(((A_, i*3, i*3+3, 9),), i, 3)
+    L2 = evaluable.LoopConcatenateCombined(((A_, i*3, i*3+3, 9),), i._name, i.length)
     self.assertIn(A_, L2._Evaluable__args)
     desired = evaluable.Tuple((A_, evaluable.ArrayFromTuple(L2, 0, (9,), int)))
     self.assertEqual(actual, desired)
 
   def test_nested_variant(self):
-    i = evaluable.Argument('i', (), int)
-    j = evaluable.Argument('j', (), int)
-    A = evaluable.LoopConcatenate((evaluable.InsertAxis(i+j, 1), i, i+1, 3,), i, 3)
-    B = evaluable.LoopConcatenate((A, j*3, j*3+3, 9,), j, 3)
+    i = evaluable.loop_index('i', 3)
+    j = evaluable.loop_index('j', 3)
+    A = evaluable.LoopConcatenate((evaluable.InsertAxis(i+j, 1), i, i+1, 3,), i._name, i.length)
+    B = evaluable.LoopConcatenate((A, j*3, j*3+3, 9,), j._name, j.length)
     actual = evaluable.Tuple((A, B))._combine_loop_concatenates(set())
-    L1 = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(i+j, 1), i, i+1, 3),), i, 3)
+    L1 = evaluable.LoopConcatenateCombined(((evaluable.InsertAxis(i+j, 1), i, i+1, 3),), i._name, i.length)
     A_ = evaluable.ArrayFromTuple(L1, 0, (3,), int)
-    L2 = evaluable.LoopConcatenateCombined(((A_, j*3, j*3+3, 9),), j, 3)
+    L2 = evaluable.LoopConcatenateCombined(((A_, j*3, j*3+3, 9),), j._name, j.length)
     self.assertNotIn(A_, L2._Evaluable__args)
     desired = evaluable.Tuple((A_, evaluable.ArrayFromTuple(L2, 0, (9,), int)))
     self.assertEqual(actual, desired)