Orszag Tang Vortex in Dedalus 3

[Sofiane Alla]

Hi everyone, 

I am building a little project in AI for fusion, and I wanted to reproduce the results of the original paper as a basis to build a FNO surrogate: methods_paper_examples/orszag_tang_vortex at master · DedalusProject/methods_paper_examples

I troubleshoot most of the compatibility issues switching from dedalus 2 to 3, but I still face major stability issues. The main problem that I have is matrix dimensionality mismatch between x_basis, y_basis and the initial conditions x and y when the dealiasing factor is 3/2. Also what is the equivalent of Domain.grid(0) in Dedalus 3? I used global_grid but it causes instabilities. 

Here's the code below (hope that helps! I will make it open source once everything is fixed): 

import numpy as np

from mpi4py import MPI

import time

from dedalus import public as de

from dedalus.extras import flow_tools

import logging

logger = logging.getLogger(__name__)

# Parameters

Lx, Ly = (1., 1.)

Prandtl_m = 1.

Prandtl = 1.

Rm = 1e4

gamma = 5/3

c_v = 1/(gamma-1)

c_p = gamma*c_v

# Create coordinate objects

coords = de.CartesianCoordinates('x', 'y')

x = coords['x']

y = coords['y']

# Create a distributor with complex data type

dist = de.Distributor(coords, dtype=np.float64)

dtype=np.float64

# Define Fourier bases for each coordinate

x_basis = de.Fourier(x, size=256, bounds=(0, Lx), dealias=3/2, dtype=dtype)

y_basis = de.Fourier(y, size=256, bounds=(0, Ly), dealias=3/2, dtype=dtype)

# Create fields for your variables

u = dist.Field(name='u', bases=(x_basis, y_basis))

v = dist.Field(name='v', bases=(x_basis, y_basis))

T = dist.Field(name='T', bases=(x_basis, y_basis))

logrho = dist.Field(name='logrho', bases=(x_basis, y_basis))

A = dist.Field(name='A', bases=(x_basis, y_basis))

# Set up the initial value problem (IVP)

problem = de.IVP([u, v, T, logrho, A], time='t')

# Define parameters and assign them to the problem's namespace

eta = 1 / Rm

nu = Prandtl_m / Rm

chi = Prandtl_m / (Prandtl * Rm)

T0 = 1. / gamma

problem.namespace['eta'] = eta

problem.namespace['nu'] = nu

problem.namespace['chi'] = chi

problem.namespace['gamma'] = gamma

problem.namespace['c_v'] = c_v

problem.namespace['T0'] = T0

# Define differential operators using the Differentiate operator

dx = lambda f: de.Differentiate(f, x)

dy = lambda f: de.Differentiate(f, y)

# Assign differential operators to the problem namespace

problem.namespace['dx'] = dx

problem.namespace['dy'] = dy

# Define substitutions directly using the operators

bx = dy(A)

by = -dx(A)

p_m = 0.5 * (dx(A)**2 + dy(A)**2)

div_u = dx(u) + dy(v)

# Assign computed quantities to the problem's namespace

problem.namespace['bx'] = bx

problem.namespace['by'] = by

problem.namespace['p_m'] = p_m

problem.namespace['div_u'] = div_u

# Viscous terms

viscous_u_lhs = nu * (-dx(dx(u)) - dy(dy(u)) - (1./3.) * dx(div_u))

viscous_v_lhs = nu * (-dx(dx(v)) - dy(dy(v)) - (1./3.) * dy(div_u))

viscous_u_rhs = nu * (

    dx(logrho)*dx(u) + dy(logrho)*dy(u) +

    dx(u)*dx(logrho) + dx(v)*dy(logrho) -

    (2./3.)*dx(logrho)*div_u

)

viscous_v_rhs = nu * (

    dx(logrho)*dx(v) + dy(logrho)*dy(v) +

    dy(u)*dx(logrho) + dy(v)*dy(logrho) -

    (2./3.)*dy(logrho)*div_u

)

viscous_heating = (

    2.*dx(u)**2 + dy(u)**2 + dx(v)**2 + 2.*dy(v)**2 +

    2.*dx(v)*dy(u) - (2./3.)*div_u**2

)

ohmic_heating = (dx(dx(A)) + dy(dy(A)))**2

# Assign viscous terms to the namespace

problem.namespace['viscous_u_lhs'] = viscous_u_lhs

problem.namespace['viscous_v_lhs'] = viscous_v_lhs

problem.namespace['viscous_u_rhs'] = viscous_u_rhs

problem.namespace['viscous_v_rhs'] = viscous_v_rhs

problem.namespace['viscous_heating'] = viscous_heating

problem.namespace['ohmic_heating'] = ohmic_heating

# Use de.math.exp for the exponential function

exp_neg_logrho = np.exp(-logrho)

problem.namespace['exp_neg_logrho'] = exp_neg_logrho

# Add equations to the problem

problem.add_equation(

    "dt(u) + dx(T) + T0*dx(logrho) + viscous_u_lhs = "

    "-T*dx(logrho) - (u*dx(u) + v*dy(u)) + viscous_u_rhs + "

    "(bx*dx(bx) + by*dy(bx) - dx(p_m))*exp_neg_logrho"

)

problem.add_equation(

    "dt(v) + dy(T) + T0*dy(logrho) + viscous_v_lhs = "

    "-T*dy(logrho) - (u*dx(v) + v*dy(v)) + viscous_v_rhs + "

    "(bx*dx(by) + by*dy(by) - dy(p_m))*exp_neg_logrho"

)

problem.add_equation(

    "dt(logrho) + div_u = -u*dx(logrho) - v*dy(logrho)"

)

problem.add_equation(

    "dt(T) + (gamma - 1)*T0*div_u - chi/c_v*(dx(dx(T)) + dy(dy(T))) = "

    "-u*dx(T) - v*dy(T) - (gamma - 1)*T*div_u + chi/c_v*(dx(T)*dx(logrho) + dy(T)*dy(logrho)) + "

    "nu/c_v*viscous_heating + eta/c_v*ohmic_heating*exp_neg_logrho"

)

problem.add_equation(

    "dt(A) - eta*(dx(dx(A)) + dy(dy(A))) = u*by - v*bx"

)

# Build solver

timestepper=de.RK443

solver = problem.build_solver(timestepper)

logger.info('Solver built')

import math

# Access evaluator

evaluator = solver.evaluator

# Update the evaluator's vars with problem's namespace

evaluator.vars.update(problem.namespace)

# Assign differential operators and functions to evaluator's vars

evaluator.vars['dx'] = dx

evaluator.vars['dy'] = dy

# Assign fields and parameters to evaluator's vars

evaluator.vars['A'] = A

evaluator.vars['logrho'] = logrho

evaluator.vars['T'] = T

evaluator.vars['T0'] = T0

# Import 'exp' function from dedalus.core.math

evaluator.vars['exp'] = np.exp

# Define 'bx' and 'by' in the evaluator's vars

bx_eval = evaluator.vars['dy'](A)

by_eval = -evaluator.vars['dx'](A)

evaluator.vars['bx'] = bx_eval

evaluator.vars['by'] = by_eval

# Initial conditions

# Retrieve 1D grids at the dealiased scales

x_grid = x_basis.global_grid(dist, scale=1)

y_grid = y_basis.global_grid(dist, scale=1)

# Use the field variables directly

logrho['g'] = np.log(25/(36*np.pi))

u['g'] = - np.sin(2*np.pi*y_grid)

v['g'] =   np.sin(2*np.pi*x_grid)

B0 = 1/np.sqrt(4*np.pi)

A['g'] = B0*( np.cos(4*np.pi*x_grid)/(4*np.pi) + np.cos(2*np.pi*y_grid)/(2*np.pi) )

# Set initial T

T['g'] = 0.0  # Assuming initial temperature perturbation is zero

# Initial timestep

dt = 0.000025

# Integration parameters

solver.stop_sim_time = 1.001

solver.stop_wall_time = np.inf

solver.stop_iteration = np.inf

# Analysis

snapshots = evaluator.add_file_handler('/content/drive/My Drive/snapshots_Re1e4_48', sim_dt=0.01, max_writes=10)

# Add fields

snapshots.add_task(u, name='u')

snapshots.add_task(v, name='v')

snapshots.add_task(T, name='T')

snapshots.add_task(logrho, name='logrho')

snapshots.add_task(A, name='A')

# Add computed quantities

snapshots.add_task(bx, name="Bx")

snapshots.add_task(by, name="By")

snapshots.add_task(np.exp(logrho), name="rho")

snapshots.add_task(T + T0, name="temp")

# CFL

from dedalus.extras.flow_tools import CFL

CFL = CFL(solver, initial_dt=dt, cadence=5, safety=0.6,

          max_change=1.5, min_change=0.5, max_dt=0.05, threshold=0.05)

# Import VectorField

from dedalus.core.field import VectorField

# Create vector field for velocity

velocity = dist.VectorField(coords, name='velocity', bases=(x_basis, y_basis))

velocity['g'][0] = u['g']

velocity['g'][1] = v['g']

# Create fields for bx and by

bx_field = dist.Field(name='bx', bases=(x_basis, y_basis))

by_field = dist.Field(name='by', bases=(x_basis, y_basis))

# Compute bx and by

bx_field['g'] = evaluator.vars['dy'](A).evaluate()['g']

by_field['g'] = (-evaluator.vars['dx'](A)).evaluate()['g']

# Create vector field for magnetic field

magnetic_field = dist.VectorField(coords, name='magnetic_field', bases=(x_basis, y_basis))

magnetic_field['g'][0] = bx_field['g']

magnetic_field['g'][1] = by_field['g']

# Add vector fields to CFL

CFL.add_velocity(velocity)

CFL.add_velocity(magnetic_field)

# Flow properties

# Import 'exp' function from dedalus.core.math

evaluator.vars['sqrt'] = math.sqrt

flow = flow_tools.GlobalFlowProperty(solver, cadence=10)

flow.add_property(np.sqrt(u*u + v*v) / eta, name='Rm')

flow.add_property(np.sqrt(bx*bx + by*by) / eta, name='S')

# Main loop

try:

    logger.info('Starting loop')

    start_time = time.time()

    while solver.proceed:

        dt = CFL.compute_timestep()

        solver.step(dt)

        if (solver.iteration - 1) % 10 == 0:

            logger.info('Iteration: %i, Time: %e, dt: %e' % (solver.iteration, solver.sim_time, dt))

            logger.info('Max Rm = %f' % flow.max('Rm'))

            logger.info('Max S = %f' % flow.max('S'))

except Exception as e:

    logger.error(f'Exception raised, triggering end of main loop: {e}')

    raise e

finally:

    end_time = time.time()

    logger.info('Iterations: %i' % solver.iteration)

    logger.info('Sim end time: %f' % solver.sim_time)

    logger.info('Run time: %.2f sec' % (end_time - start_time))

    logger.info('Run time: %f cpu-hr' % ((end_time - start_time) / 3600 * dist.comm_cart.size))

[Keaton Burns]

Hi Sofiane,

Please take a look at the tutorial notebooks here. They cover issues like setting up initial conditions, basis grids, and field scales.

Best,

-Keaton

--
You received this message because you are subscribed to the Google Groups "Dedalus Users" group.
To unsubscribe from this group and stop receiving emails from it, send an email to dedalus-user...@googlegroups.com.
To view this discussion visit https://groups.google.com/d/msgid/dedalus-users/5590fe03-7329-42aa-9d8d-cbec0db52788n%40googlegroups.com.