That was a great suggestion. I got this to work, but ran into some other issues on how to compute the right-hand side of the equation. I am attaching my script below so other users facing a similar problem can see what worked.
Additionally, if you have suggestions on how I do this postprocessing, I would highly appreciate it.
The script is below.
```
import numpy as np
import dedalus.public as d3
import logging
logger = logging.getLogger(__name__)
# -----------------------
# Simulation parameters
# -----------------------
Lx = Ly = Lz = 2*np.pi # Domain size
Nx = Ny = Nz = 64 # Resolution
Re_L = 100 # Reynolds number
C = 50 # Damping coefficient for forcing
Ek_inf = 1.5 # Target kinetic energy
dealias = 3/2 # Dealiasing factor
stop_sim_time = 800 # End time
timestepper = d3.RK222 # Time-stepping scheme
max_timestep = 0.01
dtype = np.float64
# -----------------------
# Dedalus bases and fields
# -----------------------
coords = d3.CartesianCoordinates('x', 'y', 'z')
dist = d3.Distributor(coords, dtype=dtype)
# Spectral bases
xbasis = d3.RealFourier(coords['x'], size=Nx, bounds=(0,Lx), dealias=dealias)
ybasis = d3.RealFourier(coords['y'], size=Ny, bounds=(0,Ly), dealias=dealias)
zbasis = d3.Chebyshev(coords['z'], size=Nz, bounds=(0,Lz), dealias=dealias)
# Fields
p = dist.Field(name='p', bases=(xbasis, ybasis, zbasis))
u = dist.VectorField(coords, name='u', bases=(xbasis, ybasis, zbasis))
u_prime = dist.VectorField(coords, name='u_prime', bases=(xbasis, ybasis, zbasis))
u_hp_filtered = dist.VectorField(coords, name='u_filt', bases=(xbasis, ybasis, zbasis))
u1 = dist.Field(name='u1', bases=(xbasis, ybasis, zbasis))
u2 = dist.Field(name='u2', bases=(xbasis, ybasis, zbasis))
u3 = dist.Field(name='u3', bases=(xbasis, ybasis, zbasis))
tau_p = dist.Field(name='tau_p')
tau_u1 = dist.VectorField(coords, name='tau_u1', bases=(xbasis, ybasis))
tau_u2 = dist.VectorField(coords, name='tau_u2', bases=(xbasis, ybasis))
# Reynolds stress fields
u1u1 = dist.Field(name='u1u1', bases=(xbasis,ybasis,zbasis))
u1u2 = dist.Field(name='u1u2', bases=(xbasis,ybasis,zbasis))
u1u3 = dist.Field(name='u1u3', bases=(xbasis,ybasis,zbasis))
u2u1 = dist.Field(name='u2u1', bases=(xbasis,ybasis,zbasis))
u2u2 = dist.Field(name='u2u2', bases=(xbasis,ybasis,zbasis))
u2u3 = dist.Field(name='u2u3', bases=(xbasis,ybasis,zbasis))
u3u1 = dist.Field(name='u3u1', bases=(xbasis,ybasis,zbasis))
u3u2 = dist.Field(name='u3u2', bases=(xbasis,ybasis,zbasis))
u3u3 = dist.Field(name='u3u3', bases=(xbasis,ybasis,zbasis))
# Reynolds stress fields at start of time averaging
u1u1_statStartTime = dist.Field(name='u1u1_statStartTime', bases=(xbasis,ybasis,zbasis))
u1u2_statStartTime = dist.Field(name='u1u2_statStartTime', bases=(xbasis,ybasis,zbasis))
u1u3_statStartTime = dist.Field(name='u1u3_statStartTime', bases=(xbasis,ybasis,zbasis))
u2u1_statStartTime = dist.Field(name='u2u1_statStartTime', bases=(xbasis,ybasis,zbasis))
u2u2_statStartTime = dist.Field(name='u2u2_statStartTime', bases=(xbasis,ybasis,zbasis))
u2u3_statStartTime = dist.Field(name='u2u3_statStartTime', bases=(xbasis,ybasis,zbasis))
u3u1_statStartTime = dist.Field(name='u3u1_statStartTime', bases=(xbasis,ybasis,zbasis))
u3u2_statStartTime = dist.Field(name='u3u2_statStartTime', bases=(xbasis,ybasis,zbasis))
u3u3_statStartTime = dist.Field(name='u3u3_statStartTime', bases=(xbasis,ybasis,zbasis))
# Time-averaged planar fields
u1u1_tavg = dist.Field(name='u1u1_tavg', bases=(xbasis,ybasis,zbasis))
u1u2_tavg = dist.Field(name='u1u2_tavg', bases=(xbasis,ybasis,zbasis))
u1u3_tavg = dist.Field(name='u1u3_tavg', bases=(xbasis,ybasis,zbasis))
u2u1_tavg = dist.Field(name='u2u1_tavg', bases=(xbasis,ybasis,zbasis))
u2u2_tavg = dist.Field(name='u2u2_tavg', bases=(xbasis,ybasis,zbasis))
u2u3_tavg = dist.Field(name='u2u3_tavg', bases=(xbasis,ybasis,zbasis))
u3u1_tavg = dist.Field(name='u3u1_tavg', bases=(xbasis,ybasis,zbasis))
u3u2_tavg = dist.Field(name='u3u2_tavg', bases=(xbasis,ybasis,zbasis))
u3u3_tavg = dist.Field(name='u3u3_tavg', bases=(xbasis,ybasis,zbasis))
# -----------------------
# Substitutions & differential operators
# -----------------------
x = dist.local_grid(xbasis)
y = dist.local_grid(ybasis)
z = dist.local_grid(zbasis)
ex, ey, ez = coords.unit_vector_fields(dist)
dx = lambda A: d3.Differentiate(A, coords['x'])
dy = lambda A: d3.Differentiate(A, coords['y'])
dz = lambda A: d3.Differentiate(A, coords['z'])
lift_basis = zbasis.derivative_basis(1)
lift = lambda A: d3.Lift(A, lift_basis, -1)
grad_u = d3.grad(u) + ez*lift(tau_u1)
nu = 1.0/Re_L
A_force = 1.0
Volume = Lx*Ly*Lz
# # -----------------------
# # Initial velocity field (analytical)
# # -----------------------
# U = np.sqrt(24) * (Lx/np.sqrt(Lx**2 + Ly**2))
# u['g'][0] += U*np.sin((2*np.pi*x)/Lx)*np.cos((2*np.pi*y)/Ly)*np.cos((np.pi*z)/Lz)
# u['g'][1] += -U*(Ly/Lx)*np.cos((2*np.pi*x)/Lx)*np.sin((2*np.pi*y)/Ly)*np.cos((np.pi*z)/Lz)
# u['g'][2] += 0
# # Remove x-y mean to ensure zero planar average
# u['g'][0] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][0]
# u['g'][1] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][1]
# u['g'][2] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][2]
# -----------------------
# Balanced divergence-free initial condition
# -----------------------
# Define mode amplitude
A0 = 1.0 # arbitrary RMS amplitude (we’ll scale to exact RMS later)
# Base sinusoidal modes (triple sine-cosine to ensure div-free)
ux_init = np.sin(2*np.pi*x/Lx) * np.cos(2*np.pi*y/Ly) * np.cos(np.pi*z/Lz)
uy_init = -np.cos(2*np.pi*x/Lx) * np.sin(2*np.pi*y/Ly) * np.cos(np.pi*z/Lz)
uz_init = np.sin(2*np.pi*x/Lx) * np.sin(2*np.pi*y/Ly) * np.sin(np.pi*z/Lz)
# Stack into u field
u['g'][0] = ux_init
u['g'][1] = uy_init
u['g'][2] = uz_init
# Remove x-y mean to ensure zero planar average
u['g'][0] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][0]
u['g'][1] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][1]
u['g'][2] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][2]
# Compute current RMS of each component
rms_x = np.sqrt(np.mean(u['g'][0]**2))
rms_y = np.sqrt(np.mean(u['g'][1]**2))
rms_z = np.sqrt(np.mean(u['g'][2]**2))
# Compute target RMS (mean of the three components)
target_rms = np.sqrt((rms_x**2 + rms_y**2 + rms_z**2)/3)
# Rescale each component to the same RMS
u['g'][0] *= target_rms / rms_x
u['g'][1] *= target_rms / rms_y
u['g'][2] *= target_rms / rms_z
# compute initial bulk kinetic energy and scale to match Ek_inf
Ek_init = 0.5 * np.mean(u['g'][0]**2 + u['g'][1]**2 + u['g'][2]**2)
if Ek_init > 0:
scale_init = np.sqrt(Ek_inf / Ek_init)
u['g'][0] *= scale_init
u['g'][1] *= scale_init
u['g'][2] *= scale_init
u_hp_filtered.change_scales(dealias)
u_hp_filtered['g'][0] = u['g'][0]
u_hp_filtered['g'][1] = u['g'][1]
u_hp_filtered['g'][2] = u['g'][2]
# -----------------------
# Forcing and energy diagnostics
# -----------------------
Ek = d3.integ(0.5*((u@ex)**2 + (u@ey)**2 + (u@ez)**2))/Volume
# Velocity derivatives for dissipation
dudx = dx(u@ex); dudy = dy(u@ex); dudz = dz(u@ex)
dvdx = dx(u@ey); dvdy = dy(u@ey); dvdz = dz(u@ey)
dwdx = dx(u@ez); dwdy = dy(u@ez); dwdz = dz(u@ez)
eps = 1/Re_L * d3.integ(
2*(dudx**2 + dvdy**2 + dwdz**2) +
(dvdx**2 + dwdx**2 + dudy**2 + dwdy**2 + dudz**2 + dvdz**2) +
2*(dudy*dvdx + dudz*dwdx + dvdz*dwdy)
)/Volume
#F = A_force*((eps - C*(Ek - Ek_inf)) / (2*Ek)) # Forcing amplitude
# Denominator for forcing amplitude: <u·u_hp>
Pden = d3.integ(u@u_hp_filtered)/Volume
epsilon = 1e-12
F = (eps - C*(Ek - Ek_inf)) / (Pden)
statStartTime = 20.0 # Start time for statistics
# Create a time-dependent scalar field for the step function
t_mask = dist.Field(name='t_mask')
t_mask['g'][:] = 0 # initially zero
# Accumulator for averaging time
Tstat = dist.Field(name='Tstat')
# -----------------------
# Problem definition
# -----------------------
problem = d3.IVP([p, u, tau_p, tau_u1, tau_u2,
u1u1_tavg, u1u2_tavg, u1u3_tavg,
u2u1_tavg, u2u2_tavg, u2u3_tavg,
u3u1_tavg, u3u2_tavg, u3u3_tavg],
namespace=locals())
problem.add_equation("trace(grad_u) + tau_p = 0") # Incompressibility
problem.add_equation("dt(u) - nu*div(grad_u) + grad(p) + lift(tau_u2) = - u@grad(u) + F*u_hp_filtered")
problem.add_equation("u(z=0) = 0")
problem.add_equation("u(z=Lz) = 0")
problem.add_equation("integ(p) = 0") # Pressure gauge
# Time-averaged Reynolds stress evolution
problem.add_equation("dt(u1u1_tavg) = t_mask * u1u1")
problem.add_equation("dt(u1u2_tavg) = t_mask * u1u2")
problem.add_equation("dt(u1u3_tavg) = t_mask * u1u3")
problem.add_equation("dt(u2u1_tavg) = t_mask * u2u1")
problem.add_equation("dt(u2u2_tavg) = t_mask * u2u2")
problem.add_equation("dt(u2u3_tavg) = t_mask * u2u3")
problem.add_equation("dt(u3u1_tavg) = t_mask * u3u1")
problem.add_equation("dt(u3u2_tavg) = t_mask * u3u2")
problem.add_equation("dt(u3u3_tavg) = t_mask * u3u3")
solver = problem.build_solver(timestepper)
solver.stop_sim_time = stop_sim_time
# -----------------------
# Precompute Reynolds stresses at initialization
# -----------------------
u1u1.change_scales(dealias)
u1u2.change_scales(dealias)
u1u3.change_scales(dealias)
u2u1.change_scales(dealias)
u2u2.change_scales(dealias)
u2u3.change_scales(dealias)
u3u1.change_scales(dealias)
u3u2.change_scales(dealias)
u3u3.change_scales(dealias)
u1u1['g'][:] = (u@ex*u@ex).evaluate()['g']
u1u2['g'][:] = (u@ex*u@ey).evaluate()['g']
u1u3['g'][:] = (u@ex*u@ez).evaluate()['g']
u2u1['g'][:] = (u@ey*u@ex).evaluate()['g']
u2u2['g'][:] = (u@ey*u@ey).evaluate()['g']
u2u3['g'][:] = (u@ey*u@ez).evaluate()['g']
u3u1['g'][:] = (u@ez*u@ex).evaluate()['g']
u3u2['g'][:] = (u@ez*u@ey).evaluate()['g']
u3u3['g'][:] = (u@ez*u@ez).evaluate()['g']
# -----------------------
# Analysis handlers
# -----------------------
# snapshots = solver.evaluator.add_file_handler('snapshots', sim_dt=0.25, max_writes=50)
# snapshots.add_task(p, name='p')
# snapshots.add_task(u, name='u')
# snapshots.add_task(F, name='F')
# snapshots.add_task(Ek, name='Ek')
# snapshots.add_task(eps, name='eps')
# analysis = solver.evaluator.add_file_handler('analysis', iter=1, max_writes=4000000)
# analysis.add_task(d3.Average(u1u1_tavg, ('x','y')), layout='g', name='<u1u1_ave>')
# analysis.add_task(d3.Average(u1u2_tavg, ('x','y')), layout='g', name='<u1u2_ave>')
# analysis.add_task(d3.Average(u1u3_tavg, ('x','y')), layout='g', name='<u1u3_ave>')
# analysis.add_task(d3.Average(u2u1_tavg, ('x','y')), layout='g', name='<u2u1_ave>')
# analysis.add_task(d3.Average(u2u2_tavg, ('x','y')), layout='g', name='<u2u2_ave>')
# analysis.add_task(d3.Average(u2u3_tavg, ('x','y')), layout='g', name='<u2u3_ave>')
# analysis.add_task(d3.Average(u3u1_tavg, ('x','y')), layout='g', name='<u3u1_ave>')
# analysis.add_task(d3.Average(u3u2_tavg, ('x','y')), layout='g', name='<u3u2_ave>')
# analysis.add_task(d3.Average(u3u3_tavg, ('x','y')), layout='g', name='<u3u3_ave>')
# analysis.add_tasks(solver.state, layout='g')
snapshots = solver.evaluator.add_file_handler('snapshot', sim_dt=10, max_writes=1)
snapshots.add_task(p, name='p')
snapshots.add_task(u, name='u')
snapshots.add_task(u_hp_filtered, name='u_hp_filtered')
#snapshots.add_task(F, name='F')
#snapshots.add_task(Ek, name='Ek')
#snapshots.add_task(eps, name='eps')
analysis = solver.evaluator.add_file_handler('statistics', iter=1, max_writes=4000000)
analysis.add_task(d3.Average(u1u1_tavg, ('x','y')) / (Tstat), layout='g', name='<u1u1_ave>')
analysis.add_task(d3.Average(u1u2_tavg, ('x','y')) / (Tstat), layout='g', name='<u1u2_ave>')
analysis.add_task(d3.Average(u1u3_tavg, ('x','y')) / (Tstat), layout='g', name='<u1u3_ave>')
analysis.add_task(d3.Average(u2u1_tavg, ('x','y')) / (Tstat), layout='g', name='<u2u1_ave>')
analysis.add_task(d3.Average(u2u2_tavg, ('x','y')) / (Tstat), layout='g', name='<u2u2_ave>')
analysis.add_task(d3.Average(u2u3_tavg, ('x','y')) / (Tstat), layout='g', name='<u2u3_ave>')
analysis.add_task(d3.Average(u3u1_tavg, ('x','y')) / (Tstat), layout='g', name='<u3u1_ave>')
analysis.add_task(d3.Average(u3u2_tavg, ('x','y')) / (Tstat), layout='g', name='<u3u2_ave>')
analysis.add_task(d3.Average(u3u3_tavg, ('x','y')) / (Tstat), layout='g', name='<u3u3_ave>')
analysis.add_tasks(solver.state, layout='g')
analysis.add_task(Tstat, layout='g', name='Tstat')
# -----------------------
# CFL and flow properties
# -----------------------
CFL = d3.CFL(solver, initial_dt=max_timestep, cadence=10, safety=0.4,
threshold=0.05, max_change=1.5, min_change=0.5, max_dt=max_timestep)
CFL.add_velocity(u)
timestep = CFL.compute_timestep()
flow = d3.GlobalFlowProperty(solver, cadence=10)
flow.add_property(np.sqrt(u@u)/nu, name='Re')
flow.add_property((u@ex)**2, name='u2')
flow.add_property((u@ey)**2, name='v2')
flow.add_property((u@ez)**2, name='w2')
flow.add_property(F, name='F')
flow.add_property(Ek, name='bulk_Ek')
# -----------------------
# Precompute velocity components for RMS/Reynolds stress calculations
# -----------------------
u_ex = dist.Field(name='u_ex', bases=(xbasis,ybasis,zbasis))
u_ey = dist.Field(name='u_ey', bases=(xbasis,ybasis,zbasis))
u_ez = dist.Field(name='u_ez', bases=(xbasis,ybasis,zbasis))
u_ex.change_scales(dealias)
u_ey.change_scales(dealias)
u_ez.change_scales(dealias)
# -----------------------
# Main loop
# -----------------------
try:
logger.info("Starting main loop")
while solver.proceed:
timestep = CFL.compute_timestep()
# In the main loop, update it each timestep:
t_mask['g'][:] = 1.0 if solver.sim_time >= statStartTime else 0.0
Tstat += timestep if solver.sim_time >= statStartTime else 0.0
# -----------------------
# Update mean-free and high-pass filtered velocity
# -----------------------
u_prime.change_scales(dealias)
u_hp_filtered.change_scales(dealias)
u_hp_filtered['g'][0] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][0]
u_hp_filtered['g'][1] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][1]
u_hp_filtered['g'][2] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][2]
#u_hp_filtered.high_pass_filter(0.75)
u_hp_filtered.high_pass_filter((8, 8, 1))
solver.step(timestep)
# -----------------------
# Enforce zero planar average
# -----------------------
u['g'][0] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][0]
u['g'][1] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][1]
u['g'][2] = (u - d3.Average(u, ('x','y'))).evaluate()['g'][2]
# -----------------------
# Precompute velocity components
# -----------------------
u_ex['g'][:] = u['g'][0]
u_ey['g'][:] = u['g'][1]
u_ez['g'][:] = u['g'][2]
# -----------------------
# Update Reynolds stress fields in-place
# -----------------------
u1u1['g'][:] = u_ex['g'] * u_ex['g']
u1u2['g'][:] = u_ex['g'] * u_ey['g']
u1u3['g'][:] = u_ex['g'] * u_ez['g']
u2u1['g'][:] = u_ey['g'] * u_ex['g']
u2u2['g'][:] = u_ey['g'] * u_ey['g']
u2u3['g'][:] = u_ey['g'] * u_ez['g']
u3u1['g'][:] = u_ez['g'] * u_ex['g']
u3u2['g'][:] = u_ez['g'] * u_ey['g']
u3u3['g'][:] = u_ez['g'] * u_ez['g']
# -----------------------
# Logging
# -----------------------
if (solver.iteration-1) % 10 == 0:
max_u = np.sqrt(flow.max('u2'))
max_v = np.sqrt(flow.max('v2'))
max_w = np.sqrt(flow.max('w2'))
max_F = flow.max('F')
bulk_Ek = flow.max('bulk_Ek')
max_Re = flow.max('Re')
logger.info(
f"Iteration={solver.iteration}, Time={solver.sim_time:.6e}, dt={timestep:.3e}, "
f"max(u)={max_u:.6f}, max(v)={max_v:.6f}, max(w)={max_w:.6f}, "
f"max(F)={max_F:.6f}, bulk Ek={bulk_Ek:.6f}, max(Re)={max_Re:.6f}"
)
except:
logger.error("Exception raised, stopping main loop")
raise
finally:
solver.log_stats()
```
Thank you for your help.