Improvements for solver.System

.github/workflows/release.yml

@@ -10,10 +10,14 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
       - name: Install dependencies
-        run: python3 -m pip --break-system-packages install flit
+        run: python -um pip install flit
       - name: Build package
-        run: python3 -m flit build
+        run: python -um flit build
       - name: Publish package to PyPI
         uses: pypa/[email protected]
         with:

.github/workflows/test.yaml

@@ -19,15 +19,19 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
       - name: Install build dependencies
-        run: python3 -m pip install --break-system-packages flit
+        run: python -um pip install flit
       - name: Build package
         id: build
         run: |
           # To make the wheels reproducible, set the timestamp of the (files in
           # the) generated wheels to the date of the commit.
           export SOURCE_DATE_EPOCH=`git show -s --format=%ct`
-          python3 -m flit build
+          python -um flit build
           echo wheel=`echo dist/*.whl` >> $GITHUB_OUTPUT
       - name: Upload package artifacts
         uses: actions/upload-artifact@v4
@@ -181,12 +185,16 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
       - name: Install dependencies
         run: |
-          python3 -um pip install --break-system-packages setuptools wheel
-          python3 -um pip install --break-system-packages --upgrade --upgrade-strategy eager .[docs]
+          python -um pip install setuptools wheel
+          python -um pip install --upgrade --upgrade-strategy eager .[docs]
       - name: Build docs
-        run: python3 -m sphinx -n -W --keep-going docs build/sphinx/html
+        run: python -um sphinx -n -W --keep-going docs build/sphinx/html
   build-and-test-container-image:
     name: Build container image
     needs: build-python-package

examples/adaptivity.py

@@ -69,10 +69,10 @@ def main(etype: str = 'square',
         ns.du = 'u - uexact'
 
         sqr = domain.boundary['corner'].integral('u^2 dS' @ ns, degree=degree*2)
-        cons = System(sqr, trial='u').optimize(droptol=1e-15)
+        cons = System(sqr, trial='u').solve_constraints(droptol=1e-15)
 
         sqr = domain.boundary.integral('du^2 dS' @ ns, degree=7)
-        cons = System(sqr, trial='u').optimize(droptol=1e-15, constrain=cons)
+        cons = System(sqr, trial='u').solve_constraints(droptol=1e-15, constrain=cons)
 
         res = domain.integral('∇_k(v) ∇_k(u) dV' @ ns, degree=degree*2)
         args = System(res, trial='u', test='v').solve(constrain=cons)

examples/cahnhilliard.py

@@ -114,7 +114,7 @@ def main(size: Length = parse('10cm'),
 
     For this reason, the `stab` enum in this script defines the stabilizing term
     `δψ`, rather than `f`, allowing Nutils to construct the residuals through
-    automatic differentiation using the `optimize` method.
+    symbolic differentiation.
 
     Parameters
     ----------

examples/cylinderflow.py

@@ -129,7 +129,7 @@ def main(nelems: int = 99,
     ns.uwall_i = 'rotation ε_ij x_j' # clockwise positive rotation
 
     sqr = domain.boundary['inflow'].integral('Σ_i (u_i - uinf_i)^2 dS' @ ns, degree=degree*2)
-    cons = System(sqr, trial='u').optimize(droptol=1e-15) # constrain inflow boundary to unit horizontal flow
+    cons = System(sqr, trial='u').solve_constraints(droptol=1e-15) # constrain inflow boundary to unit horizontal flow
 
     sqr = domain.integral('(.5 Σ_i (u_i - uinf_i)^2 - ∇_k(u_k) p) dV' @ ns, degree=degree*2)
     args = System(sqr, trial='u,p').solve(constrain=cons) # set initial condition to potential flow

examples/drivencavity.py

@@ -118,12 +118,12 @@ def main(nelems: int = 32,
 
     # strong enforcement of non-penetrating boundary conditions
     sqr = domain.boundary.integral('(u_k n_k)^2 dS' @ ns, degree=degree*2)
-    cons = System(sqr, trial='u').optimize(droptol=1e-15)
+    cons = System(sqr, trial='u').solve_constraints(droptol=1e-15)
 
     if strongbc:
         # strong enforcement of tangential boundary conditions
         sqr = domain.boundary.integral('(ε_ij n_i (u_j - uwall_j))^2 dS' @ ns, degree=degree*2)
-        tcons = System(sqr, trial='u').optimize(droptol=1e-15)
+        tcons = System(sqr, trial='u').solve_constraints(droptol=1e-15)
         cons['u'] = numpy.choose(numpy.isnan(cons['u']), [cons['u'], tcons['u']])
     else:
         # weak enforcement of tangential boundary conditions via Nitsche's method

examples/elasticity.py

@@ -51,7 +51,7 @@ def main(nelems: int = 24,
     ns.q_i = '-δ_i1'
 
     sqr = domain.boundary['top'].integral('u_k u_k dS' @ ns, degree=degree*2)
-    cons = System(sqr, trial='u').optimize(droptol=1e-15)
+    cons = System(sqr, trial='u').solve_constraints(droptol=1e-15)
 
     # solve for equilibrium configuration
     energy = domain.integral('(E - u_i q_i) dV' @ ns, degree=degree*2)

examples/finitestrain.py

@@ -58,7 +58,7 @@ def main(nelems: int = 20,
 
     sqr = domain.boundary['left'].integral('u_k u_k dS' @ ns, degree=degree*2)
     sqr += domain.boundary['right'].integral('((u_0 - X_1 sin(2 angle) - cos(angle) + 1)^2 + (u_1 - X_1 (cos(2 angle) - 1) + sin(angle))^2) dS' @ ns, degree=degree*2)
-    cons = System(sqr, trial='u').optimize(droptol=1e-15)
+    cons = System(sqr, trial='u').solve_constraints(droptol=1e-15)
 
     energy = domain.integral('energy dV' @ ns, degree=degree*2)
     args0 = System(energy, trial='u').solve(constrain=cons)

examples/laplace.py

@@ -72,7 +72,7 @@ def main(nelems: int = 10,
 
     sqr = domain.boundary['left'].integral('u^2 dS' @ ns, degree=degree*2)
     sqr += domain.boundary['top'].integral('(u - cosh(1) sin(x_0))^2 dS' @ ns, degree=degree*2)
-    cons = System(sqr, trial='u').optimize(droptol=1e-15)
+    cons = System(sqr, trial='u').solve_constraints(droptol=1e-15)
 
     # The unconstrained entries of `u` are to be such that the residual
     # evaluates to zero for all possible values of `v`. The resulting array `u`

examples/platewithhole.py

@@ -137,10 +137,10 @@ def main(mode: Union[FCM, NURBS] = NURBS(),
     log.info('hole radius exact up to L2 error {:.2e}'.format(radiuserr))
 
     sqr = topo.boundary['sym'].integral('(u_i n_i)^2 dS' @ ns, degree=degree*2)
-    cons = System(sqr, trial='u').optimize(droptol=1e-15)
+    cons = System(sqr, trial='u').solve_constraints(droptol=1e-15)
 
     sqr = topo.boundary['far'].integral('du_k du_k dS' @ ns, degree=20)
-    cons = System(sqr, trial='u').optimize(droptol=1e-15, constrain=cons)
+    cons = System(sqr, trial='u').solve_constraints(droptol=1e-15, constrain=cons)
 
     res = topo.integral('∇_j(v_i) σ_ij dV' @ ns, degree=degree*2)
     args = System(res, trial='u', test='v').solve(constrain=cons)

examples/poisson.py

@@ -25,7 +25,7 @@ def main(nelems: int = 32):
     J = function.J(x)
 
     sqr = topo.boundary.integral(u**2 * J, degree=2)
-    cons = System(sqr, trial='u').optimize(droptol=1e-12)
+    cons = System(sqr, trial='u').solve_constraints(droptol=1e-12)
 
     energy = topo.integral((g @ g / 2 - u) * J, degree=1)
     args = System(energy, trial='u').solve(constrain=cons)

nutils/evaluable.py

@@ -691,16 +691,21 @@
         return NotImplemented
 
 
+def assert_equal(a, b):
+    return a if a == b else AssertEqual(a, b)
+
+
+def assert_equal_tuple(A, B):
+    assert len(A) == len(B)
+    return tuple(map(assert_equal, A, B))
+
+
 class AssertEqual(Array):
     'Confirm arrays equality at runtime'
 
     a: Array
     b: Array
 
-    @classmethod
-    def map(cls, A, B):
-        return tuple(map(cls, A, B))
-
     def __post_init__(self):
         assert not _certainly_different(self.a, self.b)
 
@@ -710,16 +715,23 @@
 
     @cached_property
     def shape(self):
-        return self.map(self.a.shape, self.b.shape)
+        return assert_equal_tuple(self.a.shape, self.b.shape)
 
     def _intbounds_impl(self):
         lowera, uppera = self.a._intbounds_impl()
         lowerb, upperb = self.b._intbounds_impl()
         return max(lowera, lowerb), min(uppera, upperb)
 
     def _simplified(self):
-        if self.a == self.b:
-            return self.a
+        unique = []
+        queue = [self.a, self.b]
+        for item in queue:
+            if isinstance(item, AssertEqual):
+                queue.append(item.a)
+                queue.append(item.b)
+            elif item not in unique:
+                unique.append(item)
+        return functools.reduce(AssertEqual, unique)
 
     @property
     def dependencies(self):
@@ -1442,7 +1454,7 @@
     @cached_property
     def shape(self):
         func1, func2 = self.funcs
-        return AssertEqual.map(func1.shape, func2.shape)
+        return assert_equal_tuple(func1.shape, func2.shape)
 
     @property
     def _factors(self):
@@ -1629,7 +1641,7 @@
     @cached_property
     def shape(self):
         func1, func2 = self.funcs
-        return AssertEqual.map(func1.shape, func2.shape)
+        return assert_equal_tuple(func1.shape, func2.shape)
 
     @cached_property
     def _inflations(self):
@@ -1775,7 +1787,7 @@
         for idx, arg in zip(self.args_idx, self.args):
             for i, length in zip(idx, arg.shape):
                 n = lengths.get(i)
-                lengths[i] = length if n is None else AssertEqual(length, n)
+                lengths[i] = length if n is None else assert_equal(length, n)
         try:
             self.shape = tuple(lengths[i] for i in self.out_idx)
         except KeyError(e):
@@ -2101,7 +2113,7 @@
     def __post_init__(self):
         self.shape = self.dependencies[0].shape
         for dep in self.dependencies[1:]:
-            self.shape = AssertEqual.map(self.shape, dep.shape)
+            self.shape = assert_equal_tuple(self.shape, dep.shape)
 
     def _newargs(self, *args):
         '''

nutils/solver.py

@@ -237,29 +237,47 @@
         self.trialshapes = dict(zip(self.trials, evaluable.compile(tuple(argobjects[t].shape for t in self.trials))()))
         self.__trial_offsets = numpy.cumsum([0] + [numpy.prod(self.trialshapes[t], dtype=int) for t in self.trials])
 
-        value = functional if self.is_symmetric else evaluable.constant(numpy.nan)
+        value = functional if self.is_symmetric else ()
         block_vector = [evaluable._flat(res) for res in residuals]
         block_matrix = [[evaluable._flat(evaluable.derivative(res, argobjects[t]).simplified, 2) for t in self.trials] for res in block_vector]
-        self.__eval = evaluable.compile((tuple(tuple(map(evaluable.as_csr, row)) for row in block_matrix), tuple(block_vector), value))
 
         self.is_linear = not any(arg.name in self.trials for row in block_matrix for col in row for arg in col.arguments)
-        self.is_constant = self.is_linear and not any(col.arguments for row in block_matrix for col in row)
-        self.__matrix_cache = None
+        if self.is_linear:
+            z = {t: evaluable.zeros_like(argobjects[t]) for t in self.trials}
+            block_vector = [evaluable.replace_arguments(vector, z).simplified for vector in block_vector]
+            if self.is_symmetric:
+                value = evaluable.replace_arguments(value, z).simplified
+
+        self.__eval = evaluable.compile((tuple(tuple(map(evaluable.as_csr, row)) for row in block_matrix), tuple(block_vector), value))
+
+        self.is_constant_matrix = self.is_linear and not any(col.arguments for row in block_matrix for col in row)
+        self.is_constant = self.is_constant_matrix and not any(vec.arguments for vec in block_vector) and not (self.is_symmetric and value.arguments)
+        self.__mat_vec_val = None, None, None
 
     @property
     def __nutils_hash__(self):
         return self.__eval.__nutils_hash__
 
     @log.withcontext
     def assemble(self, arguments: Dict[str, numpy.ndarray]):
-        mat_blocks, vec_blocks, val = self.__eval(**arguments)
-        vec = numpy.concatenate(vec_blocks)
-        if self.__matrix_cache is not None:
-            mat = self.__matrix_cache
-        else:
-            mat = matrix.assemble_block_csr(mat_blocks)
+        mat, vec, val = self.__mat_vec_val
+        if vec is None:
+            mat_blocks, vec_blocks, maybe_val = self.__eval(**arguments)
+            if mat is None:
+                mat = matrix.assemble_block_csr(mat_blocks)
+            vec = numpy.concatenate(vec_blocks)
+            val = self.dtype(maybe_val) if self.is_symmetric else None
             if self.is_constant:
-                self.__matrix_cache = mat
+                vec.flags.writeable = False
+                self.__mat_vec_val = mat, vec, val
+            elif self.is_constant_matrix:
+                self.__mat_vec_val = mat, None, None
+        if self.is_linear:
+            x = numpy.concatenate([arguments[t].ravel() for t in self.trials])
+            matx = mat @ x
+            if self.is_symmetric:
+                val += vec @ x + .5 * (x @ matx) # val(x) = val(0) + vec(0) x + .5 x mat x
+            vec = vec + matx # vec(x) = vec(0) + mat x
         return mat, vec, val
 
     def prepare_solution_vector(self, arguments: Dict[str, numpy.ndarray], constrain: Dict[str, numpy.ndarray]):
@@ -389,14 +407,22 @@
 
     @cache.function
     @log.withcontext
-    def optimize(self, *, droptol: Optional[float], arguments: Dict[str, numpy.ndarray] = {}, constrain: Dict[str, numpy.ndarray] = {}, linargs: Dict[str, Any] = {}) -> Tuple[Dict[str, numpy.ndarray], float]:
-        '''Optimize singular system.
+    def solve_constraints(self, *, droptol: Optional[float], arguments: Dict[str, numpy.ndarray] = {}, constrain: Dict[str, numpy.ndarray] = {}, linargs: Dict[str, Any] = {}) -> Tuple[Dict[str, numpy.ndarray], float]:
+        '''Solve for Dirichlet constraints.
 
-        This method is similar to ``solve``, but specific to symmetric linear
-        systems, with the added ability to solve a limited class of singular
-        systems. It does so by isolating the subset of arguments that
-        contribute to the functional and solving the corresponding submatrix.
-        The remaining argument values are returned as NaN (not a number).
+        This method is similar to ``solve``, but with two key differences.
+
+        The method is limited to linear systems, but adds the ability to solve
+        a limited class of singular systems. It does so by isolating the subset
+        of arguments that contribute (up to droptol) to the residual, and
+        solving the corresponding submatrix. The remaining argument values are
+        returned as NaN (not a number).
+
+        The second key difference with solve is that the returned dictionary
+        is augmented with the remaining _constrain_ items, rather than those
+        from _arguments_, reflecting the method's main utility of forming
+        Dirichlet constraints. This allows for the aggregation of constraints
+        by calling the method multiple times in series.
 
         Parameters
         ----------
@@ -414,19 +440,23 @@
             participate in the optimization problem.
         '''
 
-        log.info(f'optimizing for argument {self._trial_info} with drop tolerance {droptol:.0e}')
+        log.info(f'{"optimizing" if self.is_symmetric else "solving"} for argument {self._trial_info} with drop tolerance {droptol:.0e}')
         if not self.is_linear:
-            raise SolverError('system is not linear')
-        if not self.is_symmetric:
-            raise SolverError('system is not symmetric')
+            raise ValueError('system is not linear')
         x, iscons, arguments = self.prepare_solution_vector(arguments, constrain)
         jac, res, val = self.assemble(arguments)
-        nosupp = ~jac.rowsupp(droptol)
-        dx = -jac.solve(res, constrain=iscons | nosupp, **_copy_with_defaults(linargs, symmetric=self.is_symmetric))
-        log.info(f'optimal value: {val+.5*(res@dx):.1e}') # val(x + dx) = val(x) + res(x) dx + .5 dx jac dx
+        data, colidx, _ = jac.export('csr')
+        mycons = numpy.ones_like(iscons)
+        mycons[colidx[abs(data) > droptol]] = False # unconstrain dofs with nonzero columns
+        mycons |= iscons
+        dx = -jac.solve(res, constrain=mycons, **_copy_with_defaults(linargs, symmetric=self.is_symmetric))
+        if self.is_symmetric:
+            log.info(f'optimal value: {val+.5*(res@dx):.1e}') # val(x + dx) = val(x) + res(x) dx + .5 dx jac dx
         x += dx
-        x[nosupp & ~iscons] = numpy.nan
-        return arguments
+        for trial, i, j in zip(self.trials, self.__trial_offsets, self.__trial_offsets[1:]):
+            log.info(f'constrained {j-i-mycons[i:j].sum()} degrees of freedom of {trial}')
+        x[mycons & ~iscons] = numpy.nan
+        return dict(constrain, **{t: arguments[t] for t in self.trials})
 
 
 @dataclass(eq=True, frozen=True)
@@ -492,11 +522,11 @@
     failrelax: float = -10.
 
     def __str__(self):
-        return 'newton'
+        return 'minimize'
 
     def __call__(self, system, *, arguments: Dict[str, numpy.ndarray] = {}, constrain: Dict[str, numpy.ndarray] = {}, linargs: Dict[str, Any] = {}) -> System.MethodIter:
         if not system.is_symmetric:
-            raise SolverError('problem is not symmetric')
+            raise ValueError('problem is not symmetric')
         x, iscons, arguments = system.prepare_solution_vector(arguments, constrain)
         jac, res, val = system.assemble(arguments)
         linargs = _copy_with_defaults(linargs, rtol=1-3, symmetric=system.is_symmetric)
@@ -889,7 +919,7 @@
         raise TypeError('unexpected keyword arguments: {}'.format(', '.join(kwargs)))
     system = System(functional, *_split_trial_test(target))
     if droptol is not None:
-        return system.optimize(arguments=arguments, constrain=constrain or {}, linargs=linargs, droptol=droptol)
+        return system.solve_constraints(arguments=arguments, constrain=constrain or {}, linargs=linargs, droptol=droptol)
     method = None if system.is_linear else Newton() if linesearch is None else LinesearchNewton(strategy=linesearch, relax0=relax0, failrelax=failrelax)
     return system.solve(arguments=arguments, constrain=constrain or {}, linargs=linargs, method=method, tol=tol)