Minimal (non)working example

TG.py

from __future__ import print_function
import warnings
from numpy import pi, zeros, sum, float64, sin, cos, prod
from spectralDNS import config, get_solver, solve

try:
    import matplotlib.pyplot as plt

except ImportError:
    warnings.warn("matplotlib not installed")
    plt = None

def initialize(solver, context):
    initialize1(solver, context)
    config.params.t = 0.0
    config.params.tstep = 0

def initialize1(solver, context):
    U, X = context.U, context.X
    U[0] = sin(X[0])*cos(X[1])*cos(X[2])
    U[1] = -cos(X[0])*sin(X[1])*cos(X[2])
    U[2] = 0
    solver.set_velocity(**context)

def energy_fourier(comm, a):
    result = 2*sum(abs(a[..., 1:-1])**2) + sum(abs(a[..., 0])**2) + sum(abs(a[..., -1])**2)
    result = comm.allreduce(result)
    return result


k = []
w = []
im1 = None
kold = zeros(1)
def update(context):
    global k, w, im1
    c = context
    params = config.params
    solver = config.solver

    U = solver.get_velocity(**c)
    kk = solver.comm.reduce(sum(U*U)/prod(params.N)/2) # Compute energy with double precision
    ww2 = energy_fourier(solver.comm, c.U_hat)/2 #prod(params.N)**2/2

    if solver.rank == 0:
        print("kk=%e, ww2 =  %e, difference is %e" % (kk,ww2,(kk - ww2)))

if __name__ == "__main__":
    config.update(
        {'nu': 1./280,             # Viscosity
         'dt': 0.1,                 # Time step
         'T': 5.0,                   # End time
         'L': [2*pi, 2.*pi, 2*pi],
         'M': [6, 6, 6],
         'planner_effort': {'fft': 'FFTW_ESTIMATE',
                            'rfftn': 'FFTW_ESTIMATE',
                            'irfftn': 'FFTW_ESTIMATE'},
        }, "triplyperiodic")
    config.triplyperiodic.add_argument("--compute_energy", type=int, default=2)
    config.triplyperiodic.add_argument("--plot_step", type=int, default=2)
    sol = get_solver(update=update,
                     mesh="triplyperiodic")

    context = sol.get_context()

    # Add curl to the stored results. For this we need to update the update_components
    # method used by the HDF5File class to compute the real fields that are stored
    #context.hdf5file.filename = "NS9"
    #context.hdf5file.results['data'].update({'curl': [context.curl]})
    def update_components(**c):
        """Overload default because we want to store the curl as well"""
        sol.get_velocity(**c)
        sol.get_pressure(**c)
        sol.get_curl(**c)

    #context.hdf5file.update_components = update_components

    initialize(sol, context)
    solve(sol, context)