diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 7494cad3d..c5cc33596 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -11,7 +11,7 @@ jobs: - name: Checkout uses: actions/checkout@v4 - name: Install dependencies - run: python3 -m pip install flit + run: python3 -m pip --break-system-packages install flit - name: Build package run: python3 -m flit build - name: Publish package to PyPI diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index ea91d6162..2a4109bba 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -20,7 +20,7 @@ jobs: - name: Checkout uses: actions/checkout@v4 - name: Install build dependencies - run: python3 -m pip install flit + run: python3 -m pip install --break-system-packages flit - name: Build package id: build run: | @@ -183,8 +183,8 @@ jobs: uses: actions/checkout@v4 - name: Install dependencies run: | - python3 -um pip install setuptools wheel - python3 -um pip install --upgrade --upgrade-strategy eager .[docs] + python3 -um pip install --break-system-packages setuptools wheel + python3 -um pip install --break-system-packages --upgrade --upgrade-strategy eager .[docs] - name: Build docs run: python3 -m sphinx -n -W --keep-going docs build/sphinx/html build-and-test-container-image: diff --git a/examples/adaptivity.py b/examples/adaptivity.py index 046f9be32..9716cfee2 100644 --- a/examples/adaptivity.py +++ b/examples/adaptivity.py @@ -1,4 +1,5 @@ -from nutils import mesh, function, solver, export, testing +from nutils import mesh, function, export, testing +from nutils.solver import System from nutils.expression_v2 import Namespace import numpy import treelog @@ -68,13 +69,13 @@ def main(etype: str = 'square', ns.du = 'u - uexact' sqr = domain.boundary['corner'].integral('u^2 dS' @ ns, degree=degree*2) - cons = solver.optimize('u,', sqr, droptol=1e-15) + cons = System(sqr, trial='u').optimize(droptol=1e-15) sqr = domain.boundary.integral('du^2 dS' @ ns, degree=7) - cons = solver.optimize('u,', sqr, droptol=1e-15, constrain=cons) + cons = System(sqr, trial='u').optimize(droptol=1e-15, constrain=cons) res = domain.integral('∇_k(v) ∇_k(u) dV' @ ns, degree=degree*2) - args = solver.solve_linear('u:v', res, constrain=cons) + args = System(res, trial='u', test='v').solve(constrain=cons) ndofs = len(args['u']) error = numpy.sqrt(domain.integral(['du^2 dV', '(du^2 + ∇_k(du) ∇_k(du)) dV'] @ ns, degree=7)).eval(**args) diff --git a/examples/burgers.py b/examples/burgers.py index e7bd0baa1..a5d4586b4 100644 --- a/examples/burgers.py +++ b/examples/burgers.py @@ -1,4 +1,5 @@ -from nutils import mesh, function, solver, export, testing +from nutils import mesh, function, export, testing +from nutils.solver import System from nutils.expression_v2 import Namespace import treelog as log import numpy @@ -40,8 +41,8 @@ def main(nelems: int = 40, ns.x = geom ns.define_for('x', gradient='∇', normal='n', jacobians=('dV', 'dS')) ns.add_field(('u', 'u0', 'v'), domain.basis(btype, degree=degree)) - ns.dt = timestep - ns.dudt = '(u - u0) / dt' + ns.add_field(('t', 't0')) + ns.dudt = '(u - u0) / (t - t0)' ns.f = '.5 u^2' ns.C = 1 ns.uinit = 'exp(-25 x^2)' @@ -50,18 +51,20 @@ def main(nelems: int = 40, res -= domain.interfaces.integral('[v] n ({f} - .5 C [u] n) dS' @ ns, degree=degree*2) sqr = domain.integral('(u - uinit)^2 dV' @ ns, degree=max(degree*2, 5)) - args = solver.optimize('u,', sqr) + args = System(sqr, trial='u').solve() + args['t'] = 0. + + system = System(res, trial='u', test='v') bezier = domain.sample('bezier', 7) - with log.iter.plain('timestep', itertools.count(step=timestep)) as times: - for t in times: - log.info('time:', round(t, 10)) + with log.iter.plain('timestep', itertools.count()) as steps: + for _ in steps: + log.info('time:', round(args['t'], 10)) x, u = bezier.eval(['x', 'u'] @ ns, **args) export.triplot('solution.png', x[:,numpy.newaxis], u, tri=bezier.tri, hull=bezier.hull, clim=(0, 1)) - if t >= endtime: + if args['t'] >= endtime: break - args['u0'] = args['u'] - args = solver.newton('u:v', res, arguments=args).solve(newtontol) + args = system.step(timestep=timestep, arguments=args, timetarget='t', historysuffix='0', tol=newtontol) return args diff --git a/examples/cahnhilliard.py b/examples/cahnhilliard.py index b1d5fb83e..4f21df762 100644 --- a/examples/cahnhilliard.py +++ b/examples/cahnhilliard.py @@ -1,4 +1,5 @@ -from nutils import mesh, function, solver, numeric, export, testing +from nutils import mesh, function, numeric, export, testing +from nutils.solver import System from nutils.expression_v2 import Namespace from nutils.SI import Length, Time, Density, Tension, Energy, Pressure, Velocity, parse import numpy @@ -178,6 +179,9 @@ def main(size: Length = parse('10cm'), ns.define_for('x', gradient='∇', normal='n', jacobians=('dV', 'dS')) ns.add_field(('φ', 'φ0'), basis) ns.add_field('η', basis * stens / epsilon) # basis scaling to give η the required unit + ns.add_field(('t', 't0')) + ns.t *= timestep + ns.t0 *= timestep ns.ε = epsilon ns.σ = stens ns.σmean = (wtensp + wtensn) / 2 @@ -188,7 +192,7 @@ def main(size: Length = parse('10cm'), ns.δψ = stab.value ns.M = mobility ns.J_i = '-M ∇_i(η)' - ns.dt = timestep + ns.dt = 't - t0' nrg_mix = domain.integral('(ψ σ / ε) dV' @ ns, degree=degree*4) nrg_iface = domain.integral('.5 σ ε ∇_k(φ) ∇_k(φ) dV' @ ns, degree=degree*4) @@ -198,14 +202,15 @@ def main(size: Length = parse('10cm'), numpy.random.seed(seed) args = dict(φ=numpy.random.normal(0, .5, basis.shape)) # initial condition + system = System(nrg / tol, trial='φ,η') + with log.iter.fraction('timestep', range(round(endtime / timestep))) as steps: for istep in steps: E = numpy.stack(function.eval([nrg_mix, nrg_iface, nrg_wall], **args)) log.user('energy: {0:,.0μJ/m} ({1[0]:.0f}% mixture, {1[1]:.0f}% interface, {1[2]:.0f}% wall)'.format(numpy.sum(E), 100*E/numpy.sum(E))) - args['φ0'] = args['φ'] - args = solver.optimize(['φ', 'η'], nrg / tol, arguments=args, tol=1) + args = system.step(timestep=1., timetarget='t', historysuffix='0', arguments=args, tol=1, maxiter=5) with export.mplfigure('phase.png') as fig: ax = fig.add_subplot(aspect='equal', xlabel='[mm]', ylabel='[mm]') diff --git a/examples/coil.py b/examples/coil.py index c92166aec..07ae8ec94 100644 --- a/examples/coil.py +++ b/examples/coil.py @@ -1,4 +1,5 @@ -from nutils import export, function, mesh, solver, testing +from nutils import export, function, mesh, testing +from nutils.solver import System from nutils.expression_v2 import Namespace import functools import numpy @@ -115,7 +116,7 @@ def main(nelems: int = 50, res = REV.integral(RZ.integral('-∇_j(Atest_i) ∇_j(A_i) dV' @ ns, degree=2*degree), degree=0) res += REV.integral(RZ['coil'].integral('μ0 Atest_i J_i dV' @ ns, degree=2*degree), degree=0) - args = solver.solve_linear('A:Atest,', res) + args = System(res, trial='A', test='Atest').solve() # Since the coordinate transformation is singular at r=0 we can't evaluate # `B` (the curl of `A`) at r=0. We circumvent this problem by projecting `B` @@ -124,7 +125,7 @@ def main(nelems: int = 50, ns.Borig = ns.B ns.add_field(('B', 'Btest'), RZ.basis('spline', degree=degree), ns.rot, dtype=complex) res = REV.integral(RZ.integral('Btest_i (B_i - Borig_i) dV' @ ns, degree=2*degree), degree=0) - args = solver.solve_linear('B:Btest,', res, arguments=args) + args = System(res, trial='B', test='Btest').solve(arguments=args) with export.mplfigure('magnetic-potential-1.png', dpi=300) as fig: ax = fig.add_subplot(111, aspect='equal', xlabel='$x_0$', ylabel='$x_2$', adjustable='datalim') diff --git a/examples/cylinderflow.py b/examples/cylinderflow.py index 247c34cc8..0821a9e7d 100644 --- a/examples/cylinderflow.py +++ b/examples/cylinderflow.py @@ -1,4 +1,5 @@ -from nutils import mesh, function, solver, export, testing, numeric +from nutils import mesh, function, export, testing, numeric +from nutils.solver import System from nutils.expression_v2 import Namespace import itertools import numpy @@ -118,8 +119,8 @@ def main(nelems: int = 99, domain.basis('spline', degree=(degree, degree-1), removedofs=((0,), None)), domain.basis('spline', degree=(degree-1, degree))]) @ J.T / detJ) ns.add_field(('p', 'q'), domain.basis('spline', degree=degree-1) / detJ) - ns.dt = timestep - ns.DuDt_i = '(u_i - u0_i) / dt + ∇_j(u_i) u_j' # material derivative + ns.add_field(('t', 't0')) + ns.DuDt_i = '(u_i - u0_i) / (t - t0) + ∇_j(u_i) u_j' # material derivative ns.σ_ij = '(∇_j(u_i) + ∇_i(u_j)) / Re - p δ_ij' ns.ω = 'ε_ij ∇_j(u_i)' ns.N = 10 * degree / elemangle # Nitsche constant based on element size = elemangle/2 @@ -128,10 +129,10 @@ 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 = solver.optimize('u,', sqr, droptol=1e-15) # constrain inflow boundary to unit horizontal flow + cons = System(sqr, trial='u').optimize(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 = solver.optimize('u,p', sqr, constrain=cons) # set initial condition to potential flow + args = System(sqr, trial='u,p').solve(constrain=cons) # set initial condition to potential flow res = domain.integral('(v_i DuDt_i + ∇_j(v_i) σ_ij + q ∇_k(u_k)) dV' @ ns, degree=degree*3) res += domain.boundary['inner'].integral('(nitsche_i (u_i - uwall_i) - v_i σ_ij n_j) dS' @ ns, degree=degree*2) @@ -142,10 +143,11 @@ def main(nelems: int = 99, steps = treelog.iter.fraction('timestep', range(round(endtime / timestep))) if endtime < float('inf') \ else treelog.iter.plain('timestep', itertools.count()) + system = System(res, trial='u,p', test='v,q') + for _ in steps: treelog.info(f'velocity divergence: {div.eval(**args):.0e}') - args['u0'] = args['u'] - args = solver.newton('u:v,p:q', residual=res, arguments=args, constrain=cons).solve(1e-10) + args = system.step(timestep=timestep, timetarget='t', historysuffix='0', arguments=args, constrain=cons, tol=1e-10) postprocess(args) return args, numpy.sqrt(domain.integral('∇_k(u_k)^2 dV' @ ns, degree=2)) @@ -161,7 +163,6 @@ def __init__(self, topo, ns, region=4., aspect=16/9, figscale=7.2, spacing=.05, self.spacing = spacing self.pstep = pstep self.vortlim = vortlim - self.t = 0. self.initialize_xgrd() self.topo = topo @@ -199,10 +200,8 @@ def __call__(self, args): ax.tricontour(*x.T, self.bezier.tri, p, numpy.arange(numpy.min(p)//1, numpy.max(p), self.pstep), colors='k', linestyles='solid', linewidths=.5, zorder=9) export.plotlines_(ax, x.T, self.bezier.hull, colors='k', linewidths=.1, alpha=.5) ax.quiver(*self.xgrd.T, *ugrd.T, angles='xy', width=1e-3, headwidth=3e3, headlength=5e3, headaxislength=2e3, zorder=9, alpha=.5, pivot='tip') - ax.plot(0, 0, 'k', marker=(3, 2, -self.t*self.ns.rotation.eval()*180/numpy.pi-90), markersize=20) - dt = self.ns.dt.eval() - self.t += dt - self.xgrd += ugrd * dt + ax.plot(0, 0, 'k', marker=(3, 2, -args['t']*self.ns.rotation.eval()*180/numpy.pi-90), markersize=20) + self.xgrd += ugrd * (args['t'] - args['t0']) self.regularize_xgrd() diff --git a/examples/drivencavity.py b/examples/drivencavity.py index efa4c0984..0009bbf22 100644 --- a/examples/drivencavity.py +++ b/examples/drivencavity.py @@ -1,4 +1,5 @@ -from nutils import mesh, function, solver, export, testing +from nutils import mesh, function, export, testing +from nutils.solver import System, LinesearchNewton from nutils.expression_v2 import Namespace import treelog as log import numpy @@ -117,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 = solver.optimize('u,', sqr, droptol=1e-15) + cons = System(sqr, trial='u').optimize(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 = solver.optimize('u,', sqr, droptol=1e-15) + tcons = System(sqr, trial='u').optimize(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 @@ -131,7 +132,7 @@ def main(nelems: int = 32, res += domain.boundary.integral('(nitsche_i (u_i - uwall_i) - v_i σ_ij n_j) dS' @ ns, degree=2*degree) with log.context('stokes'): - args0 = solver.solve_linear('u:v,p:q', res, constrain=cons) + args0 = System(res, trial='u,p', test='v,q').solve(constrain=cons) postprocess(domain, ns, **args0) # change to Navier-Stokes by adding convection @@ -141,7 +142,7 @@ def main(nelems: int = 32, res += domain.integral('.5 u_i v_i ∇_j(u_j) dV' @ ns, degree=degree*3) with log.context('navier-stokes'): - args1 = solver.newton('u:v,p:q', res, arguments=args0, constrain=cons).solve(tol=1e-10) + args1 = System(res, trial='u,p', test='v,q').solve(arguments=args0, constrain=cons, tol=1e-10, method=LinesearchNewton()) postprocess(domain, ns, **args1) return args0, args1 @@ -155,7 +156,7 @@ def postprocess(domain, ns, **arguments): # reconstruct velocity streamlines sqr = domain.integral('Σ_i (u_i - ε_ij ∇_j(ψ))^2 dV' @ ns, degree=4) - arguments = solver.optimize('ψ,', sqr, arguments=arguments) + arguments = System(sqr, trial='ψ').solve(arguments=arguments) bezier = domain.sample('bezier', 9) x, u, p, ψ = bezier.eval(['x_i', 'sqrt(u_i u_i)', 'p', 'ψ'] @ ns, **arguments) diff --git a/examples/elasticity.py b/examples/elasticity.py index 5b357c4dd..0091e42e9 100644 --- a/examples/elasticity.py +++ b/examples/elasticity.py @@ -1,4 +1,5 @@ -from nutils import mesh, function, solver, export, testing +from nutils import mesh, function, export, testing +from nutils.solver import System from nutils.expression_v2 import Namespace import treelog as log import numpy @@ -50,20 +51,18 @@ def main(nelems: int = 24, ns.q_i = '-δ_i1' sqr = domain.boundary['top'].integral('u_k u_k dS' @ ns, degree=degree*2) - cons = solver.optimize(('u',), sqr, droptol=1e-15) + cons = System(sqr, trial='u').optimize(droptol=1e-15) # solve for equilibrium configuration - internal = domain.integral('E dV' @ ns, degree=degree*2) - external = domain.integral('u_i q_i dV' @ ns, degree=degree*2) - args = solver.optimize('u,', internal - external, constrain=cons) + energy = domain.integral('(E - u_i q_i) dV' @ ns, degree=degree*2) + args = System(energy, trial='u').solve(constrain=cons) # evaluate tractions and net force if direct: ns.t_i = 'σ_ij n_j' # <-- this is an inadmissible boundary term else: - external += domain.boundary['top'].integral('u_i t_i dS' @ ns, degree=degree*2) - invcons = dict(t=numpy.choose(numpy.isnan(cons['u']), [numpy.nan, 0.])) - args = solver.solve_linear(('t',), [(internal - external).derivative('u')], constrain=invcons, arguments=args) + energy -= domain.boundary['top'].integral('u_i t_i dS' @ ns, degree=degree*2) + args = System(energy, trial='t', test='u').solve(constrain={'t': numpy.isnan(cons['u'])}, arguments=args) F = domain.boundary['top'].integrate('t_i dS' @ ns, degree=degree*2, arguments=args) log.user('total clamping force:', F) @@ -132,8 +131,8 @@ def test_poisson(self): eNpjaGBAhSBAZTEAEKAUAQ==''') with self.subTest('solution'): self.assertAlmostEqual64(args['u'], ''' - eNqTNig6vcVwwekjRuJn5Iy1zzIAwQs999MdBmWn+w0Zz7QYpoPF/v+3Pv3/f/5pBgbWMwwM8WAxiYvu - pyvOl50uPMd4puYcRN3T80Wnfc4tOG1zVvzMozMQ8wAOOSng''') + eNqTNig6vcVwwekjRuJn5Iy1zzIAwQs999MdBmWn+w0Zz7QYpoPFGBisTzMw5AMx6xkGhniwmMRF99MV + 58tOF55jPFNzDqLu6fmi0z7nFpy2OSt+5tEZiHkAKRAl5A==''') with self.subTest('traction'): self.assertAlmostEqual64(args['t'], ''' eNpjYEAF/Sc+maMJMdw0emzGgAFiMdSpn8VUV2j+yRwAoCAJFw==''') diff --git a/examples/finitestrain.py b/examples/finitestrain.py index 9567c5eda..7cf4de8f0 100644 --- a/examples/finitestrain.py +++ b/examples/finitestrain.py @@ -1,4 +1,5 @@ -from nutils import mesh, function, solver, export, testing +from nutils import mesh, function, export, testing +from nutils.solver import System, Minimize from nutils.expression_v2 import Namespace import numpy @@ -57,10 +58,10 @@ 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 = solver.optimize('u,', sqr, droptol=1e-15) + cons = System(sqr, trial='u').optimize(droptol=1e-15) energy = domain.integral('energy dV' @ ns, degree=degree*2) - args0 = solver.optimize('u,', energy, constrain=cons) + args0 = System(energy, trial='u').solve(constrain=cons) x, energy = bezier.eval(['x_i', 'energy'] @ ns, **args0) export.triplot('linear.png', x, energy, tri=bezier.tri, hull=bezier.hull, cmap='jet') @@ -68,7 +69,7 @@ def main(nelems: int = 20, ns.energy = 'λ ε_ii ε_jj + 2 μ ε_ij ε_ij' energy = domain.integral('energy dV' @ ns, degree=degree*2) - args1 = solver.minimize('u,', energy, arguments=args0, constrain=cons).solve(restol) + args1 = System(energy, trial='u').solve(arguments=args0, constrain=cons, method=Minimize(), tol=restol) x, energy = bezier.eval(['x_i', 'energy'] @ ns, **args1) export.triplot('nonlinear.png', x, energy, tri=bezier.tri, hull=bezier.hull, cmap='jet') diff --git a/examples/laplace.py b/examples/laplace.py index 08b919343..d9d75284f 100644 --- a/examples/laplace.py +++ b/examples/laplace.py @@ -1,4 +1,5 @@ -from nutils import mesh, function, solver, export, testing +from nutils import mesh, function, export, testing +from nutils.solver import System from nutils.expression_v2 import Namespace import treelog @@ -71,13 +72,13 @@ 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 = solver.optimize('u,', sqr, droptol=1e-15) + cons = System(sqr, trial='u').optimize(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` # matches `cons` in the constrained entries. - args = solver.solve_linear('u:v', res, constrain=cons) + args = System(res, trial='u', test='v').solve(constrain=cons) # Once all arguments are establised, the corresponding solution can be # vizualised by sampling values of `ns.u` along with physical coordinates diff --git a/examples/platewithhole.py b/examples/platewithhole.py index 2d0973dfd..088b61ef2 100644 --- a/examples/platewithhole.py +++ b/examples/platewithhole.py @@ -1,4 +1,5 @@ -from nutils import mesh, function, solver, export, testing +from nutils import mesh, function, export, testing +from nutils.solver import System from nutils.expression_v2 import Namespace from dataclasses import dataclass from typing import Union @@ -75,7 +76,7 @@ def generate(self, radius): topo = topo.refine(self.nrefine) bsplinebasis = topo.basis('spline', degree=2) sqr = topo.integral((function.dotarg('w', bsplinebasis) - weightfunc)**2, degree=9) - controlweights = solver.optimize('w', sqr) + controlweights = System(sqr, trial='w').solve()['w'] nurbsbasis = bsplinebasis * controlweights / weightfunc return topo.withboundary(hole='left', sym='top,bottom', far='right'), geom, nurbsbasis, 5 @@ -136,13 +137,13 @@ 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 = solver.optimize(('u',), sqr, droptol=1e-15) + cons = System(sqr, trial='u').optimize(droptol=1e-15) sqr = topo.boundary['far'].integral('du_k du_k dS' @ ns, degree=20) - cons = solver.optimize(('u',), sqr, droptol=1e-15, constrain=cons) + cons = System(sqr, trial='u').optimize(droptol=1e-15, constrain=cons) res = topo.integral('∇_j(v_i) σ_ij dV' @ ns, degree=degree*2) - args = solver.solve_linear('u:v', res, constrain=cons) + args = System(res, trial='u', test='v').solve(constrain=cons) bezier = topo.sample('bezier', 5) X, σxx = bezier.eval(['X_i', 'σ_00'] @ ns, **args) diff --git a/examples/poisson.py b/examples/poisson.py index eca7f984f..c88c8d188 100644 --- a/examples/poisson.py +++ b/examples/poisson.py @@ -1,4 +1,5 @@ -from nutils import mesh, function, solver, export, testing +from nutils import mesh, function, export, testing +from nutils.solver import System # This script demonstrates direct function manipulation without the aid of # namespace expressions. @@ -24,10 +25,10 @@ def main(nelems: int = 32): J = function.J(x) sqr = topo.boundary.integral(u**2 * J, degree=2) - cons = solver.optimize(('u',), sqr, droptol=1e-12) + cons = System(sqr, trial='u').optimize(droptol=1e-12) energy = topo.integral((g @ g / 2 - u) * J, degree=1) - args = solver.optimize(('u',), energy, constrain=cons) + args = System(energy, trial='u').solve(constrain=cons) bezier = topo.sample('bezier', 3) x, u = bezier.eval([x, u], **args) diff --git a/examples/torsion.py b/examples/torsion.py index 77d030f79..9d0a84d9a 100644 --- a/examples/torsion.py +++ b/examples/torsion.py @@ -1,4 +1,5 @@ -from nutils import mesh, function, solver, export, testing +from nutils import mesh, function, export, testing +from nutils.solver import System, Minimize from nutils.expression_v2 import Namespace import treelog as log import numpy as np @@ -85,12 +86,13 @@ def main(length: float = 2*np.pi, ns.W = 'F_ij F_ij - 3 - 2 log(J) + D (J - 1)^2' # Neo-Hookean energy density energy = topo.integral('W dV' @ ns, degree=degree*2) + system = System(energy, trial='u') args = {} clim = (0, 1) if stretch == 1 else None # fix scale for undeformed configuration for args['φ'] in np.linspace(0, rotation, round(rotation / increment) + 1): with log.context('{φ:.1f} deg', **args): - args = solver.minimize('u,', energy, arguments=args).solve(restol) + args = system.solve(arguments=args, method=Minimize(), tol=restol) x, W = bezier.eval(['x_i', 'W'] @ ns, **args) export.triplot('energy.jpg', x, W, tri=bezier.tri, hull=bezier.hull, clim=clim, cmap='inferno_r', vlabel='strain energy density') diff --git a/nutils/__init__.py b/nutils/__init__.py index bfd8fd76d..27809816f 100644 --- a/nutils/__init__.py +++ b/nutils/__init__.py @@ -1,4 +1,4 @@ 'Numerical Utilities for Finite Element Analysis' -__version__ = version = '9a36' +__version__ = version = '9a37' version_name = 'jook-sing' diff --git a/nutils/evaluable.py b/nutils/evaluable.py index b6c590212..8d4355173 100644 --- a/nutils/evaluable.py +++ b/nutils/evaluable.py @@ -85,49 +85,29 @@ def _isindex(arg): return isinstance(arg, Array) and arg.ndim == 0 and arg.dtype == int and arg._intbounds[0] >= 0 -def _equals_scalar_constant(arg: 'Array', value: Dtype): - assert isinstance(arg, Array) and arg.ndim == 0, f'arg={arg!r}' - assert arg.dtype == type(value), f'arg.dtype={arg.dtype}, type(value)={type(value)}' - return arg.isconstant and arg.eval() == value - - def _equals_simplified(arg1: 'Array', arg2: 'Array'): + 'return True if certainly equal, False if certainly not equal, None otherwise' + assert isinstance(arg1, Array), f'arg1={arg1!r}' assert isinstance(arg2, Array), f'arg2={arg2!r}' if arg1 is arg2: return True - assert equalshape(arg1.shape, arg2.shape) and arg1.dtype == arg2.dtype, f'arg1={arg1!r}, arg2={arg2!r}' + if arg1.dtype != arg2.dtype or arg1.ndim != arg2.ndim: + return False arg1 = arg1.simplified arg2 = arg2.simplified if arg1 is arg2: return True if arg1.arguments != arg2.arguments: return False - if arg1.isconstant: # implies arg2.isconstant - return numpy.all(arg1.eval() == arg2.eval()) - + if isinstance(arg1, Constant) and isinstance(arg2, Constant): + return False # values differ -def equalshape(N: typing.Tuple['Array', ...], M: typing.Tuple['Array', ...]): - '''Compare two array shapes. - Returns `True` if all indices are certainly equal, `False` if any indices are - certainly not equal, or `None` if equality cannot be determined at compile - time. - ''' - - assert isinstance(N, tuple) and all(_isindex(n) for n in N), f'N={N!r}' - assert isinstance(M, tuple) and all(_isindex(n) for n in M), f'M={M!r}' - if N == M: - return True - if len(N) != len(M): - return False - retval = True - for eq in map(_equals_simplified, N, M): - if eq == False: - return False - if eq == None: - retval = None - return retval +_certainly_different = lambda a, b: _equals_simplified(a, b) is False +_certainly_equal = lambda a, b: _equals_simplified(a, b) is True +_any_certainly_different = lambda a, b: any(map(_certainly_different, a, b)) +_all_certainly_equal = lambda a, b: all(map(_certainly_equal, a, b)) class ExpensiveEvaluationWarning(warnings.NutilsInefficiencyWarning): @@ -184,7 +164,7 @@ def simplified(obj): if retval is None: return obj if isinstance(obj, Array): - assert isinstance(retval, Array) and equalshape(retval.shape, obj.shape) and retval.dtype == obj.dtype, '{} --simplify--> {}'.format(obj, retval) + assert isinstance(retval, Array) and not _any_certainly_different(retval.shape, obj.shape) and retval.dtype == obj.dtype, '{} --simplify--> {}'.format(obj, retval) return retval def _simplified(self): @@ -200,7 +180,7 @@ def _optimized_for_numpy1(obj): if retval is None: return obj if isinstance(obj, Array): - assert isinstance(retval, Array) and equalshape(retval.shape, obj.shape), '{0}._optimized_for_numpy or {0}._simplified resulted in shape change'.format(type(obj).__name__) + assert isinstance(retval, Array) and not _any_certainly_different(retval.shape, obj.shape), '{0}._optimized_for_numpy or {0}._simplified resulted in shape change'.format(type(obj).__name__) return retval def _optimized_for_numpy(self): @@ -404,7 +384,7 @@ def align(arg, where, shape): where.append(i) if where != list(range(len(shape))): arg = Transpose(arg, util.untake(where)) - assert equalshape(arg.shape, shape), f'arg.shape={arg.shape!r}, shape={shape!r}' + assert not _any_certainly_different(arg.shape, shape), f'arg.shape={arg.shape!r}, shape={shape!r}' return arg @@ -463,7 +443,7 @@ def _assparse(self): for *indices, values in chunks: assert len(indices) == self.ndim assert all(idx.dtype == int for idx in indices) - assert all(equalshape(idx.shape, values.shape) for idx in indices) + assert not any(_any_certainly_different(idx.shape, values.shape) for idx in indices) elif chunks: assert len(chunks) == 1 chunk, = chunks @@ -595,6 +575,18 @@ def real(self): def imag(self): return imag(self) + @property + def assparse(self): + if not self.ndim: + return InsertAxis(self, constant(1)), (), () + sparse = self._assparse + if not sparse: + indices = [zeros((constant(0),), int)] * self.ndim + values = zeros((constant(0),), self.dtype) + else: + *indices, values = tuple(concatenate([_flat(array) for array in arrays]) for arrays in zip(*sparse)) + return values, tuple(indices), self.shape + @cached_property @verify_sparse_chunks def _assparse(self): @@ -671,17 +663,16 @@ def as_evaluable_array(self): @cached_property def _intbounds(self): # inclusive lower and upper bounds + lower, upper = self._intbounds_impl() + assert isinstance(lower, int) or lower == float('-inf') + assert isinstance(upper, int) or upper == float('inf') + assert lower <= upper + return lower, upper + + def _intbounds_impl(self): 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') - assert isinstance(upper, int) or upper == float('inf') - assert lower <= upper - return lower, upper - - def _intbounds_impl(self): return float('-inf'), float('inf') @property @@ -700,6 +691,46 @@ def _compile_with_out(self, builder, out, out_block_id, mode): return NotImplemented +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) + + @property + def dtype(self): + return self.a.dtype + + @cached_property + def shape(self): + return self.map(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 + + @property + def dependencies(self): + return self.a, self.b + + def evalf(self, a, b): + if a.shape != b.shape or (a != b).any(): + raise Exception('values are not equal') + return a + + class Orthonormal(Array): 'make a vector orthonormal to a subspace' @@ -711,7 +742,7 @@ class Orthonormal(Array): def __post_init__(self): assert isinstance(self.basis, Array) and self.basis.ndim >= 2 and self.basis.dtype not in (bool, complex), f'basis={self.basis!r}' assert isinstance(self.vector, Array) and self.vector.ndim >= 1 and self.vector.dtype not in (bool, complex), f'vector={self.vector!r}' - assert equalshape(self.basis.shape[:-1], self.vector.shape) + assert not _any_certainly_different(self.basis.shape[:-1], self.vector.shape) @property def dependencies(self): @@ -722,7 +753,7 @@ def shape(self): return self.vector.shape def _simplified(self): - if _equals_scalar_constant(self.shape[-1], 1): + if isunit(self.shape[-1]): return Sign(self.vector) basis, vector, where = unalign(self.basis, self.vector, naxes=self.ndim - 1) if len(where) < self.ndim - 1: @@ -737,7 +768,7 @@ def evalf(G, n): return numeric.normalize(n - v3) def _derivative(self, var, seen): - if _equals_scalar_constant(self.shape[-1], 1): + if isunit(self.shape[-1]): return zeros(self.shape + var.shape) # definitions: @@ -767,7 +798,7 @@ def _derivative(self, var, seen): Q = -einsum('Aim,Amn,AjnB->AijB', G, invGG, derivative(G, var, seen)) QN = einsum('Ai,AjiB->AjB', self, Q) - if _equals_simplified(G.shape[-1], G.shape[-2] - 1): # dim(kern(G)) = 1 + if _certainly_equal(G.shape[-1], G.shape[-2] - 1): # dim(kern(G)) = 1 # In this situation, since N is a basis for the kernel of G, we # have the identity P == N N^T which cancels the entire second term # of N' along with any reference to v', reducing it to N' = Q N. @@ -873,9 +904,8 @@ def _takediag(self, axis1, axis2): list(range(axis1)) + list(range(axis1+1, axis2)) + list(range(axis2+1, self.ndim)) + [axis1, axis2]))) def _take(self, index, axis): - if index.isconstant: - index_ = index.eval() - return constant(self.value.take(index_, axis)) + if isinstance(index, Constant): + return constant(self.value.take(index.value, axis)) def _power(self, n): if isinstance(n, Constant): @@ -945,7 +975,7 @@ def _unaligned(self): return self.func._unaligned def _simplified(self): - if _equals_scalar_constant(self.length, 0): + if iszero(self.length): return zeros_like(self) return self.func._insertaxis(self.ndim-1, self.length) @@ -959,7 +989,7 @@ def evalf(func, length): return numpy.repeat(func[..., numpy.newaxis], length, -1) def _compile_expression(self, py_self, func, length): - if _equals_scalar_constant(self.length, 1): + if isunit(self.length): return func.get_item(_pyast.Tuple((_pyast.Raw('...'), _pyast.Variable('numpy').get_attr('newaxis')))) else: return super()._compile_expression(py_self, func, length) @@ -1245,6 +1275,14 @@ def _intbounds_impl(self): def _const_uniform(self): return self.func._const_uniform + def _optimized_for_numpy(self): + if isinstance(self.func, Transpose): + return transpose(self.func.func, [self.func.axes[i] for i in self.axes]) + if isinstance(self.func, Assemble): + offsets = numpy.cumsum([0] + [index.ndim for index in self.func.indices]) + axes = [n for axis in self.axes for n in range(offsets[axis], offsets[axis+1])] + return Assemble(transpose(self.func.func, axes), tuple(self.func.indices[i] for i in self.axes), self.shape) + class Product(Array): @@ -1269,7 +1307,7 @@ def _compile_expression(self, py_self, func): return _pyast.Variable('numpy').get_attr('all' if self.dtype == bool else 'prod').call(func, axis=_pyast.LiteralInt(-1)) def _simplified(self): - if _equals_scalar_constant(self.func.shape[-1], 1): + if isunit(self.func.shape[-1]): return get(self.func, self.ndim, constant(0)) return self.func._product() @@ -1294,7 +1332,7 @@ class Inverse(Array): func: Array def __post_init__(self): - assert isinstance(self.func, Array) and self.func.dtype in (float, complex) and self.func.ndim >= 2 and _equals_simplified(self.func.shape[-1], self.func.shape[-2]), f'func={self.func!r}' + assert isinstance(self.func, Array) and self.func.dtype in (float, complex) and self.func.ndim >= 2 and not _certainly_different(self.func.shape[-1], self.func.shape[-2]), f'func={self.func!r}' @property def dependencies(self): @@ -1309,9 +1347,9 @@ def shape(self): return self.func.shape def _simplified(self): - if _equals_scalar_constant(self.func.shape[-1], 1): + if isunit(self.func.shape[-1]): return reciprocal(self.func) - if _equals_scalar_constant(self.func.shape[-1], 0): + if iszero(self.func.shape[-1]): return singular_like(self) result = self.func._inverse(self.ndim-2, self.ndim-1) if result is not None: @@ -1349,7 +1387,7 @@ class Determinant(Array): func: Array def __post_init__(self): - assert isarray(self.func) and self.func.dtype in (float, complex) and self.func.ndim >= 2 and _equals_simplified(self.func.shape[-1], self.func.shape[-2]) + assert isarray(self.func) and self.func.dtype in (float, complex) and self.func.ndim >= 2 and not _certainly_different(self.func.shape[-1], self.func.shape[-2]) @property def dependencies(self): @@ -1367,7 +1405,7 @@ def _simplified(self): result = self.func._determinant(self.ndim, self.ndim+1) if result is not None: return result - if _equals_scalar_constant(self.func.shape[-1], 1): + if isunit(self.func.shape[-1]): return Take(Take(self.func, zeros((), int)), zeros((), int)) evalf = staticmethod(numpy.linalg.det) @@ -1389,7 +1427,7 @@ class Multiply(Array): def __post_init__(self): assert isinstance(self.funcs, types.frozenmultiset), f'funcs={self.funcs!r}' func1, func2 = self.funcs - assert equalshape(func1.shape, func2.shape) and func1.dtype == func2.dtype, 'Multiply({}, {})'.format(func1, func2) + assert not _any_certainly_different(func1.shape, func2.shape) and func1.dtype == func2.dtype, 'Multiply({}, {})'.format(func1, func2) @property def dependencies(self): @@ -1404,7 +1442,7 @@ def dtype(self): @cached_property def shape(self): func1, func2 = self.funcs - return func1.shape + return AssertEqual.map(func1.shape, func2.shape) @property def _factors(self): @@ -1485,7 +1523,7 @@ def _add(self, other): def _determinant(self, axis1, axis2): axis1, axis2 = sorted([axis1, axis2]) factors = tuple(self._factors) - if all(_equals_scalar_constant(self.shape[axis], 1) for axis in (axis1, axis2)): + if all(isunit(self.shape[axis]) for axis in (axis1, axis2)): return multiply(*[determinant(f, (axis1, axis2)) for f in factors]) for i, fi in enumerate(factors): unaligned, where = unalign(fi) @@ -1576,7 +1614,7 @@ class Add(Array): def __post_init__(self): assert isinstance(self.funcs, types.frozenmultiset) and len(self.funcs) == 2, f'funcs={self.funcs!r}' func1, func2 = self.funcs - assert equalshape(func1.shape, func2.shape) and func1.dtype == func2.dtype, 'Add({}, {})'.format(func1, func2) + assert not _any_certainly_different(func1.shape, func2.shape) and func1.dtype == func2.dtype, 'Add({}, {})'.format(func1, func2) @property def dependencies(self): @@ -1591,7 +1629,7 @@ def dtype(self): @cached_property def shape(self): func1, func2 = self.funcs - return func1.shape + return AssertEqual.map(func1.shape, func2.shape) @cached_property def _inflations(self): @@ -1606,10 +1644,10 @@ def _inflations(self): parts2 = func2_inflations[axis] dofmaps = set(parts1) | set(parts2) if (len(parts1) < len(dofmaps) and len(parts2) < len(dofmaps) # neither set is a subset of the other; total may be dense - and self.shape[axis].isconstant and all(dofmap.isconstant for dofmap in dofmaps)): - mask = numpy.zeros(int(self.shape[axis]), dtype=bool) + and isinstance(self.shape[axis], Constant) and all(isinstance(dofmap, Constant) for dofmap in dofmaps)): + mask = numpy.zeros(self.shape[axis].value, dtype=bool) for dofmap in dofmaps: - mask[dofmap.eval()] = True + mask[dofmap.value] = True if mask.all(): # axis adds up to dense continue inflations.append((axis, types.frozendict((dofmap, util.sum(parts[dofmap] for parts in (parts1, parts2) if dofmap in parts)) for dofmap in dofmaps))) @@ -1736,13 +1774,12 @@ def __post_init__(self): lengths = {} for idx, arg in zip(self.args_idx, self.args): for i, length in zip(idx, arg.shape): - if i not in lengths: - lengths[i] = length - elif not _equals_simplified(lengths[i], length): - raise ValueError('Axes with index {} have different lengths.'.format(i)) - for i in self.out_idx: - if i not in lengths: - raise ValueError(f'Output axis {i} is not listed in any of the arguments.') + n = lengths.get(i) + lengths[i] = length if n is None else AssertEqual(length, n) + try: + self.shape = tuple(lengths[i] for i in self.out_idx) + except KeyError(e): + raise ValueError(f'Output axis {e} is not listed in any of the arguments.') @cached_property def _einsumfmt(self): @@ -1756,11 +1793,6 @@ def dependencies(self): def dtype(self): return self.args[0].dtype - @cached_property - def shape(self): - lengths = {i: length for idx, arg in zip(self.args_idx, self.args) for i, length in zip(idx, arg.shape)} - return tuple(lengths[i] for i in self.out_idx) - def _compile_expression(self, py_self, *args): return _pyast.Variable('numpy').get_attr('core').get_attr('multiarray').get_attr('c_einsum').call(_pyast.LiteralStr(self._einsumfmt), *args) @@ -1806,7 +1838,7 @@ def _compile_expression(self, py_self, func): return _pyast.Variable('numpy').get_attr('any' if self.dtype == bool else 'sum').call(func, axis=_pyast.LiteralInt(-1)) def _simplified(self): - if _equals_scalar_constant(self.func.shape[-1], 1): + if isunit(self.func.shape[-1]): return Take(self.func, constant(0)) return self.func._sum(self.ndim) @@ -1869,7 +1901,7 @@ class TakeDiag(Array): func: Array def __post_init__(self): - assert isinstance(self.func, Array) and self.func.ndim >= 2 and _equals_simplified(*self.func.shape[-2:]), f'func={self.func!r}' + assert isinstance(self.func, Array) and self.func.ndim >= 2 and not _certainly_different(*self.func.shape[-2:]), f'func={self.func!r}' @property def dependencies(self): @@ -1884,7 +1916,7 @@ def shape(self): return self.func.shape[:-1] def _simplified(self): - if _equals_scalar_constant(self.shape[-1], 1): + if isunit(self.shape[-1]): return Take(self.func, constant(0)) return self.func._takediag(self.ndim-1, self.ndim) @@ -1991,7 +2023,7 @@ class Power(Array): def __post_init__(self): assert isinstance(self.func, Array) and self.func.dtype != bool, f'func={self.func!r}' assert isinstance(self.power, Array) and (self.power.dtype in (float, complex) or self.power.dtype == int and self.power._intbounds[0] >= 0), f'power={self.power!r}' - assert equalshape(self.func.shape, self.power.shape) and self.func.dtype == self.power.dtype, 'Power({}, {})'.format(self.func, self.power) + assert not _any_certainly_different(self.func.shape, self.power.shape) and self.func.dtype == self.power.dtype, 'Power({}, {})'.format(self.func, self.power) @property def dependencies(self): @@ -2027,11 +2059,10 @@ def _compile_expression(self, py_self, func, power): return _pyast.Variable('numpy').get_attr('power').call(func, power) def _derivative(self, var, seen): - if self.power.isconstant: - p = self.power.eval() - return einsum('A,A,AB->AB', constant(p), power(self.func, p - (p != 0)), derivative(self.func, var, seen)) - if self.dtype == complex: - raise NotImplementedError('The complex derivative is not implemented.') + if isinstance(self.power, Constant): + p = self.power.value.copy() + p -= p != 0 # exclude zero powers from decrement to avoid potential division by zero errors + return einsum('A,A,AB->AB', self.power, power(self.func, p), derivative(self.func, var, seen)) # self = func**power # ln self = power * ln func # self` / self = power` * ln func + power * func` / func @@ -2068,12 +2099,9 @@ class Pointwise(Array): dependencies = util.abstract_property() def __post_init__(self): - shape0 = self.dependencies[0].shape - assert all(equalshape(arg.shape, shape0) for arg in self.dependencies[1:]), 'pointwise arguments have inconsistent shapes' - - @cached_property - def shape(self): - return self.dependencies[0].shape + self.shape = self.dependencies[0].shape + for dep in self.dependencies[1:]: + self.shape = AssertEqual.map(self.shape, dep.shape) def _newargs(self, *args): ''' @@ -2896,7 +2924,7 @@ class Eig(Evaluable): symmetric: bool = False def __post_init__(self): - assert isinstance(self.func, Array) and self.func.ndim >= 2 and _equals_simplified(self.func.shape[-1], self.func.shape[-2]), f'func={self.func!r}' + assert isinstance(self.func, Array) and self.func.ndim >= 2 and not _certainly_different(self.func.shape[-1], self.func.shape[-2]), f'func={self.func!r}' assert isinstance(self.symmetric, bool), f'symmetric={self.symmetric!r}' @property @@ -3083,7 +3111,7 @@ def __post_init__(self): assert isinstance(self.func, Array), f'func={self.func!r}' assert isinstance(self.dofmap, Array), f'dofmap={self.dofmap!r}' assert _isindex(self.length), f'length={self.length!r}' - assert equalshape(self.func.shape[self.func.ndim-self.dofmap.ndim:], self.dofmap.shape), f'func.shape={self.func.shape!r}, dofmap.shape={self.dofmap.shape!r}' + assert not _any_certainly_different(self.func.shape[self.func.ndim-self.dofmap.ndim:], self.dofmap.shape), f'func.shape={self.func.shape!r}, dofmap.shape={self.dofmap.shape!r}' @property def dependencies(self): @@ -3110,25 +3138,28 @@ def _inflations(self): def _simplified(self): for axis in range(self.dofmap.ndim): - if _equals_scalar_constant(self.dofmap.shape[axis], 1): + if isunit(self.dofmap.shape[axis]): return Inflate(_take(self.func, constant(0), self.func.ndim-self.dofmap.ndim+axis), _take(self.dofmap, constant(0), axis), self.length) for axis, parts in self.func._inflations: i = axis - (self.ndim-1) if i >= 0: return util.sum(Inflate(f, _take(self.dofmap, ind, i), self.length) for ind, f in parts.items()) - if self.dofmap.ndim == 0 and _equals_scalar_constant(self.dofmap, 0) and _equals_scalar_constant(self.length, 1): + if self.dofmap.ndim == 0 and iszero(self.dofmap) and isunit(self.length): return InsertAxis(self.func, constant(1)) return self.func._inflate(self.dofmap, self.length, self.ndim-1) \ or self.dofmap._rinflate(self.func, self.length, self.ndim-1) + def _optimized_for_numpy(self): + indices = [Range(n) for n in self.shape[:-1]] + [self.dofmap] + return Assemble(self.func, tuple(indices), self.shape) + def evalf(self, array, indices, length): assert indices.ndim == self.dofmap.ndim assert length.ndim == 0 if not self.dofmap.isconstant and int(length) > indices.size: warnings.warn('using explicit inflation; this is usually a bug.', ExpensiveEvaluationWarning) - inflated = numpy.zeros(array.shape[:array.ndim-indices.ndim] + (length,), dtype=self.dtype) - numpy.add.at(inflated, (slice(None),)*(self.ndim-1)+(indices,), array) - return inflated + shape = *array.shape[:array.ndim-indices.ndim], length + return numeric.accumulate(array, (slice(None),)*(self.ndim-1)+(indices,), shape) def _compile_with_out(self, builder, out, out_block_id, mode): assert mode in ('iadd', 'assign') @@ -3198,7 +3229,7 @@ def _unravel(self, axis, shape): return Inflate(unravel(self.func, axis, shape), self.dofmap, self.length) def _sign(self): - if self.dofmap.isconstant and _isunique(self.dofmap.eval()): + if isinstance(self.dofmap, Constant) and _isunique(self.dofmap.value): return Inflate(Sign(self.func), self.dofmap, self.length) @cached_property @@ -3231,8 +3262,8 @@ def dependencies(self): return self.inflateidx, self.takeidx def _simplified(self): - if self.isconstant: - return Tuple(tuple(map(constant, self.eval()))) + if isinstance(self.inflateidx, Constant) and isinstance(self.takeidx, Constant): + return Tuple(tuple(map(constant, self.evalf(self.inflateidx.value, self.takeidx.value)))) def __iter__(self): shape = ArrayFromTuple(self, index=2, shape=(), dtype=int), @@ -3272,6 +3303,81 @@ def _intbounds_tuple(self): return ((0, float('inf')),) * 3 +class Assemble(Array): + + func: Array + indices: typing.Tuple[Array, ...] + shape: typing.Tuple[Array, ...] + + def __post_init__(self): + assert len(self.indices) == len(self.shape) + assert builtins.sum(index.ndim for index in self.indices) == self.func.ndim + + @property + def dtype(self): + return self.func.dtype + + @property + def dependencies(self): + return self.func, *self.indices, *self.shape + + @staticmethod + def evalf(func, *args): + n = len(args) // 2 + indices = args[:n] + shape = args[n:] + reshaped_indices = tuple(index.reshape((1,)*offset + index.shape + (1,)*(func.ndim-index.ndim-offset)) + for offset, index in zip(util.cumsum(index.ndim for index in indices), indices)) + return numeric.accumulate(func, reshaped_indices, shape) + + def _compile_with_out(self, builder, out, out_block_id, mode): + assert mode in ('iadd', 'assign') + if mode == 'assign': + builder.get_block_for_evaluable(self, block_id=out_block_id, comment='zero').array_fill_zeros(out) + compiled_indices = [] + advanced_ndim = builtins.sum(index.ndim for index in self.indices if not isinstance(index, Range)) + i = 0 + for index in self.indices: + if isinstance(index, Range): + n = builder.compile(index.shape[0]) + compiled_indices.append(_pyast.Variable('slice').call(n)) + else: + j = i + index.ndim + compiled_index = builder.compile(index) + if index.ndim < advanced_ndim: + compiled_index = compiled_index.get_item(_pyast.Raw(','.join(['None'] * i + [':'] * index.ndim + ['None'] * (advanced_ndim-j)))) + compiled_indices.append(compiled_index) + i = j + assert i == advanced_ndim + compiled_func = builder.compile(self.func) + trans = [i for i, index in enumerate(self.indices) if not isinstance(index, Range)] + if advanced_ndim and any(numpy.diff(trans) > 1): + # see https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing + trans.extend(i for i, index in enumerate(self.indices) if isinstance(index, Range)) + compiled_func = compiled_func.get_attr('transpose').call(*[_pyast.LiteralInt(i) for i in trans]) + builder.get_block_for_evaluable(self).array_add_at(out, _pyast.Tuple(tuple(compiled_indices)), compiled_func) + + def _optimized_for_numpy(self): + if isinstance(self.func, Assemble): + nontrivial_indices = [] + for f in self.func, self: # order is important to maintain insertion order of ndim=0 indices + offset = 0 + for index in f.indices: + if index != Range(f.shape[offset]): + nontrivial_indices.append((offset, index)) + offset += index.ndim + assert offset == f.func.ndim + nontrivial_indices.sort(key=lambda item: item[0]) # order of equal offsets is maintained + indices = [] + for i, index in nontrivial_indices: + if i < len(indices): + return # inflations overlap + while i > len(indices): + indices.append(Range(self.shape[len(indices)])) + indices.append(index) + return Assemble(self.func.func, tuple(indices), self.shape) + + class Diagonalize(Array): func: Array @@ -3433,8 +3539,9 @@ def _compile_expression(self, py_self, where): return _pyast.Variable('numpy').get_attr('nonzero').call(where).get_item(_pyast.LiteralInt(0)) def _simplified(self): - if self.isconstant: - return constant(self.eval()) + if isinstance(self.where, Constant): + indices, = self.where.value.nonzero() + return constant(indices) class DerivativeTargetBase(Array): @@ -3638,9 +3745,9 @@ def _inflations(self): return tuple(inflations) def _simplified(self): - if _equals_scalar_constant(self.func.shape[-2], 1): + if isunit(self.func.shape[-2]): return get(self.func, -2, constant(0)) - if _equals_scalar_constant(self.func.shape[-1], 1): + if isunit(self.func.shape[-1]): return get(self.func, -1, constant(0)) return self.func._ravel(self.ndim-1) @@ -3649,7 +3756,7 @@ def evalf(f): return f.reshape(f.shape[:-2] + (f.shape[-2]*f.shape[-1],)) def _multiply(self, other): - if isinstance(other, Ravel) and equalshape(other.func.shape[-2:], self.func.shape[-2:]): + if isinstance(other, Ravel) and _all_certainly_equal(other.func.shape[-2:], self.func.shape[-2:]): return Ravel(multiply(self.func, other.func)) return Ravel(multiply(self.func, Unravel(other, *self.func.shape[-2:]))) @@ -3683,7 +3790,7 @@ def _rtake(self, func, axis): def _unravel(self, axis, shape): if axis != self.ndim-1: return Ravel(unravel(self.func, axis, shape)) - elif equalshape(shape, self.func.shape[-2:]): + elif _all_certainly_equal(shape, self.func.shape[-2:]): return self.func def _inflate(self, dofmap, length, axis): @@ -3737,7 +3844,7 @@ def __post_init__(self): assert isinstance(self.func, Array) and self.func.ndim > 0, f'func={self.func!r}' assert _isindex(self.sh1), f'sh1={self.sh1!r}' assert _isindex(self.sh2), f'sh2={self.sh2!r}' - assert _equals_simplified(self.func.shape[-1], self.sh1 * self.sh2) + assert not _certainly_different(self.func.shape[-1], self.sh1 * self.sh2) @property def dependencies(self): @@ -3752,9 +3859,9 @@ def shape(self): return *self.func.shape[:-1], self.sh1, self.sh2 def _simplified(self): - if _equals_scalar_constant(self.shape[-2], 1): + if isunit(self.shape[-2]): return insertaxis(self.func, self.ndim-2, constant(1)) - if _equals_scalar_constant(self.shape[-1], 1): + if isunit(self.shape[-1]): return insertaxis(self.func, self.ndim-1, constant(1)) return self.func._unravel(self.ndim-2, self.shape[-2:]) @@ -3821,11 +3928,11 @@ def _take(self, index, axis): return RavelIndex(self.ia, _take(self.ib, index, axis - self.ia.ndim), self.na, self.nb) def _rtake(self, func, axis): - if _equals_simplified(func.shape[axis], self._length): + if _certainly_equal(func.shape[axis], self._length): return _take(_take(unravel(func, axis, (self.na, self.nb)), self.ib, axis+1), self.ia, axis) def _rinflate(self, func, length, axis): - if _equals_simplified(length, self._length): + if _certainly_equal(length, self._length): return Ravel(Inflate(_inflate(func, self.ia, self.na, func.ndim - self.ndim), self.ib, self.nb)) def _unravel(self, axis, shape): @@ -3866,7 +3973,7 @@ def _unravel(self, axis, shape): return RavelIndex(Range(shape[0]), Range(shape[1]), shape[0], shape[1]) def _rtake(self, func, axis): - if _equals_simplified(self.length, func.shape[axis]): + if _certainly_equal(self.length, func.shape[axis]): return func def _rinflate(self, func, length, axis): @@ -3982,7 +4089,7 @@ def _simplified(self): ncoeffs_lower, ncoeffs_upper = self.coeffs.shape[-1]._intbounds if iszero(self.coeffs): return zeros_like(self) - elif _equals_scalar_constant(self.coeffs.shape[-1], 1): + elif isunit(self.coeffs.shape[-1]): return prependaxes(get(self.coeffs, -1, constant(0)), self.points.shape[:-1]) points, where_points = unalign(self.points, naxes=self.points.ndim - 1) coeffs, where_coeffs = unalign(self.coeffs, naxes=self.coeffs.ndim - 1) @@ -4122,7 +4229,7 @@ class PolyMul(Array): def __post_init__(self): assert isinstance(self.coeffs_left, Array) and self.coeffs_left.ndim >= 1 and self.coeffs_left.dtype == float, f'coeffs_left={self.coeffs_left!r}' assert isinstance(self.coeffs_right, Array) and self.coeffs_right.ndim >= 1 and self.coeffs_right.dtype == float, f'coeffs_right={self.coeffs_right!r}' - assert equalshape(self.coeffs_left.shape[:-1], self.coeffs_right.shape[:-1]), 'PolyMul({}, {})'.format(self.coeffs_left, self.coeffs_right) + assert not _any_certainly_different(self.coeffs_left.shape[:-1], self.coeffs_right.shape[:-1]), 'PolyMul({}, {})'.format(self.coeffs_left, self.coeffs_right) @cached_property def degree_left(self): @@ -4222,9 +4329,9 @@ def evalf(self): return poly.GradPlan(self.nvars, degree) def _simplified(self): - if iszero(self.coeffs) or _equals_scalar_constant(self.degree, 0): + if iszero(self.coeffs) or iszero(self.degree): return zeros_like(self) - elif _equals_scalar_constant(self.degree, 1): + elif isunit(self.degree): return InsertAxis(Take(self.coeffs, constant(self.nvars - 1) - Range(constant(self.nvars))), constant(1)) def _takediag(self, axis1, axis2): @@ -4314,7 +4421,7 @@ class Choose(Array): def __post_init__(self): assert isinstance(self.index, Array) and self.index.dtype == int, f'index={self.index!r}' assert isinstance(self.choices, Array), f'choices={self.choices!r}' - assert equalshape(self.choices.shape[:-1], self.index.shape) + assert not _any_certainly_different(self.choices.shape[:-1], self.index.shape) @property def dependencies(self): @@ -4377,7 +4484,7 @@ class NormDim(Array): def __post_init__(self): assert isinstance(self.length, Array) and self.length.dtype == int, f'length={self.length!r}' assert isinstance(self.index, Array) and self.index.dtype == int, f'index={self.index!r}' - assert equalshape(self.length.shape, self.index.shape) + assert not _any_certainly_different(self.length.shape, self.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: # @@ -4414,8 +4521,8 @@ def _simplified(self): 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()) + if isinstance(self.length, Constant) and isinstance(self.index, Constant): + return constant(self.evalf(self.length.value, self.index.value)) def _intbounds_impl(self): lower_length, upper_length = self.length._intbounds @@ -4685,7 +4792,7 @@ def _intbounds_impl(self): return 0, max(0, upper_length - 1) def _simplified(self): - if _equals_scalar_constant(self.length, 1): + if isunit(self.length): return Zeros((), int) @property @@ -4973,10 +5080,15 @@ def _simplified(self): elif self.index not in self.func.arguments: return Ravel(Transpose.from_end(InsertAxis(self.func, self.index.length), -2)) unaligned, where = unalign(self.func) + reinserted_unit = False if self.ndim-1 not in where: # reinsert concatenation axis, at unit length if possible so we can # insert the remainder outside of the loop - unaligned = InsertAxis(unaligned, self.func.shape[-1] if self.index in self.func.shape[-1].arguments else constant(1)) + n = self.func.shape[-1] + if self.index not in n.arguments and not isunit(n): + n = constant(1) + reinserted_unit = True + unaligned = InsertAxis(unaligned, n) where += self.ndim-1, elif where[-1] != self.ndim-1: # bring concatenation axis to the end @@ -4984,7 +5096,7 @@ def _simplified(self): unaligned = Transpose.to_end(unaligned, axis) where = (*where[:axis], *where[axis+1:], self.ndim-1) f = loop_concatenate(unaligned, self.index) - if not _equals_simplified(self.shape[-1], f.shape[-1]): + if reinserted_unit: # last axis was reinserted at unit length AND it was not unit length # originally - if it was unit length originally then we proceed only if # there are other insertions to promote, otherwise we'd get a recursion. @@ -5066,6 +5178,125 @@ def _unravel(self, axis, shape): return SearchSorted(unravel(self.arg, axis, shape), array=self.array, side=self.side, sorter=self.sorter) +class ArgSort(Array): + + array: Array + + dtype = int + + def __post_init__(self): + assert self.array.ndim + + @property + def shape(self): + return self.array.shape + + @property + def dependencies(self): + return self.array, + + def evalf(self, array): + index = numpy.argsort(array, -1) + # on some platforms (windows) argsort does not return indices as + # numpy.dtype(int), so we type cast it for consistency + return index.astype(int, copy=False) + + +class UniqueMask(Array): + + sorted_array: Array + + dtype = bool + + def __post_init__(self): + assert self.sorted_array.ndim == 1 + + @property + def shape(self): + return self.sorted_array.shape + + @property + def dependencies(self): + return self.sorted_array, + + def evalf(self, sorted_array): + mask = numpy.empty(sorted_array.shape, dtype=bool) + mask[:1] = True + numpy.not_equal(sorted_array[1:], sorted_array[:-1], out=mask[1:]) + return mask + + +class UniqueInverse(Array): + + unique_mask: Array + sorter: Array + + dtype = int + + def __post_init__(self): + assert self.unique_mask.dtype == bool and self.unique_mask.ndim == 1 + assert self.sorter.dtype == int and self.sorter.ndim == 1 + assert not _any_certainly_different(self.unique_mask.shape, self.sorter.shape), (self.unique_mask, self.sorter) + + @property + def shape(self): + return self.sorter.shape + + @property + def dependencies(self): + return self.unique_mask, self.sorter + + def evalf(self, unique_mask, sorter): + inverse = numpy.empty_like(sorter) + inverse[sorter] = numpy.cumsum(unique_mask) + inverse -= 1 + return inverse + + +def unique(array, return_index=False, return_inverse=False): + sorter = ArgSort(array) + mask = UniqueMask(Take(array, sorter)) + index = Take(sorter, Find(mask)) + unique = Take(array, index) + inverse = UniqueInverse(mask, sorter) + return (unique, index, inverse)[slice(0, 2+return_inverse, 2-return_index) if return_inverse or return_index else 0] + + +class CompressIndices(Array): + + indices: Array + length: Array + + dtype = int + ndim = 1 + + def __post_init__(self): + assert self.indices.dtype == int and self.indices.ndim == 1 + assert self.length.dtype == int and self.length.ndim == 0 + + @property + def shape(self): + return self.length + 1, + + @property + def dependencies(self): + return self.indices, self.length + + @staticmethod + def evalf(indices, length): + return numeric.compress_indices(indices, length) + + +def as_csr(array): + assert array.ndim == 2 + values, (rowindices, colindices), (nrows, ncols) = array.simplified.assparse + indices, inverse = unique(rowindices * ncols + colindices, return_inverse=True) + rowidx, colidx = divmod(indices, ncols) + rowptr = CompressIndices(rowidx, nrows) + values = Inflate(values, inverse, indices.shape[0]) + return values, rowptr, colidx, ncols + + # AUXILIARY FUNCTIONS (FOR INTERNAL USE) @@ -5089,7 +5320,7 @@ def _numpy_align(a, b): return _inflate_scalar(a, b.shape), b if not b.ndim: return a, _inflate_scalar(b, a.shape) - if equalshape(a.shape, b.shape): + if not _any_certainly_different(a.shape, b.shape): return a, b raise ValueError('incompatible shapes: {} != {}'.format(*[tuple(int(n) if n.isconstant else n for n in arg.shape) for arg in (a, b)])) @@ -5180,6 +5411,11 @@ def iszero(arg): return isinstance(arg.simplified, Zeros) +def isunit(arg): + simple = arg.simplified + return isinstance(simple, Constant) and numpy.all(simple.value == 1) + + def zeros(shape, dtype=float): return Zeros(shape, dtype) @@ -5328,7 +5564,7 @@ def stack(args, axis=0): def repeat(arg, length, axis): arg = asarray(arg) - assert _equals_scalar_constant(arg.shape[axis], 1) + assert isunit(arg.shape[axis]) return insertaxis(get(arg, axis, constant(0)), axis, length) @@ -5358,7 +5594,7 @@ def grammium(arg, axes=(-2, -1)): def sqrt_abs_det_gram(arg, axes=(-2, -1)): arg = Transpose.to_end(arg, *axes) - if _equals_simplified(arg.shape[-1], arg.shape[-2]): + if _certainly_equal(arg.shape[-1], arg.shape[-2]): return abs(determinant(arg)) else: return sqrt(abs(determinant(grammium(arg)))) @@ -5398,7 +5634,7 @@ def derivative(func, var, seen=None): else: result = func._derivative(var, seen) seen[func] = result - assert equalshape(result.shape, func.shape+var.shape) and result.dtype == func.dtype, 'bug in {}._derivative'.format(type(func).__name__) + assert not _any_certainly_different(result.shape, func.shape+var.shape) and result.dtype == func.dtype, 'bug in {}._derivative'.format(type(func).__name__) return result @@ -5445,10 +5681,10 @@ def take(arg: Array, index: Array, axis: int): assert isinstance(axis, int), f'axis={axis!r}' length = arg.shape[axis] if index.dtype == bool: - assert _equals_simplified(index.shape[0], length) + assert not _certainly_different(index.shape[0], length) index = Find(index) - elif index.isconstant: - index_ = index.eval() + elif isinstance(index, Constant): + index_ = index.value ineg = numpy.less(index_, 0) if not length.isconstant: if ineg.any(): @@ -5476,17 +5712,10 @@ def _inflate(arg: Array, dofmap: Array, length: Array, axis: int): assert _isindex(length), f'length={length!r}' assert isinstance(axis, int), f'axis={axis!r}' axis = numeric.normdim(arg.ndim+1-dofmap.ndim, axis) - assert equalshape(dofmap.shape, arg.shape[axis:axis+dofmap.ndim]) + assert not _any_certainly_different(dofmap.shape, arg.shape[axis:axis+dofmap.ndim]) return Transpose.from_end(Inflate(Transpose.to_end(arg, *range(axis, axis+dofmap.ndim)), dofmap, length), axis) -def mask(arg, mask: Array, axis: int = 0): - assert isinstance(arg, Array), f'arg={arg!r}' - assert isinstance(mask, numpy.ndarray) and mask.dtype == bool and mask.ndim == 1 and _equals_scalar_constant(arg.shape[axis], len(mask)), f'mask={mask!r}' - index, = mask.nonzero() - return _take(arg, constant(index), axis) - - def unravel(func, axis, shape): func = asarray(func) axis = numeric.normdim(func.ndim, axis) @@ -5500,11 +5729,11 @@ def ravel(func, axis): return Transpose.from_end(Ravel(Transpose.to_end(func, axis, axis+1)), axis) -def _flat(func): +def _flat(func, ndim=1): func = asarray(func) - if func.ndim == 0: + if func.ndim == ndim-1: return InsertAxis(func, constant(1)) - while func.ndim > 1: + while func.ndim > ndim: func = Ravel(func) return func @@ -5578,7 +5807,7 @@ def replace_arguments(value, arguments): ''' if isinstance(value, Argument) and value.name in arguments: v = asarray(arguments[value.name]) - assert equalshape(value.shape, v.shape), (value.shape, v.shape) + assert not _any_certainly_different(value.shape, v.shape), (value.shape, v.shape) assert value.dtype == v.dtype, (value.dtype, v.dtype) return v @@ -5660,7 +5889,7 @@ def einsum(fmt, *args, **dims): for s, arg in zip(sin, args): assert len(s) == arg.ndim for c, sh in zip(s, arg.shape): - if not _equals_simplified(shapes.setdefault(c, sh), sh): + if _certainly_different(shapes.setdefault(c, sh), sh): raise ValueError('shapes do not match for axis {0[0]}{0[1]}'.format(c)) ret = None @@ -5711,7 +5940,28 @@ def eval_sparse(funcs: AsEvaluableArray, **arguments: typing.Mapping[str, numpy. yield data +@util.single_or_multiple +def eval_coo(funcs: AsEvaluableArray, arguments: typing.Mapping[str, numpy.ndarray] = {}) -> typing.Tuple[numpy.ndarray, ...]: + '''Evaluate one or several Array objects as COO sparse data. + + Args + ---- + funcs : :class:`tuple` of Array objects + Arrays to be evaluated. + arguments : :class:`dict` (default: None) + Optional arguments for function evaluation. + + Returns + ------- + results : :class:`tuple` of (values, indices, shape) triplets + ''' + + f = compile(tuple(func.as_evaluable_array.simplified.assparse for func in funcs)) + return f(**arguments) + + @functools.lru_cache(32) +@log.withcontext def compile(func, /, *, simplify: bool = True, stats: typing.Optional[str] = None, cache_const_intermediates: bool = True): '''Returns a callable that evaluates ``func``. @@ -6032,7 +6282,9 @@ def collect(evaluable): # Compile. eval(builtins.compile(script, name, 'exec'), globals) - return globals['compiled'] + compiled = globals['compiled'] + compiled.__nutils_hash__ = types.nutils_hash(script) + return compiled def _define_loop_block_structure(targets: typing.Tuple[Evaluable, ...]) -> typing.Tuple[Evaluable, ...]: diff --git a/nutils/function.py b/nutils/function.py index 355680e12..5057bf404 100644 --- a/nutils/function.py +++ b/nutils/function.py @@ -43,8 +43,7 @@ class LowerArgs(NamedTuple): coordinates: Mapping[str, evaluable.Array] def consistent(self): - return all( - evaluable.equalshape(coords.shape[:-1], self.points_shape) + return not any(evaluable._any_certainly_different(coords.shape[:-1], self.points_shape) and space in self.transform_chains for space, coords in self.coordinates.items()) @@ -2438,7 +2437,7 @@ def __init__(self, ndofs: int, nelems: int, index: Array, coords: Array) -> None self._arg_ndofs_evaluable = evaluable.asarray(self._arg_dofs_evaluable.shape[0]) assert self._arg_dofs_evaluable.ndim == 1 assert self._arg_coeffs_evaluable.ndim == 2 - assert evaluable._equals_simplified(self._arg_dofs_evaluable.shape[0], self._arg_coeffs_evaluable.shape[0]) + assert not evaluable._certainly_different(self._arg_dofs_evaluable.shape[0], self._arg_coeffs_evaluable.shape[0]) @cached_property def _arg_dofs(self): @@ -3341,7 +3340,11 @@ def norm(x, ord=None, axis=None): raise NotImplementedError('only "ord" values of None are supported for now') if axis is None: axis = range(x.ndim) - return numpy.sqrt(numpy.sum(x * numpy.conjugate(x), axis)) + # NOTE while the sum of squares is always positive, we wrap it in + # maximum(0, ..) to guard against the situation that the function + # simplifies to a form in which round-off errors may nudge it below + # zero (e.g. (a-b)^2 -> a^2 - 2 a b + b^2). + return numpy.sqrt(numpy.maximum(0, numpy.sum(numpy.real(x)**2 + numpy.imag(x)**2, axis))) def _eig(symmetric, index, a): return evaluable.Eig(a, symmetric)[index] diff --git a/nutils/matrix/__init__.py b/nutils/matrix/__init__.py index c88eb8994..72dd0c0d5 100644 --- a/nutils/matrix/__init__.py +++ b/nutils/matrix/__init__.py @@ -6,7 +6,7 @@ the ``export`` method. """ -from .. import _util as util, sparse, warnings +from .. import _util as util, sparse, warnings, numeric import numpy import importlib import os @@ -20,30 +20,145 @@ @util.set_current @util.defaults_from_env def backend(matrix: str = 'auto'): - return importlib.import_module('._'+matrix.lower(), __name__) + if isinstance(matrix, str): + return importlib.import_module('._'+matrix.lower(), __name__) + if hasattr(matrix, 'assemble'): + return matrix + raise ValueError('matrix backend does not have an assemble function') + + +def assemble_csr(values, rowptr, colidx, ncols): + '''Create sparse matrix from CSR sparse data. + + Parameters + ---------- + values + Matrix values in the row and column positions of the subsequent + parameters; one dimensional array of any data type. + rowptr + Compressed row indices; one dimensonal integer array of a length one up + from the number of matrix rows. + colidx + Column indices; one dimensional integer array of the same length as the + values parameter. + ncols + Number of matrix columns. + ''' + + values = numpy.asarray(values) + rowptr = numpy.asarray(rowptr) + colidx = numpy.asarray(colidx) + ncols = ncols.__index__() + if not values.ndim == 1: + raise MatrixError('assemble received invalid values') + if not (rowptr.ndim == 1 and + rowptr.dtype.kind in 'ui' and + rowptr[0] == 0 and + all(rowptr[1:] >= rowptr[:-1]) and + rowptr[-1] == len(values)): + raise MatrixError('assemble received invalid row indices') + if not (colidx.ndim == 1 and + colidx.dtype.kind in 'ui' and + len(colidx) == rowptr[-1] and + all(colidx < ncols)): + raise MatrixError('assemble received invalid column indices') + colidx_is_increasing = numpy.empty((len(colidx)+1,), bool) + numpy.greater_equal(colidx[1:], colidx[:-1], out=colidx_is_increasing[1:-1]) + colidx_is_increasing[rowptr] = True + if not colidx_is_increasing.all(): + raise MatrixError('column indices are not stricty increasing') + return backend.current.assemble(values, rowptr, colidx, ncols) + + +def assemble_coo(values, rowidx, nrows, colidx, ncols): + '''Create sparse matrix from COO sparse data. + + Parameters + ---------- + values + Matrix values in the row and column positions of the subsequent + parameters; one dimensional array of any data type. + rowptr + Row indices; one dimensonal integer array of the same length as the + values parameter. + nrows + Number of matrix rows. + colidx + Column indices; one dimensional integer array of the same length as the + values parameter. + ncols + Number of matrix columns. + ''' + + return assemble_csr(values, numeric.compress_indices(rowidx, nrows), colidx, ncols) def assemble(data, index, shape): - if not isinstance(data, numpy.ndarray) or data.ndim != 1 or len(index) != 2 or len(shape) != 2: - raise MatrixError('assemble received invalid input') - n, = (index[0][1:] <= index[0][:-1]).nonzero() # index[0][n+1] <= index[0][n] - if (index[0][n+1] < index[0][n]).any() or (index[1][n+1] <= index[1][n]).any(): - raise MatrixError('assemble input must be sorted') - return backend.current.assemble(data, index, shape) + warnings.deprecation('matrix.assemble is deprecated, use matrix.assemble_coo instead') + rowidx, colidx = index + nrows, ncols = shape + return assemble_coo(data, rowidx, nrows, colidx, ncols) + + +def assemble_block_csr(blocks): + '''Create sparse block matrix from stacked CSR sparse data. + + Block columns may differ per row, but must add up the the same number of + matrix columns. The matrix type is derived from the block values; in case + of varying types Numpy's casting rules apply. + + Parameters + ---------- + blocks + Nested sequence of CSR data. The outer sequence represents the matrix + rows, the middle sequences the matrix columns, and the inner sequences + the arguments to the assemble_csr function. + ''' + + ncols = sum(n for *_, n in blocks[0]) + values = [] + rowptr = [0] + colidx = [] + for row in blocks: + nrows = len(row[0][1]) - 1 + block_data = [] + col_offset = 0 + for block_values, block_rowptr, block_colidx, block_ncols in row: + assert len(block_rowptr) - 1 == nrows, 'sparse blocks have inconsistent row sizes' + if len(block_values): + block_data.append((numpy.asarray(block_values), numpy.asarray(block_rowptr), numpy.add(block_colidx, col_offset))) + col_offset += block_ncols + assert col_offset == ncols, 'sparse blocks have inconsistent column sizes' + ptr = rowptr[-1] + if len(block_data) == 1: + (block_values, block_rowptr, block_colidx), = block_data + values.append(block_values) + rowptr.extend(block_rowptr[1:] + ptr) + colidx.append(block_colidx) + else: + for irow in range(nrows): + for block_values, block_rowptr, block_colidx in block_data: + i, j = block_rowptr[irow:irow+2] + values.append(block_values[i:j]) + colidx.append(block_colidx[i:j]) + ptr += j - i + rowptr.append(ptr) + return assemble_csr(numpy.concatenate(values), numpy.array(rowptr), numpy.concatenate(colidx), ncols) def fromsparse(data, inplace=False): - indices, values, shape = sparse.extract(sparse.prune(sparse.dedup(data, inplace=inplace), inplace=True)) - return backend.current.assemble(values, indices, shape) + (rowidx, colidx), values, (nrows, ncols) = sparse.extract(sparse.prune(sparse.dedup(data, inplace=inplace), inplace=True)) + return assemble_coo(values, rowidx, nrows, colidx, ncols) def empty(shape): - return backend.current.assemble(data=numpy.empty([0], dtype=float), index=numpy.empty([len(shape), 0], dtype=int), shape=shape) + nrows, ncols = shape + return assemble_csr(numpy.zeros(0, dtype=float), numpy.zeros(nrows+1, dtype=int), numpy.zeros(0, dtype=int), ncols) def diag(d): - assert d.ndim == 1 - return backend.current.assemble(d, index=numpy.arange(len(d))[numpy.newaxis].repeat(2, axis=0), shape=d.shape*2) + n = len(d) + return assemble_csr(d, numpy.arange(n+1), numpy.arange(n), n) def eye(n): diff --git a/nutils/matrix/_base.py b/nutils/matrix/_base.py index cd4b10d91..036cca1e1 100644 --- a/nutils/matrix/_base.py +++ b/nutils/matrix/_base.py @@ -41,9 +41,9 @@ def __init__(self, shape, dtype): self._cached_submatrix = None def __reduce__(self): - from . import assemble - data, index = self.export('coo') - return assemble, (data, index, self.shape) + from . import assemble_csr + data, colidx, rowptr = self.export('csr') + return assemble_csr, (data, rowptr, colidx, self.shape[1]) @abc.abstractmethod def __add__(self, other): @@ -169,7 +169,11 @@ def solve(self, rhs=None, *, lhs0=None, constrain=None, rconstrain=None, solver= else: assert rconstrain.shape == (nrows,) and constrain.dtype == bool I = ~rconstrain - lhs[J] += self.submatrix(I, J)._solver((rhs - self @ lhs)[I], solver, atol=atol, rtol=rtol, **solverargs) + try: + lhs[J] += self.submatrix(I, J)._solver((rhs - self @ lhs)[I], solver, atol=atol, rtol=rtol, **solverargs) + except ToleranceNotReached as e: + lhs[J] += e.best + raise ToleranceNotReached(lhs) from None return lhs def solve_leniently(self, *args, **kwargs): diff --git a/nutils/matrix/_mkl.py b/nutils/matrix/_mkl.py index 5c116b9e9..e8b3e8527 100644 --- a/nutils/matrix/_mkl.py +++ b/nutils/matrix/_mkl.py @@ -18,13 +18,11 @@ raise BackendNotAvailable('the Intel MKL matrix backend requires libmkl to be installed (try: pip install mkl)') -def assemble(data, index, shape): +def assemble(data, rowptr, colidx, ncols): # In the increments below the output dtype is set to int32 not only to avoid # an additional allocation, but crucially also to avoid truncation in case # the incremented index overflows the original type. - return MKLMatrix(data, ncols=shape[1], - rowptr=numpy.add(index[0].searchsorted(numpy.arange(shape[0]+1)), 1, dtype=numpy.int32), - colidx=numpy.add(index[1], 1, dtype=numpy.int32)) + return MKLMatrix(data, ncols=ncols, rowptr=numpy.add(rowptr, 1, dtype=numpy.int32), colidx=numpy.add(colidx, 1, dtype=numpy.int32)) class Pardiso: @@ -150,7 +148,7 @@ def __matmul__(self, other): if not isinstance(other, numpy.ndarray): raise TypeError if other.shape[0] != self.shape[1]: - raise MatrixError + raise MatrixError(f'cannot multiply {self.shape[0]}x{self.shape[1]} matrix with array of length {other.shape[0]}') x = numpy.ascontiguousarray(other.T, dtype=self.dtype) y = numpy.empty(x.shape[:-1] + self.shape[:1], dtype=self.dtype) if other.ndim == 1: diff --git a/nutils/matrix/_numpy.py b/nutils/matrix/_numpy.py index 795f4b1a8..8819ba528 100644 --- a/nutils/matrix/_numpy.py +++ b/nutils/matrix/_numpy.py @@ -4,10 +4,11 @@ import functools -def assemble(data, index, shape): +def assemble(data, rowptr, colidx, ncols): + rowidx = numpy.concatenate([numpy.full(n, i) for i, n in enumerate(numpy.diff(rowptr))]) + shape = len(rowptr)-1, ncols array = numpy.zeros(shape, dtype=data.dtype) - if len(data): - array[tuple(index)] = data + array[rowidx, colidx] = data return NumpyMatrix(array) diff --git a/nutils/matrix/_scipy.py b/nutils/matrix/_scipy.py index 3c4270b05..f2d244ac0 100644 --- a/nutils/matrix/_scipy.py +++ b/nutils/matrix/_scipy.py @@ -8,8 +8,8 @@ raise BackendNotAvailable('the Scipy matrix backend requires scipy to be installed (try: pip install scipy)') -def assemble(data, index, shape): - return ScipyMatrix(scipy.sparse.csr_matrix((data, index), shape)) +def assemble(data, rowptr, colidx, ncols): + return ScipyMatrix(scipy.sparse.csr_matrix((data, colidx, rowptr), (len(rowptr)-1, ncols))) class ScipyMatrix(Matrix): @@ -82,7 +82,8 @@ def mycallback(arg): if callback: callback(res) reformat(100 * numpy.log10(max(atol, res)) / numpy.log10(atol)) - lhs, status = solverfun(self.core, rhs, M=precon, tol=0., atol=atol, callback=mycallback, **solverargs) + solverargs['rtol' if scipy.__version__ >= '1.12' else 'tol'] = 0 + lhs, status = solverfun(self.core, rhs, M=precon, atol=atol, callback=mycallback, **solverargs) if status != 0: raise Exception('status {}'.format(status)) return lhs diff --git a/nutils/numeric.py b/nutils/numeric.py index e199a43cb..de80b37ba 100644 --- a/nutils/numeric.py +++ b/nutils/numeric.py @@ -444,11 +444,9 @@ def accumulate(data, index, shape): ... return array ''' - ndim = len(shape) - assert data.ndim == 1 - assert len(index) == ndim and all(isintarray(ind) and ind.shape == data.shape for ind in index) - if not ndim: - return data.sum() + assert len(index) == len(shape), f'{len(index)} != {len(shape)}' + if not shape: + return data.sum(dtype=data.dtype) retval = numpy.zeros(shape, data.dtype) numpy.add.at(retval, tuple(index), data) return retval @@ -678,4 +676,28 @@ def sanitize_einsum_subscripts(subscripts, *shapes): return (*in_, out) +def compress_indices(indices, length): + '''Return insertion positions in monotonically increasing index vector. + + The compressed array ``c`` is such that ``indices[c[i]:c[i+1]] == i`` for + any integer ``0 <= i < length``. Fails with ``ValueError`` if ``indices`` + are out of bounds or not monotonically increasing; otherwise fully + equivalent to ``indices.searchsorted(numpy.arange(length + 1))``. + ''' + + if not len(indices): + return numpy.zeros((length+1,), int) + if indices[0] < 0 or indices[-1] >= length: + raise ValueError('indices are out of bounds') + step = numpy.empty(len(indices)+1, dtype=int) + step[0] = indices[0] + 1 + numpy.subtract(indices[1:], indices[:-1], out=step[1:-1]) + step[-1] = length - indices[-1] + nz, = step.nonzero() + try: + return numpy.repeat(nz, step[nz]) # fails if any step is negative + except ValueError: + raise ValueError('indices are not monotomically increasing') from None + + # vim:sw=4:sts=4:et diff --git a/nutils/solver.py b/nutils/solver.py index 59a9bdda1..20760e383 100644 --- a/nutils/solver.py +++ b/nutils/solver.py @@ -20,8 +20,9 @@ >>> res = domain.integral('(basis_n,i u_,i + basis_n) d:x' @ ns, degree=2) >>> lhs = solver.solve_linear('lhs', residual=res, constrain=cons) -solve > solving 25 dof system to machine precision using arnoldi solver -solve > solver returned with residual ... +solve > solving for argument lhs (36) using direct method +solve > solve > solving 25 dof system to machine precision using arnoldi solver +solve > solve > solver returned with residual ... The coefficients ``lhs`` represent the solution to the Poisson problem. @@ -32,6 +33,7 @@ from . import function, evaluable, cache, numeric, types, _util as util, matrix, warnings, sparse from dataclasses import dataclass +from typing import Optional, Union, Tuple, Dict, Any, Iterator, Callable import abc import numpy import itertools @@ -180,6 +182,392 @@ def __call__(self, res0, dres0, res1, dres1): return min(max(scale, self.minscale), self.maxscale), scale >= self.acceptscale +# SYSTEM + +class System: + '''System of one or more variables. + + The System class represents a linear or nonlinear problem of one or more + variables, and offers several methods to solve it. The main solution method + is ``solve``, which serves as a general entry point that offloads the + problem to specialized solvers. The other prominent method is ``step``, + which adds functionality to deal with time dependent problems. + ''' + + def __init__(self, residual: Union[evaluable.AsEvaluableArray,Tuple[evaluable.AsEvaluableArray,...]], /, trial: Union[str,Tuple[str,...]], test: Optional[Union[str,Tuple[str,...]]] = None): + '''Construct a system. + + Parameters + ---------- + residual : + Any object that supports conversion to a scalar evaluable Array via + the AsEvaluableArray protocol. This is the function that the trial + arguments will seek to make constant in all the test arguments, or + zero in all the derivatives. If provided as a tuple of vectors then + these are already the derivatives to be made zero, and no test + arguments may be specified. + trial : + Names of the trial arguments, provided as a string ('arg1,arg2') or + tuple of strings (('arg1', 'arg2')). + test : + Names of the test arguments (optional) provided as a string or + tuple of strings. Must be omitted if residual is a tuple of + vectors; otherwise, may be omitted if test and trial are equal. + ''' + + self.trials = tuple(trial.split(',') if isinstance(trial, str) else trial) + + if isinstance(residual, (tuple, list)): + if test is not None: + raise ValueError('a test argument is not allowed in combination with a residual vector') + residuals = [res.as_evaluable_array for res in residual] + self.dtype = residuals[0].dtype + if not all(v.dtype == self.dtype for v in residuals): + raise ValueError('inconsistent data types in residual vector') + argobjects = _dict((arg.name, arg) for res in residuals for arg in res.arguments if isinstance(arg, evaluable.Argument)) + self.is_symmetric = False + else: + functional = residual.as_evaluable_array + self.dtype = functional.dtype + argobjects = {arg.name: arg for arg in functional.arguments if isinstance(arg, evaluable.Argument)} + tests = self.trials if test is None else tuple(test.split(',') if isinstance(test, str) else test) + residuals = [evaluable.derivative(functional, argobjects[t]) for t in tests] + self.is_symmetric = self.trials == tests + + 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) + 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 + + @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) + if self.is_constant: + self.__matrix_cache = mat + return mat, vec, val + + def prepare_solution_vector(self, arguments: Dict[str, numpy.ndarray], constrain: Dict[str, numpy.ndarray]): + arguments = arguments.copy() + x = numpy.empty(self.__trial_offsets[-1], self.dtype) + iscons = numpy.empty(self.__trial_offsets[-1], dtype=bool) + for trial, i, j in zip(self.trials, self.__trial_offsets, self.__trial_offsets[1:]): + trialshape = self.trialshapes[trial] + trialarg = x[i:j].reshape(trialshape) + trialarg[...] = arguments.get(trial, 0) + c = constrain.get(trial, False) + if c is not False: + assert c.shape == trialshape + if c.dtype != bool: + c, v = ~numpy.isnan(c), c + trialarg[c] = v[c] + trialcons = iscons[i:j].reshape(trialshape) + trialcons[...] = c + arguments[trial] = trialarg # IMPORTANT: arguments share memory with x + return x, iscons, arguments + + @property + def _trial_info(self): + return ' and '.join(t + ' (' + ','.join(map(str, shape)) + ')' for t, shape in self.trialshapes.items()) + + MethodIter = Iterator[Tuple[Dict[str, numpy.ndarray], float, float]] + + @cache.function + @log.withcontext + def solve(self, *, arguments: Dict[str, numpy.ndarray] = {}, constrain: Dict[str, numpy.ndarray] = {}, linargs: Dict[str, Any] = {}, tol: float = 0., miniter: int = 0, maxiter: Optional[int] = None, method: Optional[Callable[...,MethodIter]] = None) -> Tuple[Dict[str, numpy.ndarray], float]: + '''Solve the system. + + Determines the trial arguments for which the derivatives of the + system's scalar functional are zero in all the test arguments. + + Parameters + ---------- + arguments + Arguments required to evaluate the system. Arguments that coincide + with the system's trial arguments serve as initial guess. + constrain + Constrained values for any of the trial arguments, supplied as + float vectors with NaN entries or boolean vectors, the latter + holding values at those of the initial guess. + linargs + Keyword arguments to be passed on to the linear solver. + tol + Maximum residual norm, stopping criterion for the iterative solver. + Required for nonlinear problems. + miniter + Minimum number of iterations for the iterative solver, overruling + to ``tol`` based stopping criterion. + maxiter + Maximum number of iterations, after which a ``SolverError`` is + raised. + method + Iterative solution method, in the form of a callable that returns + an iterator of (arguments, resnorm, value) triplets. + ''' + + if method is None and not self.is_linear: + method = Newton() + log.info(f'{"optimizing" if self.is_symmetric else "solving"} for argument {self._trial_info} using {method or "direct"} method') + if method is None: + x, iscons, arguments = self.prepare_solution_vector(arguments, constrain) + jac, res, val = self.assemble(arguments) + dx = -jac.solve(res, constrain=iscons, **_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 + return arguments + if tol <= 0: + raise ValueError('iterative solver requires a strictly positive tolerance') + first = _First() + with log.iter.plain('iter', itertools.count()) as steps: + for arguments, resnorm, val in method(self, arguments=arguments, constrain=constrain, linargs=linargs): + progress = numpy.log(first(resnorm)/resnorm) / numpy.log(first(resnorm)/tol) if resnorm > tol else 1 + log.info(f'residual: {resnorm:.0e} ({100*progress:.0f}%)') + iiter = next(steps) # opens new log context + if iiter >= miniter and resnorm <= tol: + break + if maxiter is not None and iiter >= maxiter: + raise SolverError(f'failed to converge in {maxiter} iterations') + if self.is_symmetric: + log.info(f'optimal value: {val:.1e}') + return arguments + + def step(self, *, timestep: float, timetarget: str, historysuffix: str, arguments: Dict[str, numpy.ndarray], **solveargs) -> Dict[str, numpy.ndarray]: + '''Advance a time step. + + This method is best described by an example. Let ``timetarget`` equal + 't' and ``historysuffix`` '0', and let our system's trial arguments be + 'u' and 'v'. This method then creates argument copies 't0', 'u0', 'v0', + advances 't' by ``timestep``, and solves for the new 'u' and 'v'. + + Parameters + ---------- + timetarget + Name of the scalar argument that tracks the time. + timestep + Size of the time increment. + historysuffix + String suffix to add to argument names to denote their value at the + beginning of the time step. + arguments + Arguments required to evaluate the system. Arguments that coincide + with the system's trial arguments serve as initial value. + solveargs + Remaining keyword arguments are passed on to the ``solver`` method. + ''' + + arguments = arguments.copy() + for trial in self.trials: + if trial in arguments: + arguments[trial + historysuffix] = arguments[trial] + time = arguments.get(timetarget, 0.) + arguments[timetarget + historysuffix] = time + arguments[timetarget] = time + timestep + try: + return self.solve(arguments=arguments, **solveargs) + except (SolverError, matrix.MatrixError) as e: + log.error(f'error: {e}; retrying with timestep {timestep/2}') + with log.context('tic'): + halfway_arguments = self.step(timestep=timestep/2, timetarget=timetarget, historysuffix=historysuffix, arguments=arguments, **solveargs) + with log.context('toc'): + return self.step(timestep=timestep/2, timetarget=timetarget, historysuffix=historysuffix, arguments=halfway_arguments, **solveargs) + + @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. + + 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). + + Parameters + ---------- + arguments + Arguments required to evaluate the system. Any arguments that + coincide with the system's trial arguments serve as initial guess. + constrain + Constrained values for any of the trial arguments, supplied as + float vectors with NaN entries or boolean vectors, the latter + holding values at those of the initial guess. + linargs + Keyword arguments to be passed on to the linear solver. + droptol + Minimum absolute value if matrix entries for rows and columns to + participate in the optimization problem. + ''' + + log.info(f'optimizing 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') + 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 + x += dx + x[nosupp & ~iscons] = numpy.nan + return arguments + + +@dataclass(eq=True, frozen=True) +class Newton: + + def __str__(self): + return 'newton' + + def __call__(self, system, *, arguments: Dict[str, numpy.ndarray] = {}, constrain: Dict[str, numpy.ndarray] = {}, linargs: Dict[str, Any] = {}) -> System.MethodIter: + x, iscons, arguments = system.prepare_solution_vector(arguments, constrain) + linargs = _copy_with_defaults(linargs, rtol=1-3, symmetric=system.is_symmetric) + while True: + jac, res, val = system.assemble(arguments) + yield arguments, numpy.linalg.norm(res[~iscons]), val + x -= jac.solve_leniently(res, constrain=iscons, **linargs) + + +@dataclass(eq=True, frozen=True) +class LinesearchNewton: + + strategy: Callable = NormBased() + failrelax: float = 1e-6 + relax0: float = 1. + + def __str__(self): + return 'linesearch-newton' + + def __call__(self, system, *, arguments: Dict[str, numpy.ndarray] = {}, constrain: Dict[str, numpy.ndarray] = {}, linargs: Dict[str, Any] = {}) -> System.MethodIter: + x, iscons, arguments = system.prepare_solution_vector(arguments, constrain) + linargs = _copy_with_defaults(linargs, rtol=1-3, symmetric=system.is_symmetric) + jac, res, val = system.assemble(arguments) + relax = self.relax0 + while True: # newton iterations + yield arguments, numpy.linalg.norm(res[~iscons]), val + dx = -jac.solve_leniently(res, constrain=iscons, **linargs) + x += dx * relax + res0 = res + jac0dx = jac@dx # == res0 if dx was solved to infinite precision + while True: # line search + jac, res, _ = system.assemble(arguments) + relax, adjust = self._linesearch(res0, jac0dx, res, jac@dx, iscons, relax) + if not adjust: + break + if relax <= self.failrelax: + raise SolverError('stuck in local minimum') + x += dx * adjust + + def _linesearch(self, res0, dres0, res1, dres1, iscons, relax): + isdof = ~iscons + scale, accept = self.strategy(res0[isdof], dres0[isdof] * relax, res1[isdof], dres1[isdof] * relax) + if accept: + log.info('update accepted at relaxation', round(relax, 5)) + return min(relax * scale, 1), 0 + assert scale < 1 + return relax * scale, relax * (scale-1) + + +@dataclass(eq=True, frozen=True) +class Minimize: + + rampup: float = .5 + rampdown: float = -1. + failrelax: float = -10. + + def __str__(self): + return 'newton' + + 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') + 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) + relax = 0. + while True: + yield arguments, numpy.linalg.norm(res[~iscons]), val + dx = -jac.solve_leniently(res, constrain=iscons, **linargs) # baseline: vanilla Newton + # compute first two ritz values to determine approximate path of steepest descent + zres = res * ~iscons + dxnorm = numpy.linalg.norm(dx) + k0 = dx / dxnorm + k1 = -zres / dxnorm # == jac @ k0 + a = k1 @ k0 + k1 -= k0 * a # orthogonalize + c = numpy.linalg.norm(k1) + k1 /= c # normalize + b = k1 @ (jac @ k1) + # at this point k0 and k1 are orthonormal, and [k0 k1]^T jac [k0 k1] = [a c; c b] + D = numpy.hypot(b-a, 2*c) + L = numpy.array([a+b-D, a+b+D]) / 2 # 2nd order ritz values: eigenvalues of [a c; c b] + v0, v1 = zres + dx * L[:, numpy.newaxis] + V = numpy.stack([v1, -v0], axis=1) / D # ritz vectors times dx -- note: V @ L == -zres, V.sum() == dx + log.info('spectrum: {:.1e}..{:.1e} ({}definite)'.format(*L, 'positive ' if L[0] > 0 else 'negative ' if L[-1] < 0 else 'in')) + val0 = val + while True: # line search along steepest descent curve + r = numpy.exp(relax - numpy.log(D)) # == exp(relax) / D + eL = numpy.exp(-r*L) + dx -= V @ eL + x += dx + jac, res, val = system.assemble(arguments) + slope = res @ (V @ (eL*L)) + log.info('energy {:+.2e} / e{:+.1f} and {}creasing'.format(val - val0, relax, 'in' if slope > 0 else 'de')) + if numpy.isfinite(val) and numpy.isfinite(res).all() and val <= val0 and slope <= 0: + relax += self.rampup + break + relax += self.rampdown + if relax <= self.failrelax: + raise SolverError('stuck in local minimum') + dx = V @ eL # return to baseline + + +@dataclass(eq=True, frozen=True) +class Pseudotime: + + inertia: Tuple[evaluable.AsEvaluableArray,...] + timestep: float + + def __str__(self): + return 'pseudotime' + + def __call__(self, system, *, arguments: Dict[str, numpy.ndarray] = {}, constrain: Dict[str, numpy.ndarray] = {}, linargs: Dict[str, Any] = {}) -> System.MethodIter: + x, iscons, arguments = system.prepare_solution_vector(arguments, constrain) + linargs = _copy_with_defaults(linargs, rtol=1-3) + djac = self._assemble_inertia_matrix(system.trialshapes, arguments) + + first = _First() + while True: + jac, res, val = system.assemble(arguments) + resnorm = numpy.linalg.norm(res[~iscons]) + yield arguments, resnorm, val + timestep = self.timestep * (first(resnorm) / resnorm) + log.info(f'timestep: {timestep:.0e}') + x -= (jac + djac / timestep).solve_leniently(res, constrain=iscons, **linargs) + + def _assemble_inertia_matrix(self, trialshapes, arguments): + argobjs = [evaluable.Argument(t, tuple(map(evaluable.constant, shape)), float) for t, shape in trialshapes.items()] + djacobians = [[evaluable._flat(evaluable.derivative(evaluable._flat(res), argobj).simplified, 2) for argobj in argobjs] for res in self.inertia] + djac_blocks = evaluable.compile(tuple(tuple(map(evaluable.as_csr, row)) for row in djacobians)) + return matrix.assemble_block_csr(djac_blocks(**arguments)) + + # SOLVERS def solve_linear(target, residual, *, constrain = None, lhs0: types.arraydata = None, arguments = {}, **kwargs): @@ -208,28 +596,13 @@ def solve_linear(target, residual, *, constrain = None, lhs0: types.arraydata = lhs0=lhs0, arguments=arguments if lhs0 is None else {**arguments, target: lhs0}, **kwargs)[target] if lhs0 is not None: raise ValueError('lhs0 argument is invalid for a non-string target; define the initial guess via arguments instead') - target, residual = _target_helper(target, residual) - solveargs = _strip(kwargs, 'lin') + linargs = _strip(kwargs, 'lin') if kwargs: raise TypeError('unexpected keyword arguments: {}'.format(', '.join(kwargs))) - return _solve_linear(target, residual, - types.frozendict((k, types.arraydata(v)) for k, v in (constrain or {}).items()), - types.frozendict((k, types.arraydata(v)) for k, v in (arguments or {}).items()), - types.frozendict(solveargs)) - - -@cache.function -def _solve_linear(target, residual: tuple, constraints: dict, arguments: dict, solveargs: dict): - arguments, constraints = _parse_lhs_cons(constraints, target, _argobjs(residual), arguments) - jacobians = _derivative(residual, target) - if not set(target).isdisjoint(_argobjs(jacobians)): + system = System(residual, *_split_trial_test(target)) + if not system.is_linear: raise SolverError('problem is not linear') - dtype = _determine_dtype(target, residual, arguments, constraints) - lhs, vlhs = _redict(arguments, target, dtype) - mask, vmask = _invert(constraints, target) - res, jac = _integrate_blocks(residual, jacobians, arguments=lhs, mask=mask) - vlhs[vmask] -= jac.solve(res, **solveargs) - return lhs + return system.solve(arguments=arguments, constrain=constrain or {}, linargs=linargs) def newton(target, residual, *, jacobian = None, lhs0 = None, relax0: float = 1., constrain = None, linesearch=NormBased(), failrelax: float = 1e-6, arguments = {}, **kwargs): @@ -283,66 +656,14 @@ def newton(target, residual, *, jacobian = None, lhs0 = None, relax0: float = 1. failrelax=failrelax, arguments=arguments if lhs0 is None else {**arguments, target: lhs0}, **kwargs)[target] if lhs0 is not None: raise ValueError('lhs0 argument is invalid for a non-string target; define the initial guess via arguments instead') - target, residual = _target_helper(target, residual) - solveargs = _strip(kwargs, 'lin') - solveargs.setdefault('rtol', 1e-3) + if jacobian is not None: + warnings.warn('jac argument is no longer in use') + linargs = _strip(kwargs, 'lin') if kwargs: raise TypeError('unexpected keyword arguments: {}'.format(', '.join(kwargs))) - return _with_solve(_newton(target, residual, None if jacobian is None else tuple(jacobian), - types.frozendict((k, types.arraydata(v)) for k, v in (constrain or {}).items()), - types.frozendict((k, types.arraydata(v)) for k, v in (arguments or {}).items()), - linesearch, relax0, failrelax, types.frozendict(solveargs))) - - -class _newton(cache.Recursion, length=1): - - def __init__(self, target, residual, jacobian, constrain, arguments, linesearch, relax0: float, failrelax: float, solveargs): - super().__init__() - self.target = target - self.residual = residual - self.jacobian = _derivative(residual, target, jacobian) - self.lhs0, self.constrain = _parse_lhs_cons(constrain, target, _argobjs(residual), arguments) - self.dtype = _determine_dtype(target, residual, self.lhs0, self.constrain) - self.relax0 = relax0 - self.linesearch = linesearch - self.failrelax = failrelax - self.solveargs = solveargs - - def _eval(self, lhs, mask): - return _integrate_blocks(self.residual, self.jacobian, arguments=lhs, mask=mask) - - def resume(self, history): - mask, vmask = _invert(self.constrain, self.target) - if history: - lhs, info = history[-1] - lhs, vlhs = _redict(lhs, self.target, self.dtype) - res, jac = self._eval(lhs, mask) - relax = info.relax - else: - lhs, vlhs = _redict(self.lhs0, self.target, self.dtype) - res, jac = self._eval(lhs, mask) - relax = self.relax0 - yield lhs, types.attributes(resnorm=numpy.linalg.norm(res), relax=relax) - while True: - dlhs = -jac.solve_leniently(res, **self.solveargs) # compute new search vector - res0 = res - dres = jac@dlhs # == -res if dlhs was solved to infinite precision - vlhs[vmask] += relax * dlhs - res, jac = self._eval(lhs, mask) - if self.linesearch: - scale, accept = self.linesearch(res0, relax*dres, res, relax*(jac@dlhs)) - while not accept: # line search - assert scale < 1 - oldrelax = relax - relax *= scale - if relax <= self.failrelax: - raise SolverError('stuck in local minimum') - vlhs[vmask] += (relax - oldrelax) * dlhs - res, jac = self._eval(lhs, mask) - scale, accept = self.linesearch(res0, relax*dres, res, relax*(jac@dlhs)) - log.info('update accepted at relaxation', round(relax, 5)) - relax = min(relax * scale, 1) - yield lhs, types.attributes(resnorm=numpy.linalg.norm(res), relax=relax) + system = System(residual, *_split_trial_test(target)) + method = Newton() if linesearch is None else LinesearchNewton(strategy=linesearch, relax0=relax0, failrelax=failrelax) + return _with_solve(system, method, arguments, constrain or {}, linargs) def minimize(target, energy: evaluable.asarray, *, lhs0: types.arraydata = None, constrain = None, rampup: float = .5, rampdown: float = -1., failrelax: float = -10., arguments = {}, **kwargs): @@ -390,89 +711,12 @@ def minimize(target, energy: evaluable.asarray, *, lhs0: types.arraydata = None, failrelax=failrelax, arguments=arguments if lhs0 is None else {**arguments, target: lhs0}, **kwargs)[target] if lhs0 is not None: raise ValueError('lhs0 argument is invalid for a non-string target; define the initial guess via arguments instead') - if isinstance(target, str): - target = target.rstrip(',').split(',') - solveargs = _strip(kwargs, 'lin') - solveargs['symmetric'] = True + linargs = _strip(kwargs, 'lin') if kwargs: raise TypeError('unexpected keyword arguments: {}'.format(', '.join(kwargs))) - return _with_solve(_minimize(tuple(target), energy.as_evaluable_array, - types.frozendict((k, types.arraydata(v)) for k, v in (constrain or {}).items()), - types.frozendict((k, types.arraydata(v)) for k, v in (arguments or {}).items()), - rampup, rampdown, failrelax, types.frozendict(solveargs))) - - -class _minimize(cache.Recursion, length=1, version=3): - - def __init__(self, target, energy: evaluable.asarray, constrain, arguments, rampup: float, rampdown: float, failrelax: float, solveargs): - super().__init__() - if energy.shape != (): - raise ValueError('`energy` should be scalar') - self.target = target - self.energy = energy - self.residual = _derivative((energy,), target) - self.jacobian = _derivative(self.residual, target) - self.lhs0, self.constrain = _parse_lhs_cons(constrain, target, _argobjs((energy,)), arguments) - self.dtype = _determine_dtype(target, (energy,), self.lhs0, self.constrain) - self.rampup = rampup - self.rampdown = rampdown - self.failrelax = failrelax - self.solveargs = solveargs - - def _eval(self, lhs, mask): - return _integrate_blocks(self.energy, self.residual, self.jacobian, arguments=lhs, mask=mask) - - def resume(self, history): - mask, vmask = _invert(self.constrain, self.target) - if history: - lhs, info = history[-1] - lhs, vlhs = _redict(lhs, self.target, self.dtype) - nrg, res, jac = self._eval(lhs, mask) - relax = info.relax - else: - lhs, vlhs = _redict(self.lhs0, self.target, self.dtype) - nrg, res, jac = self._eval(lhs, mask) - relax = 0 - yield lhs, types.attributes(resnorm=numpy.linalg.norm(res), energy=nrg, relax=relax) - - while True: - nrg0 = nrg - dlhs = -jac.solve_leniently(res, **self.solveargs) - vlhs[vmask] += dlhs # baseline: vanilla Newton - - # compute first two ritz values to determine approximate path of steepest descent - dlhsnorm = numpy.linalg.norm(dlhs) - k0 = dlhs / dlhsnorm - k1 = -res / dlhsnorm # = jac @ k0 - a = k1 @ k0 - k1 -= k0 * a # orthogonalize - c = numpy.linalg.norm(k1) - k1 /= c # normalize - b = k1 @ (jac @ k1) - # at this point k0 and k1 are orthonormal, and [k0 k1]^T jac [k0 k1] = [a c; c b] - D = numpy.hypot(b-a, 2*c) - L = numpy.array([a+b-D, a+b+D]) / 2 # 2nd order ritz values: eigenvalues of [a c; c b] - v0, v1 = res + dlhs * L[:, numpy.newaxis] - V = numpy.array([v1, -v0]).T / D # ritz vectors times dlhs -- note: V.dot(L) = -res, V.sum() = dlhs - log.info('spectrum: {:.1e}..{:.1e} ({}definite)'.format(*L, 'positive ' if L[0] > 0 else 'negative ' if L[-1] < 0 else 'in')) - - eL = 0 - for irelax in itertools.count(): # line search along steepest descent curve - r = numpy.exp(relax - numpy.log(D)) # = exp(relax) / D - eL0 = eL - eL = numpy.exp(-r*L) - vlhs[vmask] -= V.dot(eL - eL0) - nrg, res, jac = self._eval(lhs, mask) - slope = res.dot(V.dot(eL*L)) - log.info('energy {:+.2e} / e{:+.1f} and {}creasing'.format(nrg - nrg0, relax, 'in' if slope > 0 else 'de')) - if numpy.isfinite(nrg) and numpy.isfinite(res).all() and nrg <= nrg0 and slope <= 0: - relax += self.rampup - break - relax += self.rampdown - if relax <= self.failrelax: - raise SolverError('stuck in local minimum') - - yield lhs, types.attributes(resnorm=numpy.linalg.norm(res), energy=nrg, relax=relax) + system = System(energy, *_split_trial_test(target)) + method = Minimize(rampup=rampup, rampdown=rampdown, failrelax=failrelax) + return _with_solve(system, method, arguments, constrain or {}, linargs) def pseudotime(target, residual, inertia, timestep: float, *, lhs0: types.arraydata = None, constrain = None, arguments = {}, **kwargs): @@ -511,63 +755,12 @@ def pseudotime(target, residual, inertia, timestep: float, *, lhs0: types.arrayd if lhs0 is not None: raise ValueError('lhs0 argument is invalid for a non-string target; define the initial guess via arguments instead') target, residual, inertia = _target_helper(target, residual, inertia) - solveargs = _strip(kwargs, 'lin') - solveargs.setdefault('rtol', 1e-3) + linargs = _strip(kwargs, 'lin') if kwargs: raise TypeError('unexpected keyword arguments: {}'.format(', '.join(kwargs))) - return _with_solve(_pseudotime(target, residual, inertia, timestep, - types.frozendict((k, types.arraydata(v)) for k, v in (constrain or {}).items()), - types.frozendict((k, types.arraydata(v)) for k, v in (arguments or {}).items()), - types.frozendict(solveargs))) - - -class _pseudotime(cache.Recursion, length=1): - - def __init__(self, target, residual, inertia, timestep: float, constrain, arguments, solveargs: dict): - super().__init__() - if target in arguments: - raise ValueError('`target` should not be defined in `arguments`') - if len(residual) != len(inertia): - raise Exception('length of residual and inertia do no match') - for inert, res in zip(inertia, residual): - if inert and not evaluable.equalshape(inert.shape, res.shape): - raise ValueError('expected `inertia` with shape {} but got {}'.format(res.shape, inert.shape)) - self.target = target - self.timesteptarget = '_pseudotime_timestep' - dt = evaluable.Argument(self.timesteptarget, ()) - self.residuals = residual - self.jacobians = _derivative(tuple(res + (inert/dt if inert else 0) for res, inert in zip(residual, inertia)), target) - self.lhs0, self.constrain = _parse_lhs_cons(constrain, target, _argobjs(residual+inertia), arguments) - self.dtype = _determine_dtype(target, residual+inertia, self.lhs0, self.constrain) - self.timestep = timestep - self.solveargs = solveargs - - def _eval(self, lhs, mask, timestep): - return _integrate_blocks(self.residuals, self.jacobians, arguments=dict({self.timesteptarget: timestep}, **lhs), mask=mask) - - def resume(self, history): - mask, vmask = _invert(self.constrain, self.target) - if history: - lhs, info = history[-1] - lhs, vlhs = _redict(lhs, self.target, self.dtype) - resnorm0 = info.resnorm0 - timestep = info.timestep - res, jac = self._eval(lhs, mask, timestep) - resnorm = numpy.linalg.norm(res) - else: - lhs, vlhs = _redict(self.lhs0, self.target, self.dtype) - timestep = self.timestep - res, jac = self._eval(lhs, mask, timestep) - resnorm = resnorm0 = numpy.linalg.norm(res) - yield lhs, types.attributes(resnorm=resnorm, timestep=timestep, resnorm0=resnorm0) - - while True: - vlhs[vmask] -= jac.solve_leniently(res, **self.solveargs) - timestep = self.timestep * (resnorm0/resnorm) - log.info('timestep: {:.0e}'.format(timestep)) - res, jac = self._eval(lhs, mask, timestep) - resnorm = numpy.linalg.norm(res) - yield lhs, types.attributes(resnorm=resnorm, timestep=timestep, resnorm0=resnorm0) + system = System(residual, *_split_trial_test(target)) + method = Pseudotime(inertia=inertia, timestep=timestep) + return _with_solve(system, method, arguments, constrain or {}, linargs) def thetamethod(target, residual, inertia, timestep: float, theta: float, *, lhs0: types.arraydata = None, constrain = None, newtontol: float = 1e-10, arguments = {}, newtonargs: types.frozendict = {}, timetarget: str = '_thetamethod_time', time0: float = 0., historysuffix: str = '0'): @@ -614,54 +807,35 @@ def thetamethod(target, residual, inertia, timestep: float, theta: float, *, lhs if lhs0 is not None: raise ValueError('lhs0 argument is invalid for a non-string target; define the initial condition via arguments instead') target, residual, inertia = _target_helper(target, residual, inertia) - return _thetamethod(target, residual, inertia, timestep, - types.frozendict((k, types.arraydata(v)) for k, v in (constrain or {}).items()), - types.frozendict((k, types.arraydata(v)) for k, v in (arguments or {}).items()), - theta, newtontol, types.frozendict(newtonargs), timetarget, time0, historysuffix) - - -class _thetamethod(cache.Recursion, length=1, version=1): - - def __init__(self, target, residual, inertia, timestep: float, constrain, arguments, theta: float, newtontol: float, newtonargs: types.frozendict, timetarget: str, time0: float, historysuffix: str): - super().__init__() - if len(residual) != len(inertia): - raise Exception('length of residual and inertia do no match') - for inert, res in zip(inertia, residual): - if not evaluable.equalshape(inert.shape, res.shape): - raise ValueError('expected `inertia` with shape {} but got {}'.format(res.shape, inert.shape)) - self.target = target - self.newtonargs = newtonargs - self.newtontol = newtontol - self.timestep = timestep - self.timetarget = timetarget - self.lhs0, self.constrain = _parse_lhs_cons(constrain, target, _argobjs(residual+inertia), arguments) - self.lhs0[timetarget] = numpy.array(time0) - self.old_new = [(t+historysuffix, t) for t in target] - self.old_new.append((timetarget+historysuffix, timetarget)) - subs0 = {new: evaluable.Argument(old, tuple(map(evaluable.constant, self.lhs0[new].shape))) for old, new in self.old_new} - dt = evaluable.Argument(timetarget, ()) - subs0[timetarget] - self.residuals = tuple(res * evaluable.astype(theta, res.dtype) + evaluable.replace_arguments(res, subs0) * evaluable.astype(1-theta, res.dtype) + (inert - evaluable.replace_arguments(inert, subs0)) / evaluable.astype(dt, res.dtype) for res, inert in zip(residual, inertia)) - self.jacobians = _derivative(self.residuals, target) - - def _step(self, lhs0, dt): - arguments = lhs0.copy() - arguments.update((old, lhs0[new]) for old, new in self.old_new) - arguments[self.timetarget] = lhs0[self.timetarget] + dt - try: - return newton(self.target, residual=self.residuals, jacobian=self.jacobians, constrain=self.constrain, arguments=arguments, **self.newtonargs).solve(tol=self.newtontol) - except (SolverError, matrix.MatrixError) as e: - log.error('error: {}; retrying with timestep {}'.format(e, dt/2)) - return self._step(self._step(lhs0, dt/2), dt/2) - def resume(self, history): - if history: - lhs, = history - else: - lhs = self.lhs0 - yield lhs - while True: - lhs = self._step(lhs, self.timestep) - yield lhs + argobjs = _dict((arg.name, arg) for func in residual + inertia if func is not None for arg in func.arguments if isinstance(arg, evaluable.Argument)) + argobjs.setdefault(timetarget, evaluable.Argument(timetarget, ())) + + old_new = [(t+historysuffix, t) for t in target] + old_new.append((timetarget+historysuffix, timetarget)) + + subs0 = {new: evaluable.Argument(old, argobjs[new].shape) for old, new in old_new} + dt = evaluable.Argument(timetarget, ()) - subs0[timetarget] + theta = evaluable.constant(float(theta)) + residuals = tuple(res * theta + evaluable.replace_arguments(res, subs0) * (1.-theta) + + (inert - evaluable.replace_arguments(inert, subs0)) / dt for res, inert in zip(residual, inertia)) + + arguments = arguments.copy() + arguments.setdefault(timetarget, time0) + system = System(residuals, *_split_trial_test(target)) + newtonargs = dict(newtonargs) + linesearch = newtonargs.pop('linesearch', NormBased()) + method = None if system.is_linear else Newton() if linesearch is None else LinesearchNewton(strategy=linesearch, **newtonargs) + return _thetamethod(system, arguments, + timestep=timestep, timetarget=timetarget, historysuffix=historysuffix, constrain=constrain or {}, tol=newtontol, method=method) + + +def _thetamethod(system, arguments, **stepargs): + with log.iter.plain('timestep', itertools.count()) as steps: + yield arguments + for _ in steps: + arguments = system.step(arguments=arguments, **stepargs) + yield arguments impliciteuler = functools.partial(thetamethod, theta=1) @@ -710,87 +884,39 @@ def optimize(target, functional: evaluable.asarray, *, tol: float = 0., argument constrain={} if constrain is None else {target: constrain}, relax0=relax0, linesearch=linesearch, failrelax=failrelax, **kwargs)[target] if lhs0 is not None: raise ValueError('lhs0 argument is invalid for a non-string target; define the initial guess via arguments instead') - if isinstance(target, str): - target = target.rstrip(',').split(',') - solveargs = _strip(kwargs, 'lin') - solveargs['symmetric'] = True + linargs = _strip(kwargs, 'lin') if kwargs: raise TypeError('unexpected keyword arguments: {}'.format(', '.join(kwargs))) - with log.context('optimize'): - return _optimize(tuple(target), functional.as_evaluable_array, - types.frozendict((k, types.arraydata(v)) for k, v in (constrain or {}).items()), - types.frozendict((k, types.arraydata(v)) for k, v in (arguments or {}).items()), - tol, droptol, relax0, linesearch, failrelax, types.frozendict(solveargs)) - - -@cache.function(version=1) -def _optimize(target, functional: evaluable.asarray, constrain, arguments, tol: float, droptol: float, relax0: float, linesearch, failrelax: float, solveargs): - argobjs = _argobjs((functional,)) - if any(t not in argobjs for t in target): - if not droptol: - raise ValueError('target {} does not occur in integrand; consider setting droptol>0'.format(', '.join(t for t in target if t not in argobjs))) - target = [t for t in target if t in argobjs] - if not target: - return {} - residual = _derivative((functional,), target) - jacobian = _derivative(residual, target) - lhs0, constrain = _parse_lhs_cons(constrain, target, argobjs, arguments) - dtype = _determine_dtype(target, (functional,), lhs0, constrain) - mask, vmask = _invert(constrain, target) - lhs, vlhs = _redict(lhs0, target, dtype) - val, res, jac = _integrate_blocks(functional, residual, jacobian, arguments=lhs, mask=mask) - if droptol is not None: - supp = jac.rowsupp(droptol) - res = res[supp] - jac = jac.submatrix(supp, supp) - nan = numpy.zeros_like(vmask) - nan[vmask] = ~supp # return value is set to nan if dof is not supported and not constrained - vmask[vmask] = supp # dof is computed if it is supported and not constrained - assert vmask.sum() == len(res) - resnorm = numpy.linalg.norm(res) - solveargs = dict(solveargs) - if not set(target).isdisjoint(_argobjs(jacobian)): - if tol <= 0: - raise ValueError('nonlinear optimization problem requires a nonzero "tol" argument') - solveargs.setdefault('rtol', 1e-3) - firstresnorm = resnorm - relax = relax0 - accept = True - with log.context('newton {:.0f}%', 0) as reformat: - while not numpy.isfinite(resnorm) or resnorm > tol: - if accept: - reformat(100 * numpy.log(firstresnorm/resnorm) / numpy.log(firstresnorm/tol)) - dlhs = -jac.solve_leniently(res, **solveargs) - res0 = res - dres = jac@dlhs # == -res0 if dlhs was solved to infinite precision - relax0 = 0 - vlhs[vmask] += (relax - relax0) * dlhs - relax0 = relax # currently applied relaxation - val, res, jac = _integrate_blocks(functional, residual, jacobian, arguments=lhs, mask=mask) - resnorm = numpy.linalg.norm(res) - scale, accept = linesearch(res0, relax*dres, res, relax*(jac@dlhs)) - relax = min(relax * scale, 1) - if relax <= failrelax: - raise SolverError('stuck in local minimum') - log.info('converged with residual {:.1e}'.format(resnorm)) - elif resnorm > tol: - solveargs.setdefault('atol', tol) - dlhs = -jac.solve(res, **solveargs) - vlhs[vmask] += dlhs - val += (res + jac@dlhs/2).dot(dlhs) + system = System(functional, *_split_trial_test(target)) if droptol is not None: - vlhs[nan] = numpy.nan - log.info('constrained {}/{} dofs'.format(len(vlhs)-nan.sum(), len(vlhs))) - log.info('optimum value {:.2e}'.format(val)) - return lhs + return system.optimize(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) # HELPER FUNCTIONS + def _strip(kwargs, prefix): return {key[len(prefix):]: kwargs.pop(key) for key in list(kwargs) if key.startswith(prefix)} +def _dict(items): + '''construct dictionary from items while checking that duplicate keys have matching values''' + + d = {} + for k, v in items: + v_ = d.setdefault(k, v) + if v_ is not v and v_ != v: # cheap test first to avoid equality checks + raise ValueError('incompatible items') + return d + + +def _copy_with_defaults(d, **kwargs): + kwargs.update(d) + return kwargs + + def _parse_lhs_cons(constrain, targets, argobjs, arguments): arguments = {t: numpy.asarray(a) for t, a in arguments.items()} constrain = {t: numpy.asarray(c) for t, c in constrain.items()} @@ -813,123 +939,70 @@ def _parse_lhs_cons(constrain, targets, argobjs, arguments): return arguments, constrain -def _derivative(residual, target, jacobian=None): - argobjs = _argobjs(residual) - if jacobian is None: - jacobian = tuple(evaluable.derivative(res, argobjs[t]).simplified for res in residual for t in target) - elif len(jacobian) != len(residual) * len(target): - raise ValueError('jacobian has incorrect length') - elif not all(evaluable.equalshape(jacobian[i*len(target)+j].shape, res.shape + argobjs[t].shape) for i, res in enumerate(residual) for j, t in enumerate(target)): - raise ValueError('jacobian has incorrect shape') - return jacobian - - -def _redict(lhs, targets, dtype=float): - '''copy argument dictionary referencing a newly allocated contiguous array''' - - vlhs = numpy.empty(sum(lhs[target].size for target in targets), dtype) - lhs = lhs.copy() - offset = 0 - for target in targets: - old = lhs[target] - nextoffset = offset + old.size - new = vlhs[offset:nextoffset].reshape(old.shape) - new[...] = old - new.flags.writeable = False - lhs[target] = new - offset = nextoffset - assert offset == len(vlhs) - return lhs, vlhs - - -def _invert(cons, targets): - '''invert constraints dictionary to tuple referencing a contiguous array''' - - mask = [] - vmask = numpy.empty(sum(cons[target].size for target in targets), dtype=bool) - offset = 0 - for target in targets: - c = cons[target] - nextoffset = offset + c.size - mask.append(numpy.invert(c, out=vmask[offset:nextoffset].reshape(c.shape))) - offset = nextoffset - assert offset == len(vmask) - return tuple(mask), vmask - - -def _integrate_blocks(*blocks, arguments, mask): - '''helper function for blockwise integration''' - - *scalars, residuals, jacobians = blocks - assert len(residuals) == len(mask) - assert len(jacobians) == len(mask)**2 - data = iter(evaluable.eval_sparse((*scalars, *residuals, *jacobians), **arguments)) - nrg = [sparse.toarray(next(data)) for _ in range(len(scalars))] - res = [sparse.take(next(data), [m]) for m in mask] - jac = [[sparse.take(next(data), [mi, mj]) for mj in mask] for mi in mask] - assert not list(data) - return nrg + [sparse.toarray(sparse.block(res)), matrix.fromsparse(sparse.block(jac), inplace=True)] - - -def _argobjs(funcs): - '''get :class:`evaluable.Argument` dependencies of multiple functions''' - - argobjs = {} - for func in filter(None, funcs): - for arg in func.arguments: - if isinstance(arg, evaluable.Argument): - if arg.name in argobjs: - if argobjs[arg.name] != arg: - raise ValueError('shape or dtype mismatch for argument {}: {} != {}'.format(arg.name, argobjs[arg.name], arg)) - else: - argobjs[arg.name] = arg - return argobjs - - -def _determine_dtype(targets, residuals, lhs0, constrain): - argobjs = _argobjs(residuals) - dtype = complex if ( - any(argobjs[target].dtype == complex for target in targets) - or any(residual.dtype == complex for residual in residuals if residual is not None) - or any(vec.dtype.kind == 'c' for vec in lhs0.values()) - or any(vec.dtype.kind == 'c' for vec in constrain.values()) - ) else float - if not all(argobjs[target].dtype == dtype for target in targets): - raise ValueError('All targets must have dtype {}.'.format(dtype.__name__)) - return dtype +def _split_trial_test(target): + if isinstance(target, str): + target = target.rstrip(',') + target = target.split(',') if target else [] + if not target: + raise ValueError('no targets specified') + target = [item.split(':') if isinstance(item, str) else item for item in target] + n = len(target[0]) + if not all(len(t) == n for t in target): + raise ValueError('inconsistent targets') + if n == 1: + trial, = zip(*target) + test = None + elif n == 2: + trial, test = zip(*target) + else: + raise ValueError('invalid targets') + return trial, test def _target_helper(target, *args): - targets = target.rstrip(',').split(',') if isinstance(target, str) else list(target) - is_functional = [':' in target for target in targets] - if all(is_functional): - targets, tests = zip(*[t.split(':', 1) for t in targets]) + trial, test = _split_trial_test(target) + if test is not None: arguments = function._join_arguments(arg.arguments for arg in args) - testargs = [function.Argument(t, *arguments[t]) for t in tests] + testargs = [function.Argument(t, *arguments[t]) for t in test] args = [map(arg.derivative, testargs) for arg in args] - elif any(is_functional): - raise ValueError('inconsistent targets') elif len(args) > 1: shapes = [{f.shape for f in ziparg if f is not None} for ziparg in zip(*args)] if any(len(arg) != len(shapes) for arg in args) or any(len(shape) != 1 for shape in shapes): raise ValueError('inconsistent residuals') args = [[function.zeros(shape) if f is None else f for f, (shape,) in zip(arg, shapes)] for arg in args] - return (tuple(targets), *[tuple(f.as_evaluable_array for f in arg) for arg in args]) + return trial, *[tuple(f.as_evaluable_array for f in arg) for arg in args] + + +class _First: + def __call__(self, value): + try: + return self.value + except AttributeError: + self.value = value + return value -class _with_solve(types.Immutable): +@dataclass(frozen=True) +class _with_solve: '''add a .solve method to iterables''' - def __init__(self, wrapped, item = None): - self._wrapped = wrapped - self._item = item + system: System + method: Any + arguments: Any + constrain: Any + linargs: Any + item: Optional[str] = None def __iter__(self): - return iter(self._wrapped) if self._item is None else ((res[self._item], info) for (res, info) in self._wrapped) + iters = self.method(self.system, arguments=self.arguments, constrain=self.constrain, linargs=self.linargs) + for arguments, resnorm, value in iters: + lhs = arguments if self.item is None else arguments[self.item] + info = types.attributes(resnorm=resnorm) if not self.system.is_symmetric else types.attributes(resnorm=resnorm, energy=value) + yield lhs, info def __getitem__(self, item): - assert self._item is None - return _with_solve(self._wrapped, item) + assert self.item is None + return _with_solve(self.system, self.method, self.arguments, self.constrain, self.linargs, item) def solve(self, tol, maxiter=float('inf'), miniter=0): '''execute nonlinear solver, return lhs @@ -967,7 +1040,7 @@ def solve_withinfo(self, tol, maxiter=float('inf'), miniter=0): if miniter > maxiter: raise ValueError('The minimum number of iterations cannot be larger than the maximum.') - with log.context(self._wrapped.__class__.__name__.strip('_')): + with log.context(str(self.method)): with log.context('iter {}', 0) as recontext: it = enumerate(self) iiter, (lhs, info) = next(it) diff --git a/nutils/topology.py b/nutils/topology.py index 3d02626eb..e6c7276c8 100644 --- a/nutils/topology.py +++ b/nutils/topology.py @@ -1561,6 +1561,8 @@ def trim(self, levelset, maxrefine, ndivisions=8, name='trimmed', leveltopo=None with log.iter.percentage('trimming', range(len(self)), self.references) as items: for ielem, ref in items: levels = levelset(_trim_index=ielem, **arguments) + if numpy.isnan(levels).any(): + raise Exception('levelset function evaluated to NaN values') refs.append(ref.trim(levels, maxrefine=maxrefine, ndivisions=ndivisions)) else: # `levelset` is evaluable on `leveltopo`, which must be a uniform @@ -1596,6 +1598,8 @@ def trim(self, levelset, maxrefine, ndivisions=8, name='trimmed', leveltopo=None lowered_levelset[degree] = evaluable.compile(levelset.lower(lower_args), stats=False) levels[indices] = lowered_levelset[degree](_ielem=ielem, **arguments) mask[indices] = False + if numpy.isnan(levels).any(): + raise Exception('levelset function evaluated to NaN values') refs.append(ref.trim(levels, maxrefine=maxrefine, ndivisions=ndivisions)) log.debug('cache', fcache.stats) return SubsetTopology(self, refs, newboundary=name) diff --git a/tests/test_evaluable.py b/tests/test_evaluable.py index 79d193aba..b0b108b3f 100644 --- a/tests/test_evaluable.py +++ b/tests/test_evaluable.py @@ -1,4 +1,4 @@ -from nutils import evaluable, sparse, numeric, _util as util, types, sample +from nutils import evaluable, sparse, numeric, _util as util, types, sample, matrix from nutils.testing import TestCase, parametrize import nutils_poly as poly import numpy @@ -25,14 +25,16 @@ def setUp(self): self.actual = self.op(*self.args) self.desired = self.n_op(*self.arg_values) assert numpy.isfinite(self.desired).all(), 'something is wrong with the design of this unit test' + self.varshape = not all(n.isconstant for n in self.actual.shape) + othershape = self.desired.shape if not self.varshape else () if self.actual.dtype == bool: - self.other = numpy.random.normal(size=self.desired.shape) < 0 + self.other = numpy.random.normal(size=othershape) < 0 elif self.actual.dtype == int: - self.other = numpy.random.randint(0, 100, size=self.desired.shape) + self.other = numpy.random.randint(0, 100, size=othershape) elif self.actual.dtype == float: - self.other = numpy.random.normal(size=self.desired.shape) + self.other = numpy.random.normal(size=othershape) elif self.actual.dtype == complex: - self.other = numpy.random.normal(size=self.desired.shape) + 1j * numpy.random.normal(size=self.desired.shape) + self.other = numpy.random.normal(size=othershape) + 1j * numpy.random.normal(size=othershape) self.pairs = [(i, j) for i in range(self.actual.ndim-1) for j in range(i+1, self.actual.ndim) if self.actual.shape[i] == self.actual.shape[j]] _builtin_warnings.simplefilter('ignore', evaluable.ExpensiveEvaluationWarning) @@ -40,9 +42,12 @@ def test_dtype(self): self.assertEqual(self.desired.dtype, self.actual.dtype) def test_shapes(self): - self.assertEqual(self.desired.shape, tuple(n.__index__() for n in self.actual.shape)) + evalargs = dict(zip(self.arg_names, self.arg_values)) + self.assertEqual(self.desired.shape, tuple(n.eval(**evalargs) for n in self.actual.shape)) def assertArrayAlmostEqual(self, actual, desired, decimal): + if actual.dtype != desired.dtype: + self.fail('dtypes of actual {} and desired {} are different.'.format(actual.dtype, desired.dtype)) if actual.shape != desired.shape: self.fail('shapes of actual {} and desired {} are incompatible.'.format(actual.shape, desired.shape)) error = actual - desired if not actual.dtype.kind == desired.dtype.kind == 'b' else actual ^ desired @@ -71,14 +76,17 @@ def assertFunctionAlmostEqual(self, actual, desired, decimal): with self.subTest('optimized'): self.assertArrayAlmostEqual(actual.optimized_for_numpy.eval(**evalargs), desired, decimal) with self.subTest('sparse'): - indices, values, shape = sparse.extract(evaluable.eval_sparse(actual, **evalargs)) - self.assertEqual(tuple(map(int, shape)), desired.shape) - if not indices: - dense = values.sum() - else: - dense = numpy.zeros(desired.shape, values.dtype) - numpy.add.at(dense, indices, values) - self.assertArrayAlmostEqual(dense, desired, decimal) + values, indices, shape = evaluable.eval_coo(actual, evalargs) + self.assertEqual(shape, desired.shape) + self.assertArrayAlmostEqual(numeric.accumulate(values, indices, shape), desired, decimal) + if actual.ndim == 2: + with self.subTest('csr'): + values, rowptr, colidx, ncols = evaluable.compile(evaluable.as_csr(actual))(**evalargs) + shape = len(rowptr) - 1, ncols + self.assertEqual(shape, desired.shape) + actual = numpy.zeros(shape, dtype=values.dtype) + actual[numpy.concatenate([numpy.full(n, i) for i, n in enumerate(numpy.diff(rowptr))]), colidx] = values + self.assertArrayAlmostEqual(actual, desired, decimal) def test_evalconst(self): self.assertFunctionAlmostEqual(decimal=14, @@ -155,7 +163,7 @@ def test_determinant(self): def test_take(self): indices = [0, -1] for iax, sh in enumerate(self.desired.shape): - if sh >= 2: + if sh >= 2 and self.actual.shape[iax].isconstant: self.assertFunctionAlmostEqual(decimal=14, desired=numpy.take(self.desired, indices, axis=iax), actual=evaluable.take(self.actual, evaluable.constant(indices), axis=iax)) @@ -179,7 +187,7 @@ def test_take_nomask(self): def test_take_reversed(self): indices = [-1, 0] for iax, sh in enumerate(self.desired.shape): - if sh >= 2: + if sh >= 2 and self.actual.shape[iax].isconstant: self.assertFunctionAlmostEqual(decimal=14, desired=numpy.take(self.desired, indices, axis=iax), actual=evaluable.take(self.actual, evaluable.constant(indices), axis=iax)) @@ -194,7 +202,7 @@ def test_take_duplicate_indices(self): def test_inflate(self): for iax, sh in enumerate(self.desired.shape): - dofmap = evaluable.constant(numpy.arange(int(sh)) * 2) + dofmap = evaluable.Range(self.actual.shape[iax]) * 2 desired = numpy.zeros(self.desired.shape[:iax] + (int(sh)*2-1,) + self.desired.shape[iax+1:], dtype=self.desired.dtype) desired[(slice(None),)*iax+(slice(None, None, 2),)] = self.desired self.assertFunctionAlmostEqual(decimal=14, @@ -203,12 +211,12 @@ def test_inflate(self): def test_inflate_duplicate_indices(self): for iax, sh in enumerate(self.desired.shape): - dofmap = numpy.arange(sh) % 2 + dofmap = evaluable.Range(self.actual.shape[iax]) % 2 desired = numpy.zeros(self.desired.shape[:iax] + (2,) + self.desired.shape[iax+1:], dtype=self.desired.dtype) - numpy.add.at(desired, (slice(None),)*iax+(dofmap,), self.desired) + numpy.add.at(desired, (slice(None),)*iax+(numpy.arange(sh) % 2,), self.desired) self.assertFunctionAlmostEqual(decimal=14, desired=desired, - actual=evaluable._inflate(self.actual, dofmap=evaluable.constant(dofmap), length=evaluable.constant(2), axis=iax)) + actual=evaluable._inflate(self.actual, dofmap=dofmap, length=evaluable.constant(2), axis=iax)) def test_diagonalize(self): for axis in range(self.actual.ndim): @@ -276,7 +284,7 @@ def test_power(self): def test_power0(self): if self.actual.dtype == bool: return - power = (numpy.arange(self.desired.size) % 2).reshape(self.desired.shape).astype(self.actual.dtype) + power = (numpy.arange(self.desired.size) % 2).reshape(self.desired.shape).astype(self.actual.dtype) if not self.varshape else 3 self.assertFunctionAlmostEqual(decimal=13, desired=self.desired**power, actual=self.actual**power) @@ -299,21 +307,9 @@ def test_imag(self): def test_conjugate(self): self.assertFunctionAlmostEqual(decimal=14, - desired=numpy.conjugate(self.desired), + desired=numpy.conjugate(self.desired).astype(self.desired.dtype), actual=evaluable.conjugate(self.actual)) - def test_mask(self): - for idim in range(self.actual.ndim): - if self.desired.shape[idim] <= 1: - continue - mask = numpy.ones(self.desired.shape[idim], dtype=bool) - mask[0] = False - if self.desired.shape[idim] > 2: - mask[-1] = False - self.assertFunctionAlmostEqual(decimal=14, - desired=self.desired[(slice(None,),)*idim+(mask,)], - actual=evaluable.mask(self.actual, mask, axis=idim)) - def test_ravel(self): for idim in range(self.actual.ndim-1): self.assertFunctionAlmostEqual(decimal=14, @@ -322,6 +318,8 @@ def test_ravel(self): def test_unravel(self): for idim in range(self.actual.ndim): + if not self.actual.shape[idim].isconstant: + continue length = self.desired.shape[idim] unravelshape = (length//3, 3) if (length % 3 == 0) else (length//2, 2) if (length % 2 == 0) else (length, 1) self.assertFunctionAlmostEqual(decimal=14, @@ -329,7 +327,7 @@ def test_unravel(self): actual=evaluable.unravel(self.actual, axis=idim, shape=tuple(map(evaluable.constant, unravelshape)))) def test_loopsum(self): - if self.desired.dtype == bool: + if self.desired.dtype == bool or self.varshape: return length = 3 index = evaluable.loop_index('_testindex', length) @@ -343,6 +341,8 @@ def test_loopsum(self): desired=desired) def test_loopconcatenate(self): + if not self.desired.ndim: + return length = 3 index = evaluable.loop_index('_testindex', length) for iarg, arg_value in enumerate(self.arg_values): @@ -548,7 +548,6 @@ def _check(name, op, n_op, *arg_values, hasgrad=True, zerograd=False, ndim=2): _check('eig', lambda a: evaluable.eig(a+a.swapaxes(0, 1), symmetric=True)[1], lambda a: numpy.linalg.eigh(a+a.swapaxes(0, 1))[1], ANY(4, 4), hasgrad=False) _check('eig-complex', lambda a: evaluable.eig(a+a.swapaxes(0, 1))[1], lambda a: numpy.linalg.eig(a+a.swapaxes(0, 1))[1], ANC(4, 4), hasgrad=False) _check('mod', lambda a, b: evaluable.mod(a, b), lambda a, b: numpy.mod(a, b), ANY(4), NZ(4), hasgrad=False) -_check('mask', lambda f: evaluable.mask(f, numpy.array([True, False, True, False, True, False, True]), axis=1), lambda a: a[:, ::2], ANY(4, 7, 4)) _check('ravel', lambda f: evaluable.ravel(f, axis=1), lambda a: a.reshape(4, 4, 4, 4), ANY(4, 2, 2, 4, 4)) _check('unravel', lambda f: evaluable.unravel(f, axis=1, shape=[evaluable.constant(2), evaluable.constant(2)]), lambda a: a.reshape(4, 2, 2, 4, 4), ANY(4, 4, 4, 4)) _check('ravelindex', lambda a, b: evaluable.RavelIndex(a, b, evaluable.constant(12), evaluable.constant(20)), lambda a, b: a[..., numpy.newaxis, numpy.newaxis] * 20 + b, INT(3, 4), INT(4, 5)) @@ -593,6 +592,10 @@ def _check(name, op, n_op, *arg_values, hasgrad=True, zerograd=False, ndim=2): _check('searchsorted', lambda a: evaluable.SearchSorted(evaluable.asarray(a), array=types.arraydata(numpy.linspace(0, 1, 9)), side='left', sorter=None), lambda a: numpy.searchsorted(numpy.linspace(0, 1, 9), a).astype(int), POS(4, 2)) _check('searchsorted_sorter', lambda a: evaluable.SearchSorted(evaluable.asarray(a), array=types.arraydata([.2,.8,.4,0,.6,1]), side='left', sorter=types.arraydata([3,0,2,4,1,5])), lambda a: numpy.searchsorted([.2,.8,.4,0,.6,1], a, sorter=[3,0,2,4,1,5]).astype(int), POS(4, 2)) +_check('argsort', evaluable.ArgSort, lambda a: numpy.argsort(a, axis=-1).astype(int), ANY(3, 9)) +_check('unique', evaluable.unique, numpy.unique, numpy.arange(10) % 3) +_check('unique-index', lambda a: evaluable.unique(a, return_index=True)[1], lambda a: numpy.unique(a, return_index=True)[1].astype(int), numpy.arange(10) % 3) +_check('unique-inverse', lambda a: evaluable.unique(a, return_inverse=True)[1], lambda a: numpy.unique(a, return_inverse=True)[1].astype(int), numpy.arange(10) % 3) class compile(TestCase): @@ -940,7 +943,7 @@ def test_loop_concatenate(self): '└ B = Loop\n' 'NODES\n' '%B0 = LoopConcatenate i\n' - '├ shape[0] = %A0 = Take; i:; [2,2]\n' + '├ shape[0] = %A0 = Take; i:; [0,2]\n' '│ ├ %A1 = _SizesToOffsets; i:3; [0,2]\n' '│ │ └ %A2 = InsertAxis; i:(2); [1,1]\n' '│ │ ├ 1\n' diff --git a/tests/test_matrix.py b/tests/test_matrix.py index f8e580c58..08f44ea09 100644 --- a/tests/test_matrix.py +++ b/tests/test_matrix.py @@ -3,6 +3,87 @@ from nutils import matrix, sparse, testing, warnings +class construction(testing.TestCase): + + def setUp(self): + super().setUp() + class TestBackend: + def assemble(values, rowptr, colidx, ncols): + self.assertIsInstance(values, numpy.ndarray) + self.assertEqual(values.ndim, 1) + self.assertIsInstance(rowptr, numpy.ndarray) + self.assertEqual(rowptr.ndim, 1) + self.assertIn(rowptr.dtype.kind, 'iu') + self.assertIsInstance(colidx, numpy.ndarray) + self.assertEqual(colidx.ndim, 1) + self.assertIn(colidx.dtype.kind, 'iu') + self.assertEqual(len(colidx), len(values)) + self.assertIsInstance(ncols, int) + return values, rowptr, colidx, ncols + self.enter_context(matrix.backend(TestBackend)) + + def test_assemble_csr(self): + orig_values = [10., 20., 30.] + orig_rowptr = [0, 2, 3] + orig_colidx = [0, 2, 1] + orig_ncols = 3 + values, rowptr, colidx, ncols = matrix.assemble_csr(orig_values, orig_rowptr, orig_colidx, orig_ncols) + self.assertEqual(values.tolist(), orig_values) + self.assertEqual(rowptr.tolist(), orig_rowptr) + self.assertEqual(colidx.tolist(), orig_colidx) + self.assertEqual(ncols, orig_ncols) + + def test_assemble_coo(self): + orig_values = [10., 20., 30.] + orig_colidx = [0, 2, 1] + orig_ncols = 3 + values, rowptr, colidx, ncols = matrix.assemble_coo(orig_values, [0, 0, 1], 2, orig_colidx, orig_ncols) + self.assertEqual(values.tolist(), orig_values) + self.assertEqual(rowptr.tolist(), [0, 2, 3]) + self.assertEqual(colidx.tolist(), orig_colidx) + self.assertEqual(ncols, orig_ncols) + + def test_fromsparse(self): + data = sparse.prune(sparse.fromarray(numpy.eye(2, 3, 1)), inplace=True) + values, rowptr, colidx, ncols = matrix.fromsparse(data) + self.assertEqual(values.tolist(), [1., 1.]) + self.assertEqual(rowptr.tolist(), [0, 1, 2]) + self.assertEqual(colidx.tolist(), [1, 2]) + self.assertEqual(ncols, 3) + + def test_empty(self): + values, rowptr, colidx, ncols = matrix.empty((3, 2)) + self.assertEqual(values.tolist(), []) + self.assertEqual(rowptr.tolist(), [0, 0, 0, 0]) + self.assertEqual(colidx.tolist(), []) + self.assertEqual(ncols, 2) + + def test_diag(self): + values, rowptr, colidx, ncols = matrix.diag([10., 20., 30.]) + self.assertEqual(values.tolist(), [10., 20., 30.]) + self.assertEqual(rowptr.tolist(), [0, 1, 2, 3]) + self.assertEqual(colidx.tolist(), [0, 1, 2]) + self.assertEqual(ncols, 3) + + def test_eye(self): + values, rowptr, colidx, ncols = matrix.eye(3) + self.assertEqual(values.tolist(), [1., 1., 1.]) + self.assertEqual(rowptr.tolist(), [0, 1, 2, 3]) + self.assertEqual(colidx.tolist(), [0, 1, 2]) + self.assertEqual(ncols, 3) + + def test_assemble_block_csr(self): + A00 = [10., 20., 30.], [0, 2, 3], [0, 1, 0], 2 # 2x2 + A01 = [], [0, 0, 0], [], 2 # 2x2 + A10 = [40., 60.], [0, 1, 2], [0, 0], 1 # 2x1 + A11 = [50.], [0, 1, 1], [2], 3 # 2x3 + values, rowptr, colidx, ncols = matrix.assemble_block_csr([[A00, A01], [A10, A11]]) + self.assertEqual(values.tolist(), [10., 20., 30., 40., 50., 60.]) + self.assertEqual(rowptr.tolist(), [0, 2, 3, 5, 6]) + self.assertEqual(colidx.tolist(), [0, 1, 0, 0, 3, 0]) + self.assertEqual(ncols, 4) + + @testing.parametrize class backend(testing.TestCase): @@ -104,7 +185,7 @@ def test_div(self): def test_add(self): j = self.n//2 v = 10. - other = matrix.assemble(numpy.array([v]*self.n), numpy.array([numpy.arange(self.n), [j]*self.n]), shape=(self.n, self.n)) + other = matrix.assemble_coo(numpy.full(self.n, v), numpy.arange(self.n), self.n, numpy.full(self.n, j), self.n) add = self.matrix + other numpy.testing.assert_equal(actual=add.export('dense'), desired=self.exact + numpy.eye(self.n)[j]*v) with self.assertRaises(TypeError): @@ -115,7 +196,7 @@ def test_add(self): def test_sub(self): j = self.n//2 v = 10. - other = matrix.assemble(numpy.array([v]*self.n), numpy.array([numpy.arange(self.n), [j]*self.n]), shape=(self.n, self.n)) + other = matrix.assemble_coo(numpy.full(self.n, v), numpy.arange(self.n), self.n, numpy.full(self.n, j), self.n) sub = self.matrix - other numpy.testing.assert_equal(actual=sub.export('dense'), desired=self.exact - numpy.eye(self.n)[j]*v) with self.assertRaises(TypeError): @@ -124,13 +205,13 @@ def test_sub(self): self.matrix - matrix.eye(self.n+1) def test_transpose(self): - asym = matrix.assemble(numpy.arange(1, 7), numpy.array([[0, 0, 0, 1, 1, 2], [0, 1, 2, 1, 2, 2]]), shape=(3, 3)) + asym = matrix.assemble_coo(numpy.arange(1, 7), [0, 0, 0, 1, 1, 2], 3, [0, 1, 2, 1, 2, 2], 3) exact = numpy.array([[1, 2, 3], [0, 4, 5], [0, 0, 6]], dtype=float) transpose = asym.T numpy.testing.assert_equal(actual=transpose.export('dense'), desired=exact.T) def test_rowsupp(self): - sparse = matrix.assemble(numpy.array([1e-10, 0, 1, 1]), numpy.array([[0, 0, 2, 2], [0, 1, 1, 2]]), shape=(3, 3)) + sparse = matrix.assemble_coo([1e-10, 0, 1, 1], [0, 0, 2, 2], 3, [0, 1, 1, 2], 3) self.assertEqual(tuple(sparse.rowsupp(tol=1e-5)), (False, False, True)) self.assertEqual(tuple(sparse.rowsupp(tol=0)), (True, False, True)) @@ -152,7 +233,7 @@ def test_multisolve(self): self.assertLess(numpy.max(res), 1e-9) def test_singular(self): - singularmatrix = matrix.assemble(numpy.arange(self.n)-self.n//2, numpy.arange(self.n)[numpy.newaxis].repeat(2, 0), shape=(self.n, self.n)) + singularmatrix = matrix.diag(numpy.arange(self.n)-self.n//2) rhs = numpy.ones(self.n) for args in self.solve_args: with self.subTest(args.get('solver', 'direct')), self.assertRaises(matrix.MatrixError): @@ -188,7 +269,7 @@ def test_submatrix(self): numpy.testing.assert_equal(actual=array, desired=self.exact[numpy.ix_(rows, cols)]) def test_submatrix_specialcases(self): - mat = matrix.assemble(numpy.array([1, 2, 3, 4]), numpy.array([[0, 0, 2, 2], [0, 2, 0, 2]]), (3, 3)) + mat = matrix.assemble_coo([1, 2, 3, 4], [0, 0, 2, 2], 3, [0, 2, 0, 2], 3) self.assertAllEqual(mat.export('dense'), [[1, 0, 2], [0, 0, 0], [3, 0, 4]]) self.assertAllEqual(mat.submatrix([0, 2], [0, 1, 2]).export('dense'), [[1, 0, 2], [3, 0, 4]]) self.assertAllEqual(mat.submatrix([0, 1, 2], [0, 2]).export('dense'), [[1, 2], [0, 0], [3, 4]]) diff --git a/tests/test_numeric.py b/tests/test_numeric.py index 5cf61e07d..6f2b887e8 100644 --- a/tests/test_numeric.py +++ b/tests/test_numeric.py @@ -392,3 +392,48 @@ def test_invalid(self): self.sanitize_einsum_subscripts('i...j->...', (2,)) with self.assertRaisesRegex(ValueError, "einstein sum subscripts string included output subscript 'j' which never appeared in an input"): self.sanitize_einsum_subscripts('i->j', (2,)) + + +class accumulate(TestCase): + + def test_scalar(self): + A = numeric.accumulate(numpy.array([1, 2, 3]), (), ()) + self.assertEqual(A, 6) + + def test_vector(self): + A = numeric.accumulate(numpy.array([1, 2, 3]), (numpy.array([0, 3, 0]),), (4,)) + self.assertEqual(A.tolist(), [4, 0, 0, 2]) + + def test_matrix(self): + A = numeric.accumulate(numpy.array([1, 2, 3]), (numpy.array([0, 1, 0]), numpy.array([0, 3, 0])), (2,4)) + self.assertEqual(A.tolist(), [[4, 0, 0, 0], [0, 0, 0, 2]]) + + def test_matrix_2dindex(self): + A = numeric.accumulate(numpy.array([[1, 2], [3, 4]]), (numpy.array([[0], [1]]), numpy.array([[0, 3]])), (2,4)) + self.assertEqual(A.tolist(), [[1, 0, 0, 2], [3, 0, 0, 4]]) + + def test_matrix_slice(self): + A = numeric.accumulate(numpy.array([[1, 2], [3, 4]]), (numpy.array([0, 1]), slice(0, 4, 3)), (2,4)) + self.assertEqual(A.tolist(), [[1, 0, 0, 2], [3, 0, 0, 4]]) + + +class compress_indices(TestCase): + + def test_valid(self): + length = 8 + for indices in [0, 1, 2], [0, 2, 2, 2, 5, 6, 7], []: + c = numeric.compress_indices(indices, length) + self.assertEqual(len(c), length+1) + for n in range(length): + i, j = c[n:n+2] + self.assertEqual(list(indices[i:j]), [n] * (j-i)) + + def test_not_monotonic(self): + with self.assertRaisesRegex(ValueError, 'indices are not monotomically increasing'): + numeric.compress_indices([0, 2, 1], 3) + + def test_out_of_bounds(self): + length = 2 + for indices in [0, 1, 2], [-1, 0, 1]: + with self.assertRaisesRegex(ValueError, 'indices are out of bounds'): + numeric.compress_indices(indices, length) diff --git a/tests/test_solver.py b/tests/test_solver.py index 553969b17..110c97db3 100644 --- a/tests/test_solver.py +++ b/tests/test_solver.py @@ -45,19 +45,21 @@ def _test_recursion_cache(testcase, solver_iter): with tmpcache(): for i, length in enumerate(lengths): with testcase.subTest(lengths=lengths, step=i): - if length == 0: # treated as special case because assertLogs would fail - read(0) - continue - with testcase.assertLogs('nutils', 'DEBUG') as cm: - v = read(length) + if length <= 1: + if hasattr(testcase, 'assertNoLogs'): # Python >= 3.10 + with testcase.assertNoLogs('nutils', 'DEBUG'): + v = read(length) + else: + v = read(length) + else: + with testcase.assertLogs('nutils', 'DEBUG') as cm: + v = read(length) + testcase.assertRegex('\n'.join(cm.output), '\\[cache\\.function [0-9a-f]{40}\\] ' + + ('load' if i and max(lengths[:i]) > 1 else 'failed to load')) _assert_almost_equal(testcase, v, reference[:length]) - testcase.assertRegex('\n'.join(cm.output), '\\[cache\\.Recursion [0-9a-f]{40}\\] start iterating') - testcase.assertRegex('\n'.join(cm.output), '\\[cache\\.Recursion [0-9a-f]{40}\\.0000\\] load' if i and max(lengths[:i]) > 0 - else '\\[cache\\.Recursion [0-9a-f]{40}\\.0000\\] cache exhausted') def _test_solve_cache(testcase, solver_gen): - _test_recursion_cache(testcase, solver_gen) with testcase.subTest('solve'), tmpcache(): v1 = _edit(solver_gen().solve(1e-5)) with testcase.assertLogs('nutils', 'DEBUG') as cm: @@ -263,16 +265,6 @@ def test_linear(self): cons = solver.optimize('dofs', err, droptol=1e-15) numpy.testing.assert_almost_equal(cons, numpy.take([1, numpy.nan], [0, 1, 1, 0, 1, 1, 0, 1, 1]), decimal=15) - def test_nonlinear(self): - err = self.domain.boundary['bottom'].integral('(u + .25 u^3 - 1.25)^2' @ self.ns, degree=6) - cons = solver.optimize('dofs', err, droptol=1e-15, tol=1e-15) - numpy.testing.assert_almost_equal(cons, numpy.take([1, numpy.nan], [0, 1, 1, 0, 1, 1, 0, 1, 1]), decimal=15) - - def test_nonlinear_multipleroots(self): - err = self.domain.boundary['bottom'].integral('(u + u^2 - .75)^2' @ self.ns, degree=2) - cons = solver.optimize('dofs', err, droptol=1e-15, lhs0=numpy.ones(len(self.ubasis)), tol=1e-10) - numpy.testing.assert_almost_equal(cons, numpy.take([.5, numpy.nan], [0, 1, 1, 0, 1, 1, 0, 1, 1]), decimal=15) - def test_nanres(self): err = self.domain.integral('(sqrt(1 - u) - .5)^2' @ self.ns, degree=2) dofs = solver.optimize('dofs', err, tol=1e-10) @@ -280,12 +272,8 @@ def test_nanres(self): def test_unknowntarget(self): err = self.domain.integral('(sqrt(1 - u) - .5)^2' @ self.ns, degree=2) - with self.assertRaises(ValueError): + with self.assertRaises(KeyError): dofs = solver.optimize(['dofs', 'other'], err, tol=1e-10) - dofs = solver.optimize(['dofs', 'other'], err, tol=1e-10, droptol=1e-10) - self.assertEqual(list(dofs), ['dofs']) - dofs = solver.optimize(['other'], err, tol=1e-10, droptol=1e-10) - self.assertEqual(dofs, {}) with self.assertRaises(KeyError): dofs = solver.optimize('other', err, tol=1e-10, droptol=1e-10) @@ -331,7 +319,7 @@ def check(self, method, theta): residual = topo.integral('-e_n sin(t) dV' @ ns, degree=0) timestep = 0.1 udesired = numpy.array([0.]) - uactualiter = iter(method(target='u', residual=residual, inertia=inertia, timestep=timestep, lhs0=udesired, timetarget='t')) + uactualiter = iter(method(target='u', residual=residual, inertia=inertia, timestep=timestep, lhs0=udesired.copy(), timetarget='t')) for i in range(5): with self.subTest(i=i): uactual = next(uactualiter) @@ -343,3 +331,117 @@ def test_impliciteuler(self): def test_cranknicolson(self): self.check(solver.cranknicolson, theta=0.5) + + +class system_finitestrain(TestCase): + + def setUp(self): + super().setUp() + domain, geom = mesh.rectilinear([numpy.linspace(0, 1, 9)] * 2) + u = function.dotarg('u', domain.basis('std', degree=2), shape=(2,)) + Geom = geom * [1.1, 1] + u + self.cons = solver.optimize('u,', domain.boundary['left,right'].integral((u**2).sum(0), degree=4), droptol=1e-15) + strain = .5 * (function.outer(Geom.grad(geom), axis=1).sum(0) - function.eye(2)) + energy = domain.integral(((strain**2).sum([0, 1]) + 20*(numpy.linalg.det(Geom.grad(geom))-1)**2)*function.J(geom), degree=6) + self.system = solver.System(energy, trial='u') + self.residual = evaluable.compile(energy.derivative('u').as_evaluable_array) + + def assert_resnorm(self, args, tol): + resnorm = numpy.linalg.norm(self.residual(**args)[numpy.isnan(self.cons['u'])]) + self.assertLess(resnorm, tol) + + def test_direct(self): + with self.assertRaises(ValueError): + self.system.solve(constrain=self.cons) + + def test_newton(self): + args = self.system.solve(constrain=self.cons, method=solver.Newton(), tol=1e-10, maxiter=7) + self.assert_resnorm(args, tol=1e-10) + + def test_newton_linesearch(self): + args = self.system.solve(constrain=self.cons, method=solver.LinesearchNewton(), tol=1e-10, maxiter=7) + self.assert_resnorm(args, tol=1e-10) + + def test_minimize(self): + args = self.system.solve(constrain=self.cons, method=solver.Minimize(), tol=1e-10, maxiter=13) + self.assert_resnorm(args, tol=1e-10) + + +class system_navierstokes(TestCase): + + def setUp(self): + super().setUp() + viscosity = 1e-3 + domain, geom = mesh.rectilinear([numpy.linspace(0, 1, 9)] * 2) + gauss = domain.sample('gauss', 5) + uin = geom[1] * (1-geom[1]) + dx = function.J(geom) + ubasis = domain.basis('std', degree=2) + pbasis = domain.basis('std', degree=1) + u = function.dotarg('u', ubasis, shape=(2,)) + v = function.dotarg('v', ubasis, shape=(2,)) + p = function.dotarg('p', pbasis) + q = function.dotarg('q', pbasis) + self.cons = solver.optimize('u,', domain.boundary['top,bottom'].integral((u**2).sum(), degree=4), droptol=1e-10) + self.cons = solver.optimize('u,', domain.boundary['left'].integral((u[0]-uin)**2 + u[1]**2, degree=4), droptol=1e-10, constrain=self.cons) + res = gauss.integral((viscosity * numpy.einsum('ij,ij', v.grad(geom), u.grad(geom) + u.grad(geom).T) - v.div(geom) * p + q * u.div(geom)) * dx) + self.arguments = solver.solve_linear('u:v,p:q', residual=res, constrain=self.cons) + res += gauss.integral(numpy.einsum('i,ij,j', v, u.grad(geom), u) * dx) + self.system = solver.System(res, trial='u,p', test='v,q') + self.inertia = gauss.integral(.5 * (u**2).sum(-1) * dx).derivative('u'), numpy.zeros(pbasis.shape) + self.residuals = evaluable.compile((res.derivative('v').as_evaluable_array, res.derivative('q').as_evaluable_array)) + + def assert_resnorm(self, args, tol): + ures, pres = self.residuals(**args) + resnorm = numpy.linalg.norm(numpy.concatenate([ures[numpy.isnan(self.cons['u'])], pres])) + self.assertLess(resnorm, tol) + + def test_direct(self): + with self.assertRaises(ValueError): + self.system.solve(constrain=self.cons) + + def test_newton(self): + args = self.system.solve(arguments=self.arguments, constrain=self.cons, method=solver.Newton(), tol=1e-10, maxiter=3) + self.assert_resnorm(args, tol=1e-10) + + def test_newton_linesearch_normbased(self): + args = self.system.solve(arguments=self.arguments, constrain=self.cons, method=solver.LinesearchNewton(strategy=solver.NormBased()), tol=1e-10, maxiter=3) + self.assert_resnorm(args, tol=1e-10) + + def test_newton_linesearch_medianbased(self): + args = self.system.solve(arguments=self.arguments, constrain=self.cons, method=solver.LinesearchNewton(strategy=solver.MedianBased()), tol=1e-10, maxiter=3) + self.assert_resnorm(args, tol=1e-10) + + def test_newton_tolnotreached(self): + with self.assertLogs('nutils', logging.WARNING) as cm: + args = self.system.solve(arguments=self.arguments, constrain=self.cons, method=solver.Newton(), tol=1e-10, maxiter=3, linargs=dict(rtol=1e-99)) + for msg in cm.output: + self.assertIn('solver failed to reach tolerance', msg) + self.assert_resnorm(args, tol=1e-10) + + def test_pseudotime(self): + args = self.system.solve(arguments=self.arguments, constrain=self.cons, method=solver.Pseudotime(inertia=self.inertia, timestep=1), tol=1e-10, maxiter=6) + self.assert_resnorm(args, tol=1e-10) + + +class system_burgers(TestCase): + + def setUp(self): + super().setUp() + domain, geom = mesh.rectilinear([10], periodic=(0,)) + basis = domain.basis('discont', degree=1) + ns = Namespace() + ns.x = geom + ns.define_for('x', gradient='∇', normal='n', jacobians=('dV', 'dS')) + ns.add_field(('u', 'u0', 'v'), basis) + ns.add_field(('t', 't0')) + ns.dudt = '(u - u0) / (t - t0)' + ns.f = '.5 u^2' + residual = domain.integral('(v dudt - ∇_0(v) f) dV' @ ns, degree=2) + residual += domain.interfaces.integral('-[v] n_0 ({f} - .5 [u] n_0) dS' @ ns, degree=4) + self.arguments = {'u': numpy.sin(numpy.arange(len(basis)))} # "random" initial vector + self.system = solver.System(residual, trial='u', test='v') + + def test_step(self): + args = self.system.step(timestep=100, timetarget='t', historysuffix='0', arguments=self.arguments, method=solver.LinesearchNewton(), tol=1e-10) + self.assertAlmostEqual64(args['u'], 'eNpzNBA1NjHuNHQ3FDsTfCbAuNz4nUGZgeyZiDOZxlONmQwU9W3OFJ/pNQAADZIOPA==')