Update a function within the loop that depends on other functions

Hello FEniCSx community. I am working on a turbulent model based on Dokken tutorial. My question is, what is the correct way to implement a function that depends on two other functions that are solved using another model (k-epsilon model) at each time step?

import gmsh
import numpy as np
import tqdm.autonotebook

from mpi4py import MPI
from petsc4py import PETSc

from basix.ufl import element
from dolfinx.fem import (Constant, Function, functionspace,
                         assemble_scalar, dirichletbc, form, locate_dofs_topological, set_bc)
from dolfinx.fem.petsc import (apply_lifting, assemble_matrix, assemble_vector,
                               create_vector, create_matrix, set_bc)
from dolfinx.io import (VTXWriter, distribute_entity_data, gmshio)
from ufl import (FacetNormal, Identity, Measure, TestFunction, TrialFunction,
                 as_vector, div, dot, ds, dx, inner, lhs, grad, nabla_grad, rhs, sym, system)

gmsh.initialize()

L = 2.2
H = 0.41
c_x = c_y = 0.2
r = 0.05
gdim = 2
mesh_comm = MPI.COMM_WORLD
model_rank = 0
if mesh_comm.rank == model_rank:
    rectangle = gmsh.model.occ.addRectangle(0, 0, 0, L, H, tag=1)
    obstacle = gmsh.model.occ.addDisk(c_x, c_y, 0, r, r)

if mesh_comm.rank == model_rank:
    fluid = gmsh.model.occ.cut([(gdim, rectangle)], [(gdim, obstacle)])
    gmsh.model.occ.synchronize()

fluid_marker = 1
if mesh_comm.rank == model_rank:
    volumes = gmsh.model.getEntities(dim=gdim)
    assert (len(volumes) == 1)
    gmsh.model.addPhysicalGroup(volumes[0][0], [volumes[0][1]], fluid_marker)
    gmsh.model.setPhysicalName(volumes[0][0], fluid_marker, "Fluid")

inlet_marker, outlet_marker, wall_marker, obstacle_marker = 2, 3, 4, 5
inflow, outflow, walls, obstacle = [], [], [], []
if mesh_comm.rank == model_rank:
    boundaries = gmsh.model.getBoundary(volumes, oriented=False)
    for boundary in boundaries:
        center_of_mass = gmsh.model.occ.getCenterOfMass(boundary[0], boundary[1])
        if np.allclose(center_of_mass, [0, H / 2, 0]):
            inflow.append(boundary[1])
        elif np.allclose(center_of_mass, [L, H / 2, 0]):
            outflow.append(boundary[1])
        elif np.allclose(center_of_mass, [L / 2, H, 0]) or np.allclose(center_of_mass, [L / 2, 0, 0]):
            walls.append(boundary[1])
        else:
            obstacle.append(boundary[1])
    gmsh.model.addPhysicalGroup(1, walls, wall_marker)
    gmsh.model.setPhysicalName(1, wall_marker, "Walls")
    gmsh.model.addPhysicalGroup(1, inflow, inlet_marker)
    gmsh.model.setPhysicalName(1, inlet_marker, "Inlet")
    gmsh.model.addPhysicalGroup(1, outflow, outlet_marker)
    gmsh.model.setPhysicalName(1, outlet_marker, "Outlet")
    gmsh.model.addPhysicalGroup(1, obstacle, obstacle_marker)
    gmsh.model.setPhysicalName(1, obstacle_marker, "Obstacle")

res_min = r / 3
if mesh_comm.rank == model_rank:
    distance_field = gmsh.model.mesh.field.add("Distance")
    gmsh.model.mesh.field.setNumbers(distance_field, "EdgesList", obstacle)
    threshold_field = gmsh.model.mesh.field.add("Threshold")
    gmsh.model.mesh.field.setNumber(threshold_field, "IField", distance_field)
    gmsh.model.mesh.field.setNumber(threshold_field, "LcMin", res_min)
    gmsh.model.mesh.field.setNumber(threshold_field, "LcMax", 0.25 * H)
    gmsh.model.mesh.field.setNumber(threshold_field, "DistMin", r)
    gmsh.model.mesh.field.setNumber(threshold_field, "DistMax", 2 * H)
    min_field = gmsh.model.mesh.field.add("Min")
    gmsh.model.mesh.field.setNumbers(min_field, "FieldsList", [threshold_field])
    gmsh.model.mesh.field.setAsBackgroundMesh(min_field)

if mesh_comm.rank == model_rank:
    gmsh.option.setNumber("Mesh.Algorithm", 8)
    gmsh.option.setNumber("Mesh.RecombinationAlgorithm", 2)
    gmsh.option.setNumber("Mesh.RecombineAll", 1)
    gmsh.option.setNumber("Mesh.SubdivisionAlgorithm", 1)
    gmsh.model.mesh.generate(gdim)
    gmsh.model.mesh.setOrder(2)
    gmsh.model.mesh.optimize("Netgen")

mesh, _, ft = gmshio.model_to_mesh(gmsh.model, mesh_comm, model_rank, gdim=gdim)
ft.name = "Facet markers"

t = 0
T = 1                       # Final time
dt = 1 / 1600                 # Time step size
num_steps = int(T / dt)
k = Constant(mesh, PETSc.ScalarType(dt))
mu = Constant(mesh, PETSc.ScalarType(0.001))  # Dynamic viscosity
rho = Constant(mesh, PETSc.ScalarType(1))     # Density

v_cg2 = element("Lagrange", mesh.topology.cell_name(), 2, shape=(mesh.geometry.dim, ))
s_cg1 = element("Lagrange", mesh.topology.cell_name(), 1)
V = functionspace(mesh, v_cg2)
Q = functionspace(mesh, s_cg1)
W = functionspace(mesh, s_cg1)

fdim = mesh.topology.dim - 1

# Define boundary conditions


class InletVelocity():
    def __init__(self, t):
        self.t = t

    def __call__(self, x):
        values = np.zeros((gdim, x.shape[1]), dtype=PETSc.ScalarType)
        values[0] = 4 * 1.5 * np.sin(self.t * np.pi / 8) * x[1] * (0.41 - x[1]) / (0.41**2)
        return values


# Inlet
u_inlet = Function(V)
inlet_velocity = InletVelocity(t)
u_inlet.interpolate(inlet_velocity)
bcu_inflow = dirichletbc(u_inlet, locate_dofs_topological(V, fdim, ft.find(inlet_marker)))
# Walls
u_nonslip = np.array((0,) * mesh.geometry.dim, dtype=PETSc.ScalarType)
bcu_walls = dirichletbc(u_nonslip, locate_dofs_topological(V, fdim, ft.find(wall_marker)), V)
# Obstacle
bcu_obstacle = dirichletbc(u_nonslip, locate_dofs_topological(V, fdim, ft.find(obstacle_marker)), V)
bcu = [bcu_inflow, bcu_obstacle, bcu_walls]
# Outlet
bcp_outlet = dirichletbc(PETSc.ScalarType(0), locate_dofs_topological(Q, fdim, ft.find(outlet_marker)), Q)
bcp = [bcp_outlet]

u = TrialFunction(V)
v = TestFunction(V)
u_ = Function(V)
u_.name = "u"
u_s = Function(V)
u_n = Function(V)
u_n1 = Function(V)
p = TrialFunction(Q)
q = TestFunction(Q)
p_ = Function(Q)
p_.name = "p"
phi = Function(Q)
k_n = Function(W); k_n.name = "k"
e_n = Function(W); e_n.name = "e"

u_ref = 2.0
C_mu = 0.09
l_0 = 0.1
k_0 = 0.01 * u_ref**2
e_0 = C_mu**0.75 * (k_0**1.5) / l_0
k_n.interpolate(lambda x: np.full(x.shape[1], k_0))
e_n.interpolate(lambda x: np.full(x.shape[1], e_0))
k_n.x.scatter_forward()
e_n.x.scatter_forward()

mu_t = rho * C_mu * k_n**2 / e_n

f = Constant(mesh, PETSc.ScalarType((0, 0)))
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 * inner((mu + mu_t) * grad(u + u_n), grad(v)) * dx - dot(p_, div(v)) * dx
F1 += dot(f, 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)

a3 = form(rho * dot(u, v) * dx)
L3 = form(rho * dot(u_s, v) * dx - k * dot(nabla_grad(phi), v) * dx)
A3 = assemble_matrix(a3)
A3.assemble()
b3 = create_vector(L3)

# 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(A3)
solver3.setType(PETSc.KSP.Type.CG)
pc3 = solver3.getPC()
pc3.setType(PETSc.PC.Type.SOR)

from pathlib import Path
folder = Path("/home/gerard/Desktop/results")
folder.mkdir(exist_ok=True, parents=True)
vtx_u = VTXWriter(mesh.comm, folder / "dfg2D-3-u.bp", [u_], engine="BP4")
vtx_p = VTXWriter(mesh.comm, folder / "dfg2D-3-p.bp", [p_], engine="BP4")
vtx_u.write(t)
vtx_p.write(t)
progress = tqdm.autonotebook.tqdm(desc="Solving PDE", total=num_steps)
for i in range(num_steps):
    progress.update(1)
    # Update current time step
    t += dt
    # Update inlet velocity
    inlet_velocity.t = t
    u_inlet.interpolate(inlet_velocity)

    # 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.x.petsc_vec)
    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, phi.x.petsc_vec)
    phi.x.scatter_forward()

    p_.x.petsc_vec.axpy(1, phi.x.petsc_vec)
    p_.x.scatter_forward()

    # Step 3: Velocity correction step
    with b3.localForm() as loc:
        loc.set(0)
    assemble_vector(b3, L3)
    b3.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
    solver3.solve(b3, u_.x.petsc_vec)
    u_.x.scatter_forward()

    # Write solutions to file
    vtx_u.write(t)
    vtx_p.write(t)

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

progress.close()
vtx_u.close()
vtx_p.close()

I have also done it using a Python function, but since it returns a vector, I am not sure if there are any problems reading it in the variational form.

def update_eddy_viscosity(k_func, e_func):
    k_vals = k_func.x.array[:]
    e_vals = e_func.x.array[:]
    l_star = 1.0
    mu_t_vals = rho * C_mu * (k_vals**2 / e_vals)
    mu_t_vals = np.minimum(mu_t_vals, rho * l_star * np.sqrt(k_vals))
    mu_t_vals = np.maximum(mu_t_vals, mu)
    mu_t.x.array[:] = mu_t_vals
    mu_t.x.scatter_forward()

Note that the mu_t function must be limited, which is why the idea is to calculate it as an array. I am not attaching the implementation of the k-eps model so as not to make the post too long.

Instead of working on the numpy-array itself it is safer to define a UFL-expression that describes the operation you would like to do.

In the snippets you have provided above, it a bit unclear how your mu_t relates to the mu_t in the variational form as:

indicates that you have a DOLFINx function, while:

is a UFL expression above.

Note thatmu_t shouldn’t necessarily be in the same function space as any of the other variables, as it contains an algebraic expression depending on k-\epsilon.

I will illustrate how to use ufl for some of these operations below.

from mpi4py import MPI
import dolfinx
import numpy as np
import basix.ufl
import ufl

mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 4, 4)
el = basix.ufl.element("Lagrange", mesh.basix_cell(), 1)
V = dolfinx.fem.functionspace(mesh, el)

k = dolfinx.fem.Function(V)
k.interpolate(lambda x: 2 + np.sin(np.pi * x[0]) * np.sin(np.pi * x[1]))

eps = dolfinx.fem.Function(V)
eps.interpolate(lambda x: 0.01 + x[0] ** 2)

C_mu = dolfinx.fem.Constant(mesh, 0.09)
dt = 0.01
mu = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.8))  # Dynamic viscosity
rho = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1))  # Density
l_star = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0))

mu_t = rho * C_mu * k**2 / eps

bounded_lower = ufl.min_value(mu_t, rho * l_star * ufl.sqrt(k))
bounded_upper = ufl.max_value(bounded_lower, mu)

out_el = basix.ufl.element("Lagrange", mesh.basix_cell(), 3)
Q = dolfinx.fem.functionspace(mesh, out_el)
uh = dolfinx.fem.Function(Q)
try:
    uh.interpolate(
        dolfinx.fem.Expression(bounded_upper, Q.element.interpolation_points())
    )
except TypeError:
    uh.interpolate(
        dolfinx.fem.Expression(bounded_upper, Q.element.interpolation_points)
    )

with dolfinx.io.VTXWriter(MPI.COMM_WORLD, "mwe_output.bp", [uh]) as writer:
    writer.write(0.0)

With this code, you can choose to:

  1. Interpolate bounded_upper into a suitable function space.
  2. Use the expression bounded_upper directly in the variational form

Thanks for your reply, Dokken. According to the two options, which one is correct for the mut variable to be updated at each time step when k and epsilon are updated? Because I tried both ways, but it is not updated at each time step. Or should I do the calculation inside the loop?

Another question: is it necessary to define all parameters such as mu, rho, and dt as constants? Is there any difference in the solution?

Could you share your attempts? As either solution should reflect changes in mu_t if written correctly.

No, they can be whatever.

Yes, there will be a big difference in the two solutions.

Option 1:

Here you introduce an interpolation of a potentially higher order function into a lower order space.
This will affect the representation of bound_upper and quality of solution. One need to choose this space with care.

Option 2:

This is the most accurate one, as it will evaluate this expression at the relevant quadrature points.
Note that some formulations might get rather expensive in this way, as the quadrature estimate can go off the charts. See for instance: The UFL forms — FEniCS Workshop on how to estimate and limit quadrature rules.

To make Option 1 and 2 equivalent: Choose a quadrature space basix.ufl — Basix 0.10.0 documentation of appropriate order for option 1. Interpolate into that and use it in the form.

Hi @dokken, I have left the constant velocity to only show the code for the k-epsilon model. I mention that it is not updated because there are no changes in the solution with or without mut.

from petsc4py import PETSc
import numpy as np
import tqdm.autonotebook
import gmsh
from dolfinx.io import VTXWriter, gmshio
from basix.ufl import element
from dolfinx.fem import (functionspace, Function, locate_dofs_topological, dirichletbc, Constant, form, assemble_scalar)
from dolfinx.fem.petsc import assemble_matrix, assemble_vector, apply_lifting, create_matrix, create_vector, set_bc
from dolfinx import io, log, mesh, default_scalar_type
from ufl import TrialFunction, TestFunction, dot, grad, nabla_grad, inner, dx, div, lhs, rhs, sqrt, sym, Identity, \
    CellDiameter, max_value, min_value

# For MPI-based parallelization
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()

gmsh.initialize()

L = 2.2
H = 0.41
c_x = c_y = 0.2
r = 0.05
gdim = 2
mesh_comm = MPI.COMM_WORLD
model_rank = 0
if mesh_comm.rank == model_rank:
    rectangle = gmsh.model.occ.addRectangle(0, 0, 0, L, H, tag=1)
    obstacle = gmsh.model.occ.addDisk(c_x, c_y, 0, r, r)

if mesh_comm.rank == model_rank:
    fluid = gmsh.model.occ.cut([(gdim, rectangle)], [(gdim, obstacle)])
    gmsh.model.occ.synchronize()

fluid_marker = 1
if mesh_comm.rank == model_rank:
    volumes = gmsh.model.getEntities(dim=gdim)
    assert (len(volumes) == 1)
    gmsh.model.addPhysicalGroup(volumes[0][0], [volumes[0][1]], fluid_marker)
    gmsh.model.setPhysicalName(volumes[0][0], fluid_marker, "Fluid")

inlet_marker, outlet_marker, wall_marker, obstacle_marker = 2, 3, 4, 5
inflow, outflow, walls, obstacle = [], [], [], []
if mesh_comm.rank == model_rank:
    boundaries = gmsh.model.getBoundary(volumes, oriented=False)
    for boundary in boundaries:
        center_of_mass = gmsh.model.occ.getCenterOfMass(boundary[0], boundary[1])
        if np.allclose(center_of_mass, [0, H / 2, 0]):
            inflow.append(boundary[1])
        elif np.allclose(center_of_mass, [L, H / 2, 0]):
            outflow.append(boundary[1])
        elif np.allclose(center_of_mass, [L / 2, H, 0]) or np.allclose(center_of_mass, [L / 2, 0, 0]):
            walls.append(boundary[1])
        else:
            obstacle.append(boundary[1])
    gmsh.model.addPhysicalGroup(1, walls, wall_marker)
    gmsh.model.setPhysicalName(1, wall_marker, "Walls")
    gmsh.model.addPhysicalGroup(1, inflow, inlet_marker)
    gmsh.model.setPhysicalName(1, inlet_marker, "Inlet")
    gmsh.model.addPhysicalGroup(1, outflow, outlet_marker)
    gmsh.model.setPhysicalName(1, outlet_marker, "Outlet")
    gmsh.model.addPhysicalGroup(1, obstacle, obstacle_marker)
    gmsh.model.setPhysicalName(1, obstacle_marker, "Obstacle")

res_min = r / 3
if mesh_comm.rank == model_rank:
    distance_field = gmsh.model.mesh.field.add("Distance")
    gmsh.model.mesh.field.setNumbers(distance_field, "EdgesList", obstacle)
    threshold_field = gmsh.model.mesh.field.add("Threshold")
    gmsh.model.mesh.field.setNumber(threshold_field, "IField", distance_field)
    gmsh.model.mesh.field.setNumber(threshold_field, "LcMin", res_min)
    gmsh.model.mesh.field.setNumber(threshold_field, "LcMax", 0.25 * H)
    gmsh.model.mesh.field.setNumber(threshold_field, "DistMin", r)
    gmsh.model.mesh.field.setNumber(threshold_field, "DistMax", 2 * H)
    min_field = gmsh.model.mesh.field.add("Min")
    gmsh.model.mesh.field.setNumbers(min_field, "FieldsList", [threshold_field])
    gmsh.model.mesh.field.setAsBackgroundMesh(min_field)

if mesh_comm.rank == model_rank:
    gmsh.option.setNumber("Mesh.Algorithm", 8)
    gmsh.option.setNumber("Mesh.RecombinationAlgorithm", 2)
    gmsh.option.setNumber("Mesh.RecombineAll", 1)
    gmsh.option.setNumber("Mesh.SubdivisionAlgorithm", 1)
    gmsh.model.mesh.generate(gdim)
    gmsh.model.mesh.setOrder(2)
    gmsh.model.mesh.optimize("Netgen")

mesh, _, ft = gmshio.model_to_mesh(gmsh.model, mesh_comm, model_rank, gdim=gdim)
ft.name = "Facet markers"

# Define time step
t = 0
T = 1                       # Final time
dt = 1 / 1600                 # Time step size
num_steps = int(T / dt)
u_ref = 2.0    # reference velocity (m s-1)

# k-eps model data
C_mu = 0.09
sigma_k = 1.0
sigma_e = 1.3
sigma_t = 1.0
C_e1 = 1.44
C_e2 = 1.92
omega_k = 0.5
omega_eps = 0.1

# Constant data
Pr = 0.708    # Prandtl number
Kappa = 0.41  # Constant von Kármán

k = Constant(mesh, PETSc.ScalarType(dt))
mu = Constant(mesh, PETSc.ScalarType(0.001))  # Dynamic viscosity
rho = Constant(mesh, PETSc.ScalarType(1))     # Density

# Define functions space
v_cg2 = element("Lagrange", mesh.topology.cell_name(), 2, shape=(mesh.geometry.dim,))
s_cg1 = element("Lagrange", mesh.topology.cell_name(), 1)
V = functionspace(mesh, v_cg2)      # Space for speed
W_k = functionspace(mesh, s_cg1)    # Space for kinetic
W_e = functionspace(mesh, s_cg1)    # Space for epsilon

fdim = mesh.topology.dim - 1

class Velocity:
    def __call__(self, x):
        values = np.zeros((2, x.shape[1]), dtype=PETSc.ScalarType)
        values[0] = u_ref
        return values

class InletKinetic:
    def __call__(self, x):
        values = np.zeros((1, x.shape[1]), dtype=PETSc.ScalarType)
        values[0] = 0.01 * u_ref**2
        return values

class InletEpsilon:
    def __call__(self, x):
        y = np.maximum(x[1], 1.0)
        k_in = 0.01 * u_ref**2
        values = np.zeros((1, x.shape[1]), dtype=PETSc.ScalarType)
        values[0] = C_mu**0.75 * k_in**1.5 / (Kappa * y)
        return values

# Boundary condition
k_inlet = Function(W_k)
inlet_kinetic = InletKinetic()
k_inlet.interpolate(inlet_kinetic)
bck_inflow = dirichletbc(k_inlet, locate_dofs_topological(W_k, fdim, ft.find(inlet_marker)))

e_inlet = Function(W_e)
inlet_kinetic = InletEpsilon()
e_inlet.interpolate(inlet_kinetic)
bce_inflow = dirichletbc(e_inlet, locate_dofs_topological(W_e, fdim, ft.find(inlet_marker)))

bc_k = [bck_inflow]
bc_e = [bce_inflow]

k_trial = TrialFunction(W_k); w_k = TestFunction(W_k)       # Kinetic energy
e_trial = TrialFunction(W_e); w_e = TestFunction(W_e)       # Dissipation energy

u_ = Function(V)
velocity = Velocity()
u_.interpolate(velocity)

# functions with DOFs at the previous step
k_n = Function(W_k); k_n.name = "k"
e_n = Function(W_e); e_n.name = "e"
k_ = Function(W_k); e_ = Function(W_e)

# Define the Initial Value for k-eps and temperature
k_0 = 0.01 * u_ref**2
k_n.interpolate(lambda x: np.full(x.shape[1], k_0))
l_0 = 10
e_0 = C_mu**0.75 * (k_0**1.5) / l_0
e_n.interpolate(lambda x: np.full(x.shape[1], e_0))
k_n.x.scatter_forward()
e_n.x.scatter_forward()

mu_t = C_mu * (k_n ** 2 / e_n)

l_star = Constant(mesh, default_scalar_type(1.0))
bounded_lower = min_value(mu_t, rho * l_star * sqrt(k))
bounded_upper = max_value(bounded_lower, mu)
mut = bounded_upper

S_sq = 2 * inner(sym(nabla_grad(u_)), sym(nabla_grad(u_)))
prod_k = mu_t * S_sq
S = 0.5 * (nabla_grad(u_) + nabla_grad(u_).T)
react_k = e_n / k_n

Qu = sqrt(dot(u_, u_))

FK = rho * (1.0 / dt) * dot((k_trial - k_n), w_k) * dx
FK += rho * dot(u_, nabla_grad(k_trial)) * w_k * dx
FK += inner((mu + mut / sigma_k) * grad(k_trial), grad(w_k)) * dx
FK -= prod_k * w_k * dx
FK += rho * react_k * k_trial * w_k * dx
a_k = form(lhs(FK))
L_k = form(rhs(FK))
A_k = create_matrix(a_k)
b_k = create_vector(L_k)

FE = rho * (1.0 / dt) * dot((e_trial - e_n), w_e) * dx
FE += rho * dot(u_, grad(e_trial)) * w_e * dx
FE += inner((mu + mut / sigma_e) * grad(e_trial), grad(w_e)) * dx
FE -= C_e1 * react_k * prod_k * w_e * dx
FE += rho * C_e2 * react_k * e_trial * w_e * dx
a_e = form(lhs(FE))
L_e = form(rhs(FE))
A_e = create_matrix(a_e)
b_e = create_vector(L_e)

# Solver for step 2
solver_k = PETSc.KSP().create(mesh.comm)
solver_k.setOperators(A_k)
solver_k.setType("preonly")
# Pre-conditioner
pc_k = solver_k.getPC()
pc_k.setType("lu")
pc_k.setFactorSolverType("mumps")
pc_k.setFactorSetUpSolverType()

# Solver for step 3
solver_e = PETSc.KSP().create(mesh.comm)
solver_e.setOperators(A_e)
solver_e.setType("preonly")
# Pre-conditioner
pc_e = solver_e.getPC()
pc_e.setType("lu")
pc_e.setFactorSolverType("mumps")
pc_e.setFactorSetUpSolverType()

from pathlib import Path
folder = Path("/home/gerard/Desktop/k_epsilon/")
folder.mkdir(exist_ok=True, parents=True)
vtx_k = VTXWriter(mesh.comm, folder / "k.bp", k_n, engine="BP4")
vtx_e = VTXWriter(mesh.comm, folder / "e.bp", e_n, engine="BP4")
#vtx_mu = VTXWriter(mesh.comm, folder / "mut.bp", mut, engine="BP4")
vtx_k.write(t)
vtx_e.write(t)
progress = tqdm.autonotebook.tqdm(desc="Solving PDE", total=num_steps)

for i in range(num_steps):
    progress.update(1)
    t += dt

    A_k.zeroEntries()
    assemble_matrix(A_k, a_k, bcs=bc_k)
    A_k.assemble()
    with b_k.localForm() as loc:
        loc.set(0)
    assemble_vector(b_k, L_k)
    apply_lifting(b_k, [a_k], [bc_k])
    b_k.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
    set_bc(b_k, bc_k)
    solver_k.solve(b_k, k_.x.petsc_vec)
    k_.x.scatter_forward()

    A_e.zeroEntries()
    assemble_matrix(A_e, a_e, bcs=bc_e)
    A_e.assemble()
    with b_e.localForm() as loc:
        loc.set(0)
    assemble_vector(b_e, L_e)
    apply_lifting(b_e, [a_e], [bc_e])
    b_e.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
    set_bc(b_e, bc_e)
    solver_e.solve(b_e, e_.x.petsc_vec)

    k_n.x.array[:] = k_.x.array[:]
    e_n.x.array[:] = e_.x.array[:]

    # Write solutions to file
    vtx_k.write(t)
    vtx_e.write(t)
    #vtx_mu.write(t)

progress.close()
vtx_e.close()
vtx_k.close()
#vtx_mu.close()

I tried to save it in a file, but I get an error. I think that to do that, I have to interpolate in a suitable space. That’s option 1 that you mentioned. Is there any way to do that?

So I could leave it as a float variable and use it in the variational form, right?