Parallelization error when customizing NewtonSolver

Legacy DOLFIN
1

Hi everyone,

I’m new to Fenics and have literally no knowledge to openmpi. I am trying to run the cahn-hillilard demo with a customized NewtonSolver which tries to clamp c within [0,1]. Here’s the customized solver:

class CustomSolver(NewtonSolver):
    def __init__(self):
        NewtonSolver.__init__(self)
    def _clip(self, x):
        MPI.barrier(comm)
        if rank == 0:
            _size = int(round(x.size()/2))
            print(_size)
            _c = np.reshape(x[0::2], (_size, 1)) # even index in x is c
            _mu = np.reshape(x[1::2], (_size, 1)) # od index in x is mu
            _tol = 1e-6
            _c = _c.clip(max=1.0-_tol, min=_tol)
            _x_clipped = np.concatenate((_c, _mu), axis=1).flatten()
            x = Vector(MPI.comm_self,_c.size*2)
            x.set_local(_x_clipped)
            return x
    def update_solution(self, x, dx, relaxation_parameter, nonlinear_problem,
                        iteration):
        x.axpy(-0.01, dx)
        x = self._clip(x)

The code runs without error when using only 1 core, but I got error when using 4 cores with mpirun, basically the size of x is wrong (basically 1/4 of the true size). I guess this is because the x array is now on 4 cores. How can I clamp the x array (or more precisely the c array that is part of x) when using mpirun? Thanks in advance!

2

Without a minimal reproducible example, it is unlikely that anyone will have the time/bandwidth to help you. I would suggest making a simple example, for instance by drawing inspiration from
https://jsdokken.com/dolfinx-tutorial/chapter4/newton-solver.html

3

Hi dokken, thanks for your quick reply and sorry for not stating the issue properly. Here’s the minimal work example, basically copied from demo-chan-hillilard.py from https://fenicsproject.org/olddocs/dolfin/2019.1.0/python/demos/cahn-hilliard/demo_cahn-hilliard.py.html

import random
from dolfin import *

# Class representing the intial conditions
class InitialConditions(UserExpression):
    def __init__(self, **kwargs):
        random.seed(2 + MPI.rank(MPI.comm_world))
        super().__init__(**kwargs)
    def eval(self, values, x):
        values[0] = 0.63 + 0.02*(0.5 - random.random())
        values[1] = 0.0
    def value_shape(self):
        return (2,)


# Class for interfacing with the Newton solver
class CahnHilliardEquation(NonlinearProblem):
    def __init__(self, a, L):
        NonlinearProblem.__init__(self)
        self.L = L
        self.a = a
    def F(self, b, x):
        assemble(self.L, tensor=b)
    def J(self, A, x):
        assemble(self.a, tensor=A)


# Model parameters
lmbda  = 1.0e-02  # surface parameter
dt     = 5.0e-06  # time step
theta  = 0.5      # time stepping family, e.g. theta=1 -> backward Euler, theta=0.5 -> Crank-Nicolson

# Form compiler options
parameters["form_compiler"]["optimize"]     = True
parameters["form_compiler"]["cpp_optimize"] = True


# Create mesh and build function space
mesh = UnitSquareMesh.create(200, 200, CellType.Type.quadrilateral)
P1 = FiniteElement("Lagrange", mesh.ufl_cell(), 1)
ME = FunctionSpace(mesh, P1*P1)


# Define trial and test functions
du    = TrialFunction(ME)
q, v  = TestFunctions(ME)

# Define functions
u   = Function(ME)  # current solution
u0  = Function(ME)  # solution from previous converged step

# Split mixed functions
dc, dmu = split(du)
c,  mu  = split(u)
c0, mu0 = split(u0)


# Create intial conditions and interpolate
u_init = InitialConditions(degree=1)
u.interpolate(u_init)
u0.interpolate(u_init)



# Compute the chemical potential df/dc
c = variable(c)
f    = 100*c**2*(1-c)**2
dfdc = diff(f, c)


# mu_(n+theta)
mu_mid = (1.0-theta)*mu0 + theta*mu

# which is then used in the definition of the variational forms::

# Weak statement of the equations
L0 = c*q*dx - c0*q*dx + dt*dot(grad(mu_mid), grad(q))*dx
L1 = mu*v*dx - dfdc*v*dx - lmbda*dot(grad(c), grad(v))*dx
L = L0 + L1


# Compute directional derivative about u in the direction of du (Jacobian)
a = derivative(L, u, du)


# Create nonlinear problem and Newton solver
problem = CahnHilliardEquation(a, L)


comm = MPI.comm_world
rank = MPI.rank(comm)
import numpy as np
class CustomSolver(NewtonSolver):
    def __init__(self):
        NewtonSolver.__init__(self)
    def _clip(self, x):
        MPI.barrier(comm)
        if rank == 0:
            _size = int(round(x.size()/2))
            print(_size)
            _c = np.reshape(x[0::2], (_size, 1)) # even index in x is c
            _mu = np.reshape(x[1::2], (_size, 1)) # od index in x is mu
            _tol = 1e-6
            _c = _c.clip(max=1.0-_tol, min=_tol)
            _x_clipped = np.concatenate((_c, _mu), axis=1).flatten()
            x = Vector(MPI.comm_self,_c.size*2)
            x.set_local(_x_clipped)
            return x
    def update_solution(self, x, dx, relaxation_parameter, nonlinear_problem,
                        iteration):
        x.axpy(-0.01, dx)
        x = self._clip(x)

solver = CustomSolver()
solver.parameters["linear_solver"] = "lu"
solver.parameters["convergence_criterion"] = "incremental"
solver.parameters["relative_tolerance"] = 1e-6

# Output file
file = File("output.pvd", "compressed")

# Step in time
t = 0.0
T = 50*dt
while (t < T):
    t += dt
    u0.vector()[:] = u.vector()
    solver.solve(problem, u.vector())
    file << (u.split()[0], t)

when I type python3 demo_cahn-hilliard.py everything works fine, but when I type mpirun -np 4 demo_cahn-hilliard.py the following error pops out:

Traceback (most recent call last): File "/Users/archieyao/Desktop/test/demo_cahn-hilliard.py", line 128, in <module> solver.solve(problem, u.vector()) File "/Users/archieyao/Desktop/test/demo_cahn-hilliard.py", line 112, in update_solution x = self._clip(x) ^^^^^^^^^^^^^ File "/Users/archieyao/Desktop/test/demo_cahn-hilliard.py", line 101, in _clip _c = np.reshape(x[0::2], (_size, 1)) # even index in x is c ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/archieyao/tools/miniconda3/envs/ocv_pf/lib/python3.12/site-packages/numpy/core/fromnumeric.py", line 285, in reshape return _wrapfunc(a, 'reshape', newshape, order=order) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/archieyao/tools/miniconda3/envs/ocv_pf/lib/python3.12/site-packages/numpy/core/fromnumeric.py", line 59, in _wrapfunc return bound(*args, **kwds) ^^^^^^^^^^^^^^^^^^^^ ValueError: cannot reshape array of size 10089 into shape (40401,1)

4

First of all, you cannot do the clip on only one rank, as the vector is distributed.
You need to call this on every process, with
_size = int(round(x.local_size()/2))

5

Thanks dokken, I have already realized this is caused by the MPI thing,

I guess this is because the x array is now on 4 cores. How can I clamp the x array (or more precisely the c array that is part of x) when using mpirun?

but I’m really a green hand in MPI so I don’t know how to resolve this bug. Is there any tutorial / example code that I can have a look on calling every process in Fenics?

6

Simply change the code to:

  def _clip(self, x):
       _size = int(round(x.local_size()/2))
       print(_size)
       _c = np.reshape(x[0::2], (_size, 1)) # even index in x is c
       _mu = np.reshape(x[1::2], (_size, 1)) # od index in x is mu
       _tol = 1e-6
       _c = _c.clip(max=1.0-_tol, min=_tol)
       _x_clipped = np.concatenate((_c, _mu), axis=1).flatten()
       x.set_local(_x_clipped)
       return x

i.e. no special handling on the main process, and do not create a new x.
There are some assumptions you have on the indexing of c and mu in the sub-spaces, that I don’t believe is true for every kind of mixed space, but if it works for you I guess its fine.

1 Like
8

ah I realized it may be wrong to assume

_c = np.reshape(x[0::2], (_size, 1)) # even index in x is c
_mu = np.reshape(x[1::2], (_size, 1)) # od index in x is mu

during parallelization, because if core 0 received 10737 (or any other odd number) of x, then the portion of x on core 1 should be

_c = np.reshape(x[1::2], (_size, 1)) # odd index in x is c
_mu = np.reshape(x[0::2], (_size, 1)) # even index in x is mu

Is there a way that I can gather the solution x and do the clip on a single core?

9

Alternatively, you can also use PETSc’s bound-constrained solver such as TAO or SNES VI.

1 Like