dedalus_ab_trial_NOT_FINAL.py

import os
import numpy as np
import matplotlib.pyplot as plt

from dedalus import public as de
from dedalus.extras import plot_tools

data = np.load('enmean_9900.npy',allow_pickle=True)
data = data.item()

# loading data for rhobar and rho total from our loaded dictionary
rhobar = np.nanmean(data['rho'][1:-1,1:-1,1:-1],axis=(0,1))
rhot = data['rhot'][1:-1,1:-1,1:-1]
u = data['u'][1:-1,1:-1,1:-1]
v = data['v'][1:-1,1:-1,1:-1]
w = data['w'][1:-1,1:-1,1:-1]

# Create bases and domain
x_basis = de.Fourier('x',151 , interval=(0, 45000))
y_basis = de.Fourier('y', 151, interval=(0, 45000))
z_basis = de.Chebyshev('z', 51, interval=(0, 20000))
domain = de.Domain([x_basis, y_basis,z_basis], grid_dtype=np.float64)

# Poisson equation
problem = de.LBVP(domain, variables=['ab','abz'])
ncc = domain.new_field(name='rhobar')
ncc['g'] = rhobar
ncc.meta['x','y']['constant'] = True
problem.parameters['rhobar'] = ncc

ncc1 = domain.new_field(name='rhot')
ncc1['g'] = rhot
# ncc1.meta['x','y']['constant'] = True
problem.parameters['rhot'] = ncc1

rhozbar = ncc.differentiate(z=1)
rhozzbar = ncc.differentiate(z=2)

rhozbar.require_grid_space()
rhozzbar.require_grid_space()

problem.parameters['rhozbar'] = rhozbar
problem.parameters['rhozzbar'] = rhozzbar
problem.parameters['g'] = 9.81

problem.add_equation("-rhobar * (dx(dx(ab)) + dy(dy(ab))) -2*dz(ab)*rhozbar - ab*rhozzbar - rhobar*dz(abz)  = g* (dx(dx(rhot)) + dy(dy(rhot)))")
problem.add_equation("abz - dz(ab) = 0")
problem.add_bc("left(ab) = 0")
problem.add_bc("right(ab) = 0")

# Build solver
solver = problem.build_solver()
solver.solve()