# Navier-Stokes with a Lorentz force term

**URL:** <https://fenicsproject.discourse.group/t/navier-stokes-with-a-lorentz-force-term/10168>\
**Category:** variational formulation\
**Created:** [January 17, 2023, 4:53pm UTC](https://fenicsproject.discourse.group/t/navier-stokes-with-a-lorentz-force-term/10168 "2023-01-17T16:53:33Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![James\_Dark](https://yyz2.discourse-cdn.com/free1/user_avatar/fenicsproject.discourse.group/james_dark/32/2761_2.png) [@James\_Dark](https://fenicsproject.discourse.group/u/James_Dark)\
**Post date:** [January 17, 2023, 4:53pm UTC](https://fenicsproject.discourse.group/t/navier-stokes-with-a-lorentz-force-term/10168/1 "2023-01-17T16:53:33Z")

</div>

Hi fenics experts!

I’m trying to model the stationary Navier-stokes equations with a Lorentz force term. I was wondering if I have the correct variational formulation here and if you have any advice for running the simulation more efficiently as it is currently very computationally intensive.

Governing equations:

\rho(\mathbf{u}\cdot \nabla\mathbf{u}) = -\nabla p + \mu\nabla^{2}\mathbf{u} + \mathbf{J} \times \mathbf{B}

\nabla \cdot \mathbf{u} = 0

Where:  
\mathbf{J} = \sigma(-\nabla \Phi + \mathbf{u} \times \mathbf{B})

I think I can expand the Lorentz force term as:  
\begin{align\*} \mathbf{J} \times \mathbf{B} &= -\mathbf{B} \times \sigma (-\nabla \Phi + \mathbf{u} \times \mathbf{B})\\ &= \sigma[(-\mathbf{B} \times -\nabla \Phi) + (-\mathbf{B} \times (\mathbf{u} \times \mathbf{B})]\\ &= \sigma[(\mathbf{B} \times \nabla \Phi) -(\mathbf{B} \cdot \mathbf{B})\mathbf{u} + (\mathbf{B} \cdot \mathbf{u})\mathbf{B}] \end{align\*}

Additionally, assuming the fluid is electrically neutral thus, \nabla \cdot \mathbf{J} = 0. This leads to an equation for the electric potential \Phi:

\nabla^{2}\Phi = \nabla \cdot (\mathbf{u} \times \mathbf{B})

Assuming Dirichlet BC’s on all walls of the domain for the velocity and pressure, and assuming electrically insulated walls (\partial \phi / \partial n=0). I think the weak formulation is as follows for the momentum:

\int\_{\Omega} \rho (\mathbf{u}\cdot \nabla\mathbf{u}) \cdot \mathbf{v} - \int\_{\Omega} p\nabla\cdot\mathbf{v} + \int\_{\Omega} \mu\nabla\mathbf{u}\cdot\nabla\mathbf{v} - \\ \ \sigma \left(\int\_{\Omega} (\mathbf{B} \times \nabla \Phi) \cdot \mathbf{v} + \int\_{\Omega} (\mathbf{B} \cdot \mathbf{B})\mathbf{u} \cdot \mathbf{v} - \int\_{\Omega} (\mathbf{B} \cdot \mathbf{u})\mathbf{B} \cdot \mathbf{v} \right) = 0 \ \ \ \forall \mathbf{v} \in V

for continuity:

\int\_{\Omega} q (\nabla \cdot \mathbf{u}) = 0 \ \ \ \forall \mathbf{q} \in Q

and for the electrical potential:

\int\_{\Omega} \nabla\mathbf{\phi} \cdot \nabla \mathbf{q2} - \int\_{\Omega} ((\nabla \times \mathbf{u}) \cdot \mathbf{B}) \cdot \mathbf{q2} + \int\_{\Omega} ((\nabla \times \mathbf{B}) \cdot \mathbf{u}) \cdot \mathbf{q2} = 0 \ \ \ \forall \mathbf{q2} \in Q2

I have implemented it in the following example:

Would love to have your feedback and apologies for the length of the post

````python
from fenics import *

N = 15
mesh = UnitCubeMesh(N, N, N)
V_ele = VectorElement("CG", mesh.ufl_cell(), 3)
Q_ele = FiniteElement("CG", mesh.ufl_cell(), 1)
W = FunctionSpace(mesh, MixedElement([V_ele, Q_ele, Q_ele]))

fixed_wall_locations = "near(x[1], 0.) | near(x[1], 1.) | near(x[2], 0.) | near(x[2], 1.)"
inlet = "near(x[0], 0.)"
outlet = "near(x[0], 1.)"
velocity_inlet = DirichletBC(W.sub(0), Constant((10, 0, 0)), inlet)
pressure_outlet = DirichletBC(W.sub(1), Constant(0), outlet)
noslip = DirichletBC(W.sub(0), (0, 0, 0), fixed_wall_locations)
bcu = [velocity_inlet, pressure_outlet, noslip]

func = Function(W)
u, p, phi = split(func)
v, q, q_2 = TestFunctions(W)
B = Constant((0, -1, 0))
sigma = Constant(1)

# Momentum
F = (
inner(dot(grad(u), u), v) * dx
- inner(p, div(v)) * dx
+ inner(grad(u), grad(v)) * dx
+ sigma * (inner(cross(B, grad(phi)), v) * dx
+ inner(u * dot(B, B), v) * dx
- inner(B * dot(B, u), v) * dx)
)

# CFD continuity
F += - inner(q, div(u)) * dx

# electric continuity
F += (inner(grad(phi), grad(q_2)) * dx
- inner(dot(B, curl(u)) + dot(u, curl(B)), q_2) * dx
)

solve(F == 0, func, bcu, solver_parameters={'newton_solver': {'linear_solver': 'mumps'}})```
````

---

<div class="post-metadata">

**Author:** ![nate](https://yyz2.discourse-cdn.com/free1/user_avatar/fenicsproject.discourse.group/nate/32/17_2.png) [@nate](https://fenicsproject.discourse.group/u/nate)\
**Post date:** [January 17, 2023, 4:56pm UTC](https://fenicsproject.discourse.group/t/navier-stokes-with-a-lorentz-force-term/10168/2 "2023-01-17T16:56:30Z")

</div>

> [@James\_Dark](#):
>
> ```auto
> V_ele = VectorElement("CG", mesh.ufl_cell(), 3)
> Q_ele = FiniteElement("CG", mesh.ufl_cell(), 1)
> 
> ```

Is this a typo? Your element pair for velocity-pressure doesn’t appear stable.

Regarding computational expense, large 3D problems typically benefit from iterative solvers. Mumps will not scale optimally. There’s a vast literature on the topic of preconditioning the Navier-Stokes system. Preconditioning your Poisson type problem for the electric potential should be trivial. I highly recommend `dolfinx` for block assembly facilitating preconditioning the individual blocks and tying them together with whichever clever method you like. See for example the [Stokes demo](https://github.com/FEniCS/dolfinx/blob/main/python/demo/demo_stokes.py#L161-L259).

---

<div class="post-metadata">

**Author:** ![James\_Dark](https://yyz2.discourse-cdn.com/free1/user_avatar/fenicsproject.discourse.group/james_dark/32/2761_2.png) [@James\_Dark](https://fenicsproject.discourse.group/u/James_Dark)\
**Post date:** [January 17, 2023, 5:09pm UTC](https://fenicsproject.discourse.group/t/navier-stokes-with-a-lorentz-force-term/10168/3 "2023-01-17T17:09:53Z")

</div>

Thanks so much for getting back nate!

Is there a more optimal element-pair that would be better to use?

Yeah, I have read a lot about preconditioning, but I’m not sure how to implement it. I actually haven’t used fenicsx either yet, perhaps this would be a good example to get started.

---

<div class="post-metadata">

**Author:** ![nate](https://yyz2.discourse-cdn.com/free1/user_avatar/fenicsproject.discourse.group/nate/32/17_2.png) [@nate](https://fenicsproject.discourse.group/u/nate)\
**Post date:** [January 17, 2023, 5:11pm UTC](https://fenicsproject.discourse.group/t/navier-stokes-with-a-lorentz-force-term/10168/4 "2023-01-17T17:11:58Z")

</div>

I would have expected the standard Taylor-Hood pair, e.g.

```auto
V_ele = VectorElement("CG", mesh.ufl_cell(), 2)
Q_ele = FiniteElement("CG", mesh.ufl_cell(), 1)

```

If you can’t solve with this pair there’s clearly something I’m missing from trying to understand your code.

---

<div class="post-metadata">

**Author:** ![James\_Dark](https://yyz2.discourse-cdn.com/free1/user_avatar/fenicsproject.discourse.group/james_dark/32/2761_2.png) [@James\_Dark](https://fenicsproject.discourse.group/u/James_Dark)\
**Post date:** [January 19, 2023, 1:56pm UTC](https://fenicsproject.discourse.group/t/navier-stokes-with-a-lorentz-force-term/10168/5 "2023-01-19T13:56:22Z")

</div>

Hi Nate!

Yes, the standard Taylor-Hood pair works much better, although I have to increase the relative tolerance rather high (1e-02) for the model to converge in all cases. Not sure how normal this is.

I have also started adapting the stokes demo, first converting it to Navier-Stokes before trying to implement the MHD effects too however I’m having some issues.

I have the variational form as

```python
a = form([[inner(grad(u), grad(v)) * dx, inner(dot(grad(u), u), v) * dx, inner(p, div(v)) * dx],
          [inner(div(u), q) * dx, None]])

```

However, I think I have the form wrong as I’m getting errors.

Heres the code used and error message:

```python
import numpy as np

import ufl
from dolfinx import cpp as _cpp
from dolfinx import fem
from dolfinx.fem import (Constant, Function, FunctionSpace, dirichletbc,
                         extract_function_spaces, form,
                         locate_dofs_geometrical, locate_dofs_topological)
from dolfinx.io import XDMFFile
from dolfinx.mesh import (CellType, GhostMode, create_rectangle,
                          locate_entities_boundary)
from ufl import div, dx, grad, inner, dot, nabla_grad

from mpi4py import MPI
from petsc4py import PETSc

# Create mesh
msh = create_rectangle(MPI.COMM_WORLD,
                       [np.array([0, 0]), np.array([1, 1])],
                       [100, 100],
                       CellType.triangle, GhostMode.none)

def noslip_boundary(x):
    return np.logical_or(np.logical_or(np.isclose(x[0], 0.0),
                                       np.isclose(x[0], 1.0)),
                         np.isclose(x[1], 0.0))

def lid(x):
    return np.isclose(x[1], 1.0)

def lid_velocity_expression(x):
    return np.stack((np.ones(x.shape[1]), np.zeros(x.shape[1])))

P2 = ufl.VectorElement("Lagrange", msh.ufl_cell(), 2)
P1 = ufl.FiniteElement("Lagrange", msh.ufl_cell(), 1)
V, Q = FunctionSpace(msh, P2), FunctionSpace(msh, P1)

# boundary conditions:
noslip = np.zeros(msh.geometry.dim, dtype=PETSc.ScalarType)
facets = locate_entities_boundary(msh, 1, noslip_boundary)
bc0 = dirichletbc(noslip, locate_dofs_topological(V, 1, facets), V)

lid_velocity = Function(V)
lid_velocity.interpolate(lid_velocity_expression)
facets = locate_entities_boundary(msh, 1, lid)
bc1 = dirichletbc(lid_velocity, locate_dofs_topological(V, 1, facets))

bcs = [bc0, bc1]

# Define variational problem
(u, p) = ufl.TrialFunction(V), ufl.TrialFunction(Q)
(v, q) = ufl.TestFunction(V), ufl.TestFunction(Q)
f = Constant(msh, (PETSc.ScalarType(0), PETSc.ScalarType(0)))

a = form([[inner(grad(u), grad(v)) * dx, inner(dot(grad(u), u), v) * dx, inner(p, div(v)) * dx],
          [inner(div(u), q) * dx, None]])

L = form([inner(f, v) * dx, inner(Constant(msh, PETSc.ScalarType(0)), q) * dx])

# block-diagonal preconditioner:
a_p11 = form(inner(p, q) * dx)
a_p = [[a[0][0], None],
       [None, a_p11]]

# ### Monolithic block iterative solver
A = fem.petsc.assemble_matrix_block(a, bcs=bcs)
A.assemble()

P = fem.petsc.assemble_matrix_block(a_p, bcs=bcs)
P.assemble()
b = fem.petsc.assemble_vector_block(L, a, bcs=bcs)

# Set near nullspace for pressure
null_vec = A.createVecLeft()
offset = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
null_vec.array[offset:] = 1.0
null_vec.normalize()
nsp = PETSc.NullSpace().create(vectors=[null_vec])
assert nsp.test(A)
A.setNullSpace(nsp)

# Build IndexSets for each field (global dof indices for each field)
V_map = V.dofmap.index_map
Q_map = Q.dofmap.index_map
offset_u = V_map.local_range[0] * V.dofmap.index_map_bs + Q_map.local_range[0]
offset_p = offset_u + V_map.size_local * V.dofmap.index_map_bs
is_u = PETSc.IS().createStride(V_map.size_local * V.dofmap.index_map_bs, offset_u, 1, comm=PETSc.COMM_SELF)
is_p = PETSc.IS().createStride(Q_map.size_local, offset_p, 1, comm=PETSc.COMM_SELF)

# Create Krylov solver
ksp = PETSc.KSP().create(msh.comm)
ksp.setOperators(A, P)
ksp.setTolerances(rtol=1e-9)
ksp.setType("minres")
ksp.getPC().setType("fieldsplit")
ksp.getPC().setFieldSplitType(PETSc.PC.CompositeType.ADDITIVE)
ksp.getPC().setFieldSplitIS(
    ("u", is_u),
    ("p", is_p))

# Configure velocity and pressure sub KSPs
ksp_u, ksp_p = ksp.getPC().getFieldSplitSubKSP()
ksp_u.setType("preonly")
ksp_u.getPC().setType("gamg")
ksp_p.setType("preonly")
ksp_p.getPC().setType("jacobi")

# Monitor the convergence of the KSP
opts = PETSc.Options()
opts["ksp_monitor"] = None
opts["ksp_view"] = None
ksp.setFromOptions()

# Compute solution
x = A.createVecRight()
ksp.solve(b, x)

```

and error:

```auto
Traceback (most recent call last):
  File "/home/fenics/shared/test.py", line 59, in <module>
    a = form([[inner(grad(u), grad(v)) * dx, inner(dot(grad(u), u), v) * dx, inner(p, div(v)) * dx],
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/fem/forms.py", line 166, in form
    return _create_form(form)
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/fem/forms.py", line 163, in _create_form
    return list(map(lambda sub_form: _create_form(sub_form), form))
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/fem/forms.py", line 163, in <lambda>
    return list(map(lambda sub_form: _create_form(sub_form), form))
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/fem/forms.py", line 163, in _create_form
    return list(map(lambda sub_form: _create_form(sub_form), form))
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/fem/forms.py", line 163, in <lambda>
    return list(map(lambda sub_form: _create_form(sub_form), form))
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/fem/forms.py", line 161, in _create_form
    return _form(form)
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/fem/forms.py", line 135, in _form
    ufcx_form, module, code = jit.ffcx_jit(mesh.comm, form,
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/jit.py", line 56, in mpi_jit
    return local_jit(*args, **kwargs)
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/jit.py", line 204, in ffcx_jit
    r = ffcx.codegeneration.jit.compile_forms([ufl_object], parameters=p_ffcx, **p_jit)
  File "/usr/local/lib/python3.10/dist-packages/ffcx/codegeneration/jit.py", line 168, in compile_forms
    impl = _compile_objects(decl, forms, form_names, module_name, p, cache_dir,
  File "/usr/local/lib/python3.10/dist-packages/ffcx/codegeneration/jit.py", line 232, in _compile_objects
    _, code_body = ffcx.compiler.compile_ufl_objects(ufl_objects, prefix=module_name, parameters=parameters)
  File "/usr/local/lib/python3.10/dist-packages/ffcx/compiler.py", line 97, in compile_ufl_objects
    analysis = analyze_ufl_objects(ufl_objects, parameters)
  File "/usr/local/lib/python3.10/dist-packages/ffcx/analysis.py", line 85, in analyze_ufl_objects
    form_data = tuple(_analyze_form(form, parameters) for form in forms)
  File "/usr/local/lib/python3.10/dist-packages/ffcx/analysis.py", line 85, in <genexpr>
    form_data = tuple(_analyze_form(form, parameters) for form in forms)
  File "/usr/local/lib/python3.10/dist-packages/ffcx/analysis.py", line 168, in _analyze_form
    form_data = ufl.algorithms.compute_form_data(
  File "/usr/local/lib/python3.10/dist-packages/ufl/algorithms/compute_form_data.py", line 407, in compute_form_data
    check_form_arity(preprocessed_form, self.original_form.arguments(), complex_mode) # Currently testing how fast this is
  File "/usr/local/lib/python3.10/dist-packages/ufl/algorithms/check_arities.py", line 177, in check_form_arity
    check_integrand_arity(itg.integrand(), arguments, complex_mode)
  File "/usr/local/lib/python3.10/dist-packages/ufl/algorithms/check_arities.py", line 159, in check_integrand_arity
    arg_tuples = map_expr_dag(rules, expr, compress=False)
  File "/usr/local/lib/python3.10/dist-packages/ufl/corealg/map_dag.py", line 36, in map_expr_dag
    result, = map_expr_dags(function, [expression], compress=compress,
  File "/usr/local/lib/python3.10/dist-packages/ufl/corealg/map_dag.py", line 99, in map_expr_dags
    r = handlers[v._ufl_typecode_](v, *[vcache[u] for u in v.ufl_operands])
  File "/usr/local/lib/python3.10/dist-packages/ufl/algorithms/check_arities.py", line 63, in product
    raise ArityMismatch("Multiplying expressions with overlapping form argument number {0}, argument is {1}.".format(x[0].number(), _afmt(x)))
  File "/usr/local/lib/python3.10/dist-packages/ufl/algorithms/check_arities.py", line 19, in _afmt
    return tuple("conj({0})".format(arg) if conj else str(arg)
  File "/usr/local/lib/python3.10/dist-packages/ufl/algorithms/check_arities.py", line 20, in <genexpr>
    for arg, conj in atuple)
TypeError: cannot unpack non-iterable bool object

```

---

<div class="post-metadata">

**Author:** ![nate](https://yyz2.discourse-cdn.com/free1/user_avatar/fenicsproject.discourse.group/nate/32/17_2.png) [@nate](https://fenicsproject.discourse.group/u/nate)\
**Post date:** [January 19, 2023, 3:38pm UTC](https://fenicsproject.discourse.group/t/navier-stokes-with-a-lorentz-force-term/10168/6 "2023-01-19T15:38:25Z")

</div>

> [@James\_Dark](#):
>
> I have to increase the relative tolerance rather high (1e-02) for the model to converge in all cases. Not sure how normal this is.

If you truly mean _relative_ tolerance, this would not be acceptable for me. I would not consider the solver to have converged and assume that I have not found a suitable approximation of the true solution. See for example [Default absolute tolerance and relative tolerance - #4 by nate](https://fenicsproject.discourse.group/t/default-absolute-tolerance-and-relative-tolerance/3829/4).

> [@James\_Dark](#):
>
> I have the variational form as
> 
> ```auto
> a = form([[inner(grad(u), grad(v)) * dx, inner(dot(grad(u), u), v) * dx, inner(p, div(v)) * dx],
> [inner(div(u), q) * dx, None]])
> 
> ```

You’ve not written the block formulation correctly. The first row has three columns and the second row has two.

---

<div class="post-metadata">

**Author:** ![James\_Dark](https://yyz2.discourse-cdn.com/free1/user_avatar/fenicsproject.discourse.group/james_dark/32/2761_2.png) [@James\_Dark](https://fenicsproject.discourse.group/u/James_Dark)\
**Post date:** [January 19, 2023, 4:07pm UTC](https://fenicsproject.discourse.group/t/navier-stokes-with-a-lorentz-force-term/10168/7 "2023-01-19T16:07:53Z")

</div>

Yeah I thought the relative tolerance was too high, the solver gets stuck after a few iterations otherwise, I will experiment with alternate initial guesses and relaxation parameters and such.

regarding the variation form for a, should it be like this then:

```python
a = form([[inner(dot(grad(u), u), v) * dx, inner(grad(u), grad(v)) * dx, inner(p, div(v)) * dx],
          [inner(div(u), q) * dx, None, None]])

```

---

<div class="post-metadata">

**Author:** ![nate](https://yyz2.discourse-cdn.com/free1/user_avatar/fenicsproject.discourse.group/nate/32/17_2.png) [@nate](https://fenicsproject.discourse.group/u/nate)\
**Post date:** [January 19, 2023, 4:18pm UTC](https://fenicsproject.discourse.group/t/navier-stokes-with-a-lorentz-force-term/10168/8 "2023-01-19T16:18:14Z")

</div>

Perhaps familiarise yourself with what the block matrix representation of the discrete FEM discretisation of the weak formulation should look like. As a starting point perhaps you could take a look at [Preconditioned iterative methods for Stokes flow problems arising in computational geodynamics - ScienceDirect](https://www.sciencedirect.com/science/article/abs/pii/S003192010800191X). There may be better introductions which are less technical; however, this paper comes to mind.

---

<div class="post-metadata">

**Author:** ![James\_Dark](https://yyz2.discourse-cdn.com/free1/user_avatar/fenicsproject.discourse.group/james_dark/32/2761_2.png) [@James\_Dark](https://fenicsproject.discourse.group/u/James_Dark)\
**Post date:** [January 25, 2023, 10:37am UTC](https://fenicsproject.discourse.group/t/navier-stokes-with-a-lorentz-force-term/10168/9 "2023-01-25T10:37:14Z")

</div>

Is it possible to solve this problem in transient using a backward difference approximation?

I’ve tried to adapt the tutorial example, but not sure how to integrate the poison-type problem into the solving process? I have this so far:

```plaintext
from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
import tqdm.autonotebook

from dolfinx.fem import (
    Constant,
    Function,
    FunctionSpace,
    dirichletbc,
    form,
    locate_dofs_geometrical,
)
from dolfinx.fem.petsc import (
    assemble_matrix,
    assemble_vector,
    apply_lifting,
    create_vector,
    set_bc,
    create_matrix,
)
from dolfinx.io import XDMFFile, VTXWriter
from dolfinx.mesh import create_box, create_rectangle, create_unit_cube
from ufl import (
    FiniteElement,
    TestFunction,
    TrialFunction,
    VectorElement,
    div,
    dot,
    dx,
    inner,
    lhs,
    nabla_grad,
    rhs,
    grad,
    cross,
    curl,
)

mesh = create_box(
    MPI.COMM_WORLD, [np.array([0, 0, 0]), np.array([1, 1, 1])], [10, 10, 10]
)
t = 0
T = 1
dt = 1 / 100 # Time step size
num_steps = int(T / dt)
k = Constant(mesh, PETSc.ScalarType(dt))
mu = Constant(mesh, PETSc.ScalarType(1)) # Dynamic viscosity
rho = Constant(mesh, PETSc.ScalarType(1)) # Density

v_cg2 = VectorElement("CG", mesh.ufl_cell(), 2)
s_cg1 = FiniteElement("CG", mesh.ufl_cell(), 1)
s2_cg1 = FiniteElement("CG", mesh.ufl_cell(), 1)
V = FunctionSpace(mesh, v_cg2)
Q = FunctionSpace(mesh, s_cg1)
Q2 = FunctionSpace(mesh, s2_cg1)

def walls(x):
    walls_1 = np.logical_or(np.isclose(x[1], 0), np.isclose(x[1], 1))
    walls_2 = np.logical_or(np.isclose(x[2], 0), np.isclose(x[2], 1))
    return np.logical_or(walls_1, walls_2)

def hartmann_walls(x):
    return np.logical_or(np.isclose(x[1], 0), np.isclose(x[1], 1))

def inlet(x):
    return np.isclose(x[0], 0)

def outlet(x):
    return np.isclose(x[0], 1)

def inlet_velocity(x):
    values = np.zeros((mesh.geometry.dim, x.shape[1]), dtype=PETSc.ScalarType)
    values[0] = 1

    return values

u_noslip = np.array((0,) * mesh.geometry.dim, dtype=PETSc.ScalarType)
bc_noslip = dirichletbc(u_noslip, locate_dofs_geometrical(V, walls), V)

u_inlet = Function(V)
u_inlet.interpolate(inlet_velocity)
bc_inflow = dirichletbc(u_inlet, locate_dofs_geometrical(V, inlet))
bc_outflow = dirichletbc(PETSc.ScalarType(0), locate_dofs_geometrical(Q, outlet), Q)

bc_fully_conductive = dirichletbc(
    PETSc.ScalarType(0), locate_dofs_geometrical(Q2, hartmann_walls), Q2
)

bcu = [bc_noslip, bc_inflow]
bcp = [bc_outflow]
bcphi = [bc_fully_conductive]

u = TrialFunction(V)
v = TestFunction(V)
u_ = Function(V)
u_.name = "u"
u_s = Function(V)
u_s2 = Function(V)
u_n = Function(V)
u_n1 = Function(V)
p = TrialFunction(Q)
q = TestFunction(Q)
p_ = Function(Q)
p_.name = "p"
p_n = Function(Q)
phi = TrialFunction(Q2)
q2 = TestFunction(Q2)
phi_ = Function(Q2)
phi_.name = "phi"
phi_n = Function(Q2)

B = Constant(mesh, PETSc.ScalarType((0, -1, 0)))
Ha = Constant(mesh, (PETSc.ScalarType(10)))
N = Ha**2

f = N * (
    inner(cross(B, grad(phi_)), v) * dx
    + inner(u_s * dot(B, B), v) * dx
    - inner(B * dot(B, u_s), v) * dx
)

F1 = rho / k * dot(u - u_n, v) * dx
F1 += inner(dot(1.5 * u_n - 0.5 * u_n1, 0.5 * nabla_grad(u + u_n)), v) * dx
F1 += 0.5 * mu * inner(grad(u + u_n), grad(v)) * dx - dot(p_, div(v)) * dx
a1 = form(lhs(F1))
L1 = form(rhs(F1))
A1 = create_matrix(a1)
b1 = create_vector(L1)

a2 = form(dot(grad(p), grad(q)) * dx)
L2 = form(-rho / k * dot(div(u_s), q) * dx)
A2 = assemble_matrix(a2, bcs=bcp)
A2.assemble()
b2 = create_vector(L2)

F2 = (
    inner(grad(phi), grad(q2)) * dx
    - inner(dot(B, curl(u_s)) + dot(u_s, curl(B)), q2) * dx
)
a3 = form(lhs(F2))
L3 = form(rhs(F2))
A3 = create_matrix(a3)
b3 = create_vector(L3)

a4 = form(rho * dot(u, v) * dx)
L4 = form(rho * dot(u_s, v) * dx - k * dot(nabla_grad(p_n), v) * dx + f)
A4 = assemble_matrix(a4)
A4.assemble()
b4 = create_vector(L4)

# Solver for step 1
solver1 = PETSc.KSP().create(mesh.comm)
solver1.setOperators(A1)
solver1.setType(PETSc.KSP.Type.BCGS)
pc1 = solver1.getPC()
pc1.setType(PETSc.PC.Type.JACOBI)

# Solver for step 2
solver2 = PETSc.KSP().create(mesh.comm)
solver2.setOperators(A2)
solver2.setType(PETSc.KSP.Type.MINRES)
pc2 = solver2.getPC()
pc2.setType(PETSc.PC.Type.HYPRE)
pc2.setHYPREType("boomeramg")

# Solver for step 3
solver3 = PETSc.KSP().create(mesh.comm)
solver3.setOperators(A2)
solver3.setType(PETSc.KSP.Type.MINRES)
pc3 = solver3.getPC()
pc3.setType(PETSc.PC.Type.HYPRE)
pc3.setHYPREType("boomeramg")

# Solver for step 4
solver4 = PETSc.KSP().create(mesh.comm)
solver4.setOperators(A4)
solver4.setType(PETSc.KSP.Type.CG)
pc4 = solver4.getPC()
pc4.setType(PETSc.PC.Type.SOR)

results_folder = "Results/MHD_testing/"
u_xdmf = XDMFFile(mesh.comm, results_folder + "u.xdmf", "w")
u_xdmf.write_mesh(mesh)
u_xdmf.write_function(u_, t)
p_xdmf = XDMFFile(mesh.comm, results_folder + "p.xdmf", "w")
p_xdmf.write_mesh(mesh)
p_xdmf.write_function(p_, t)
phi_xdmf = XDMFFile(mesh.comm, results_folder + "phi.xdmf", "w")
phi_xdmf.write_mesh(mesh)
phi_xdmf.write_function(phi_, t)

progress = tqdm.autonotebook.tqdm(desc="Solving Navier-Stokes", total=num_steps)
for i in range(num_steps):
    progress.update(1)
    # Update current time step
    t += dt

    # Step 1: Tentative velocity step
    A1.zeroEntries()
    assemble_matrix(A1, a1, bcs=bcu)
    A1.assemble()
    with b1.localForm() as loc:
        loc.set(0)
    assemble_vector(b1, L1)
    apply_lifting(b1, [a1], [bcu])
    b1.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
    set_bc(b1, bcu)
    solver1.solve(b1, u_s.vector)
    u_s.x.scatter_forward()

    # Step 2: Pressure corrrection step
    with b2.localForm() as loc:
        loc.set(0)
    assemble_vector(b2, L2)
    apply_lifting(b2, [a2], [bcp])
    b2.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
    set_bc(b2, bcp)
    solver2.solve(b2, p_n.vector)
    p_n.x.scatter_forward()

    p_.vector.axpy(1, p_n.vector)
    p_.x.scatter_forward()

    # Step 4: Phi correction step
    with b3.localForm() as loc:
        loc.set(0)
    assemble_vector(b3, L3)
    apply_lifting(b3, [a3], [bcphi])
    b3.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
    set_bc(b3, bcphi)
    solver3.solve(b3, phi_n.vector)
    phi_n.x.scatter_forward()

    phi_.vector.axpy(1, phi_n.vector)
    phi_.x.scatter_forward()

    # Step 3: Velocity correction step
    with b4.localForm() as loc:
        loc.set(0)
    assemble_vector(b4, L4)
    b4.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
    solver4.solve(b4, u_.vector)
    u_.x.scatter_forward()

    # Write solutions to file
    u_xdmf.write_function(u_, t)
    p_xdmf.write_function(p_, t)
    phi_xdmf.write_function(phi_, t)

    # Update variable with solution from this time step
    with u_.vector.localForm() as loc_, u_n.vector.localForm() as loc_n, u_n1.vector.localForm() as loc_n1:
        loc_n.copy(loc_n1)
        loc_.copy(loc_n)

u_xdmf.close()
p_xdmf.close()
phi_xdmf.close()

```