Adaptive time step for linear solver

General
1

Hello,

I am running the diffusion problem below with a Newton solver, increasing time step over time (example 1). I am looking for a way to similarly adapt the time step for a linear solver (based on either number of iterations, or convergence, or tolerance… just some measure of time step suitability). I have implemented the linear problem (example 2) but cannot find a way to access solver information that would inform my time step decisions. Any suggestion as to how I could go about this?

Many thanks,

example 1:


L = 1
W = 0.2
diffusion_time_step=0.01
end_time=0.5
simu_start_time=0
max_dif_tstep=0.2


import numpy as np
import ufl

from mpi4py import MPI
from petsc4py.PETSc import ScalarType

from dolfinx import mesh, fem, plot, io, nls
from dolfinx.fem import FunctionSpace,VectorFunctionSpace, Function,TensorFunctionSpace, Constant
from ufl import TestFunction, TrialFunction, dx, inner,grad, Measure,derivative,system
from petsc4py import PETSc
import time

startT = time.time()

domain = mesh.create_box(MPI.COMM_WORLD, [np.array([0,0,0]), np.array([L, W, W])],
                  [20,20,20], cell_type=mesh.CellType.hexahedron)

V=VectorFunctionSpace(domain, ("CG", 1), dim=3)


p = TestFunction(V)
prc_prev = Function(V)
prc = Function(V)            
prc_mid = 0.5*(prc + prc_prev)

dt=Constant(domain, diffusion_time_step)

def set_time_step(dt,converged,n):
      if converged and n<5:
            new_dt=min(max_dif_tstep,2*dt) 
      else: new_dt=dt
      return new_dt

D_p = Constant(domain,2.4e3) # Diffusion coefficient for PTHrP
D_pr = Constant(domain,2.4e-15) # Diffusion coefficient for PTHrP
kap = Constant(domain,1e5) # Binding association rate for PTHrP-receptor
kdp = Constant(domain,1.86e-3)  # Disassociation rate for PTHrP-receptor 
psi_p = Constant(domain,0.01) # Decay rate for PTHrP
psi_pr = Constant(domain,1e-50) # Decay rate for PTHrP receptor
psi_cp = Constant(domain,1e-15) # Decay rate for PTHrP-receptor Complex


xdmf = io.XDMFFile(domain.comm, "MiniExampleN.xdmf", "w")
xdmf.write_mesh(domain)

# Supply term 
def sup_func(x):
    return np.exp(-100*(x[0]%0.1+x[1]%0.1+x[2]%0.1))
p_i = Function(V.sub(0).collapse()[0]) 
p_i.interpolate(sup_func)
p_sup = p_i

# PTHrP receptor presentation
def rec_func(x):
    return np.exp(-10*(x[0]%0.15+x[1]%0.15+x[2]%0.15))
rp_i = Function(V.sub(0).collapse()[0]) 
rp_i.interpolate(rec_func) 

subspace,submap=V.sub(1).collapse()
prc.x.array[submap]=rp_i.x.array
prc_prev.x.array[submap]=rp_i.x.array

# Variational problem
F = (
                (- (prc[0] - prc_prev[0])*p[0]
                - (prc[1] - prc_prev[1])*p[1]
                - (prc[2] - prc_prev[2])*p[2])*dx()
            + (- dt*inner(D_p*grad(prc_mid[0]), grad(p[0])))*dx()
            + (+ dt*p_sup*p[0])*dx()
            + (+ dt*(- kap*prc_prev[0]*prc_mid[1]*p[0]
                    - kap*prc_mid[0]*prc_prev[1]*p[1]
                    + kap*prc_prev[0]*prc_prev[1]*p[2]
                    + kdp*prc_prev[2]*p[0]
                    + kdp*prc_prev[2]*p[1]
                    - kdp*prc_mid[2]*p[2]))*dx()
            + (+ dt*(- psi_p*prc_mid[0]*p[0]
                    - psi_pr*prc_mid[1]*p[1]
                    - psi_cp*prc_mid[2]*p[2]))*dx()
            )

# Create boundary condition
tdim = V.mesh.topology.dim
fdim = tdim - 1
domain.topology.create_connectivity(fdim, tdim)
boundary_facets = mesh.exterior_facet_indices(V.mesh.topology)
fixed_dofs_0 = fem.locate_dofs_topological(V.sub(0), fdim, boundary_facets)
fixed_dofs_1 = fem.locate_dofs_topological(V.sub(1), fdim, boundary_facets)
fixed_dofs_2 = fem.locate_dofs_topological(V.sub(2), fdim, boundary_facets)
bc0 = fem.dirichletbc(PETSc.ScalarType(0), fixed_dofs_0, V.sub(0))
bc1 = fem.dirichletbc(PETSc.ScalarType(0), fixed_dofs_1, V.sub(1))
bc2 = fem.dirichletbc(PETSc.ScalarType(0), fixed_dofs_2, V.sub(2))
        
# Define problem       
problem=fem.petsc.NonlinearProblem(F, prc, bcs=[bc0,bc1,bc2]) 
solver = nls.petsc.NewtonSolver(V.mesh.comm, problem)
solver.convergence_criterion = "incremental"
solver.atol=1e-8
solver.rtol=1e-8
# solver.max_it=25
# solver.ksp_type='cg'
# solver.krylov_solver.atol=1e-4
# solver.krylov_solver.rtol=1e-5


# Solve problem
simu_time=simu_start_time+0.0001
while simu_time < end_time:
            n, converged = solver.solve(prc)
            prc.x.scatter_forward()
            prc_prev.x.array[:]=prc.x.array
            if(simu_time%0.2<0.0001):
                xdmf.write_function(prc, simu_time)
            dt.value=set_time_step(dt.value,converged,n)
            print(dt.value,flush=True)
            simu_time += dt.value

xdmf.close()
endT = time.time()
print("The time of execution of above program is :",(endT-startT),flush=True)

example 2:


L = 1
W = 0.2
diffusion_time_step=0.05
end_time=1
simu_start_time=0
max_dif_tstep=0.05


import numpy as np
import ufl

from mpi4py import MPI
from petsc4py.PETSc import ScalarType

from dolfinx import mesh, fem, plot, io
from dolfinx.fem import FunctionSpace,VectorFunctionSpace, Function,TensorFunctionSpace, Constant
from ufl import TestFunction, TrialFunction, dx, inner,grad, Measure,derivative,system
from petsc4py import PETSc
import time

startT = time.time()

domain = mesh.create_box(MPI.COMM_WORLD, [np.array([0,0,0]), np.array([L, W, W])],
                  [20,20,20], cell_type=mesh.CellType.hexahedron)

V=VectorFunctionSpace(domain, ("CG", 1), dim=3)

xdmf = io.XDMFFile(domain.comm, "MiniExampleL.xdmf", "w")
xdmf.write_mesh(domain)

p = TestFunction(V)
prc_prev = Function(V)
prc = TrialFunction(V)            
prc_mid = 0.5*(prc + prc_prev)
prco = Function(V) 
prc_uh= Function(V)

def set_time_step(dt,converged,n):
      if converged and n<5:
            new_dt=min(max_dif_tstep,2*dt) 
      else: new_dt=dt
      return new_dt

dt=Constant(domain, diffusion_time_step)

D_p = Constant(domain,2.4e3) # Diffusion coefficient for PTHrP
D_pr = Constant(domain,2.4e-15) # Diffusion coefficient for PTHrP
kap = Constant(domain,1e5) # Binding association rate for PTHrP-receptor
kdp = Constant(domain,1.86e-3)  # Disassociation rate for PTHrP-receptor 
psi_p = Constant(domain,0.01) # Decay rate for PTHrP
psi_pr = Constant(domain,1e-50) # Decay rate for PTHrP receptor
psi_cp = Constant(domain,1e-15) # Decay rate for PTHrP-receptor Complex

# Supply term 
def sup_func(x):
    return np.exp(-100*(x[0]%0.1+x[1]%0.1+x[2]%0.1))
p_i = Function(V.sub(0).collapse()[0]) 
p_i.interpolate(sup_func)
p_sup = p_i

# PTHrP receptor presentation
def rec_func(x):
    return np.exp(-10*(x[0]%0.15+x[1]%0.15+x[2]%0.15))
rp_i = Function(V.sub(0).collapse()[0]) 
rp_i.interpolate(rec_func) 

subspace,submap=V.sub(1).collapse()
prco.x.array[submap]=rp_i.x.array
prc_prev.x.array[submap]=rp_i.x.array

# Variational problem
F = (
                (- (prc[0] - prc_prev[0])*p[0]
                - (prc[1] - prc_prev[1])*p[1]
                - (prc[2] - prc_prev[2])*p[2])*dx()
            + (- dt*inner(D_p*grad(prc_mid[0]), grad(p[0])))*dx()
            + (+ dt*p_sup*p[0])*dx()
            + (+ dt*(- kap*prc_prev[0]*prc_mid[1]*p[0]
                    - kap*prc_mid[0]*prc_prev[1]*p[1]
                    + kap*prc_prev[0]*prc_prev[1]*p[2]
                    + kdp*prc_prev[2]*p[0]
                    + kdp*prc_prev[2]*p[1]
                    - kdp*prc_mid[2]*p[2]))*dx()
            + (+ dt*(- psi_p*prc_mid[0]*p[0]
                    - psi_pr*prc_mid[1]*p[1]
                    - psi_cp*prc_mid[2]*p[2]))*dx()
            )

# Create boundary condition
tdim = V.mesh.topology.dim
fdim = tdim - 1
domain.topology.create_connectivity(fdim, tdim)
boundary_facets = mesh.exterior_facet_indices(V.mesh.topology)
fixed_dofs_0 = fem.locate_dofs_topological(V.sub(0), fdim, boundary_facets)
fixed_dofs_1 = fem.locate_dofs_topological(V.sub(1), fdim, boundary_facets)
fixed_dofs_2 = fem.locate_dofs_topological(V.sub(2), fdim, boundary_facets)
bc0 = fem.dirichletbc(PETSc.ScalarType(0), fixed_dofs_0, V.sub(0))
bc1 = fem.dirichletbc(PETSc.ScalarType(0), fixed_dofs_1, V.sub(1))
bc2 = fem.dirichletbc(PETSc.ScalarType(0), fixed_dofs_2, V.sub(2))
        
# Define problem       
(a, L) = system(F)
problem=fem.petsc.LinearProblem(a, L, bcs=[bc0,bc1,bc2],petsc_options={"ksp_type": "cg","pc_type": "hypre","pc_hypre_type": "boomeramg", "ksp_rtol":1e-8, "ksp_atol":1e-8}) 



# Solve problem
simu_time=simu_start_time
while simu_time < end_time:
            prc = problem.solve()
            prc.x.scatter_forward()
            prc_prev.x.array[:]=prc.x.array
            prco.x.array[:]=prc.x.array
            if(simu_time%0.2<0.01 or (0.2-simu_time%0.2)<0.01):
                xdmf.write_function(prco, simu_time)
            simu_time += dt.value
            # print(simu_time%0.5,flush=True)
           

xdmf.close()
endT = time.time()
print("The time of execution of above program is :",(endT-startT),flush=True)
2

I would suggest making a custom Newton solver, as for instance described at: Custom Newton solvers — FEniCSx tutorial

3

Thank you, this seems to be working well.
It is a bit slower than the ‘LinearProblem’ implementation though.

If this custom solver associated with an adaptive time step is to be used in parallel, is it likely to cause issues, such as different partitions being solved at different time points? Considering I am storing the final output only, I would think it does not matter too much how many intermediate steps each partition needed to reach it but I might be missing something?

4

LinearProblem is to solve linear problems (Ax=b)

A newton solver solves F(u)=0 by solving Jdu=-F(u_k), u_k+1=u_k+du

A NewtonSolver would converge in 1 iteration if it encounters a LinearProblem (if you use the steepest descent direction), and thus could achieve similar performance if you set the underlying ksp solver with the same parameters.

You should use a synced time step for all processes.
Each solve command solves the distributed system, not one problem per partition (to do so you would have to implement a domain decomposition method).

As Im not sure what you want to change the step, i cannot give specific suggestions as to what should be sync with for instance an allgather/allreduce

5

Thank you for the insight, I will investigate this further.