def init_problem(self):
        # Initial conditions
        self.u = self.solver.state['u']['g']
        self.v = self.solver.state['v']['g']
        self.w = self.solver.state['w']['g']
        z = self.domain.grid(2)
        b = self.solver.state['b']
        bz = self.solver.state['bz']

        # Random perturbations, initialized globally for same results in parallel
        gshape = self.domain.dist.grid_layout.global_shape(scales=1)
        slices = self.domain.dist.grid_layout.slices(scales=1)
        rand = np.random.RandomState(seed=23)
        noise = rand.standard_normal(gshape)[slices]

        # Linear background + perturbations damped at walls
        zb, zt = self.z_basis.interval
        pert = 1e-3 * noise * (zt - z) * (z - zb)
        b['g'] = -self.F * (z - pert)
        b.differentiate('z', out=bz)

        # Flow properties
        self.CFL = flow_tools.CFL(self.solver,
                                  initial_dt=1e-4,
                                  cadence=5,
                                  safety=0.5,
                                  max_change=1.5,
                                  min_change=0.5,
                                  max_dt=0.05)
        self.CFL.add_velocities(('u', 'v', 'w'))
        self.flow = flow_tools.GlobalFlowProperty(self.solver, cadence=10)
        self.flow.add_property("sqrt(u*u + v*v + w*w) / R", name='Re')
Example #2
0
    def run(self, initial_dt=1e-4, sim_time=100):

        # Integration parameters
        self.solver.stop_sim_time = self.solver.sim_time + sim_time
        self.solver.stop_wall_time = np.inf  # 60 * 60.
        self.solver.stop_iteration = np.inf

        # CFL
        CFL = flow_tools.CFL(self.solver,
                             initial_dt=initial_dt,
                             cadence=5,
                             safety=1.5,
                             max_change=1.5,
                             min_change=0.5,
                             max_dt=0.05)
        CFL.add_velocities(('u', 'v', 'w'))

        flow = flow_tools.GlobalFlowProperty(self.solver, cadence=10)
        flow.add_property("0.5*(u*u + v*v + w*w)", name='KE')

        self.CFL = CFL
        self.flow = flow

        try:
            logger.info('Starting loop')
            start_run_time = time.time()
            while self.solver.ok:
                dt = CFL.compute_dt()
                self.solver.step(dt)
                if (self.solver.iteration - 1) % 100 == 0:
                    logger.info(
                        'Iteration: %i, Time: %e, dt: %e' %
                        (self.solver.iteration, self.solver.sim_time, dt))
                    logger.info('Average KE = %e' % flow.volume_average('KE'))
        except:
            logger.error('Exception raised, triggering end of main loop.')
            raise
        finally:
            end_run_time = time.time()
            logger.info('Iterations: %i' % self.solver.iteration)
            logger.info('Sim end time: %f' % self.solver.sim_time)
            logger.info('Run time: %.2f sec' % (end_run_time - start_run_time))
            logger.info('Run time: %f cpu-hr' %
                        ((end_run_time - start_run_time) / hour *
                         self.domain.dist.comm_cart.size))
def Rayleigh_Benard(Rayleigh=1e6,
                    Prandtl=1,
                    nz=64,
                    nx=None,
                    aspect=4,
                    fixed_flux=False,
                    fixed_T=True,
                    viscous_heating=False,
                    restart=None,
                    run_time=23.5,
                    run_time_buoyancy=50,
                    run_time_iter=np.inf,
                    max_writes=10,
                    max_slice_writes=10,
                    data_dir='./',
                    coeff_output=True,
                    verbose=False,
                    no_join=False):
    import os
    from dedalus.tools.config import config

    config['logging']['filename'] = os.path.join(data_dir, 'logs/dedalus_log')
    config['logging']['file_level'] = 'DEBUG'

    import mpi4py.MPI
    if mpi4py.MPI.COMM_WORLD.rank == 0:
        if not os.path.exists('{:s}/'.format(data_dir)):
            os.makedirs('{:s}/'.format(data_dir))
        logdir = os.path.join(data_dir, 'logs')
        if not os.path.exists(logdir):
            os.mkdir(logdir)
    logger = logging.getLogger(__name__)
    logger.info("saving run in: {}".format(data_dir))

    import time
    from dedalus import public as de
    from dedalus.extras import flow_tools
    from dedalus.tools import post

    # input parameters
    logger.info("Ra = {}, Pr = {}".format(Rayleigh, Prandtl))

    # Parameters
    Lz = 1.
    Lx = aspect * Lz

    if nx is None:
        nx = int(nz * aspect)

    logger.info("resolution: [{}x{}]".format(nx, nz))
    # Create bases and domain
    x_basis = de.Fourier('x', nx, interval=(0, Lx), dealias=3 / 2)
    z_basis = de.Chebyshev('z', nz, interval=(0, Lz), dealias=3 / 2)
    domain = de.Domain([x_basis, z_basis], grid_dtype=np.float64)

    if fixed_flux:
        T_bc_var = 'Tz'
    elif fixed_T:
        T_bc_var = 'T'

    # 2D Boussinesq hydrodynamics
    problem = de.IVP(domain, variables=['Tz', 'T', 'p', 'u', 'w', 'Oy'])
    problem.meta['p', T_bc_var, 'Oy', 'w']['z']['dirichlet'] = True

    T0_z = domain.new_field()
    T0_z.meta['x']['constant'] = True
    T0_z['g'] = -1
    T0 = domain.new_field()
    T0.meta['x']['constant'] = True
    T0['g'] = Lz / 2 - domain.grid(-1)
    problem.parameters['T0'] = T0
    problem.parameters['T0_z'] = T0_z

    problem.parameters['P'] = (Rayleigh * Prandtl)**(-1 / 2)
    problem.parameters['R'] = (Rayleigh / Prandtl)**(-1 / 2)
    problem.parameters['F'] = F = 1

    problem.parameters['Lx'] = Lx
    problem.parameters['Lz'] = Lz
    problem.substitutions['plane_avg(A)'] = 'integ(A, "x")/Lx'
    problem.substitutions['vol_avg(A)'] = 'integ(A)/Lx/Lz'

    problem.substitutions['v'] = '0'
    problem.substitutions['Ox'] = '0'
    problem.substitutions['Oz'] = '(dx(v) )'
    problem.substitutions['Kx'] = '( -dz(Oy))'
    problem.substitutions['Ky'] = '(dz(Ox) - dx(Oz))'
    problem.substitutions['Kz'] = '(dx(Oy) )'

    problem.substitutions['vorticity'] = 'Oy'
    problem.substitutions['enstrophy'] = 'Oy**2'

    problem.substitutions['u_fluc'] = '(u - plane_avg(u))'
    problem.substitutions['w_fluc'] = '(w - plane_avg(w))'
    problem.substitutions['KE'] = '(0.5*(u*u+w*w))'

    problem.substitutions['sigma_xz'] = '(dx(w) + Oy + dx(w))'
    problem.substitutions['sigma_xx'] = '(2*dx(u))'
    problem.substitutions['sigma_zz'] = '(2*dz(w))'

    if viscous_heating:
        problem.substitutions[
            'visc_heat'] = 'R*(sigma_xz**2 + sigma_xx*dx(u) + sigma_zz*dz(w))'
        problem.substitutions['visc_flux_z'] = 'R*(u*sigma_xz + w*sigma_zz)'
    else:
        problem.substitutions['visc_heat'] = '0'
        problem.substitutions['visc_flux_z'] = '0'

    problem.substitutions['conv_flux_z'] = '(w*T + visc_flux_z)/P'
    problem.substitutions['kappa_flux_z'] = '(-Tz)'

    problem.add_equation("Tz - dz(T) = 0")
    problem.add_equation(
        "dt(T) - P*(dx(dx(T)) + dz(Tz)) + w*T0_z    = -(u*dx(T) + w*Tz)  - visc_heat"
    )
    # O == omega = curl(u);  K = curl(O)
    problem.add_equation("dt(u)  + R*Kx  + dx(p)              =  v*Oz - w*Oy ")
    problem.add_equation("dt(w)  + R*Kz  + dz(p)    -T        =  u*Oy - v*Ox ")
    problem.add_equation("dx(u) + dz(w) = 0")
    problem.add_equation("Oy - dz(u) + dx(w) = 0")

    problem.add_bc("right(p) = 0", condition="(nx == 0)")
    if fixed_flux:
        problem.add_bc("left(Tz)  = 0")
        problem.add_bc("right(Tz) = 0")
    elif fixed_T:
        problem.add_bc("left(T)  = 0")
        problem.add_bc("right(T) = 0")
    problem.add_bc("left(Oy) = 0")
    problem.add_bc("right(Oy) = 0")
    problem.add_bc("left(w)  = 0")
    problem.add_bc("right(w) = 0", condition="(nx != 0)")

    # Build solver
    ts = de.timesteppers.RK443
    cfl_safety = 1

    solver = problem.build_solver(ts)
    logger.info('Solver built')

    checkpoint = solver.evaluator.add_file_handler(data_dir + 'checkpoints',
                                                   wall_dt=8 * 3600,
                                                   max_writes=1)
    checkpoint.add_system(solver.state, layout='c')

    # Initial conditions
    x = domain.grid(0)
    z = domain.grid(1)
    T = solver.state['T']
    Tz = solver.state['Tz']

    # Random perturbations, initialized globally for same results in parallel
    noise = global_noise(domain, scale=1, frac=0.25)

    if restart is None:
        # Linear background + perturbations damped at walls
        zb, zt = z_basis.interval
        pert = 1e-3 * noise * (zt - z) * (z - zb)
        T['g'] = pert
        T.differentiate('z', out=Tz)
    else:
        logger.info("restarting from {}".format(restart))
        checkpoint.restart(restart, solver)

    # Integration parameters
    solver.stop_sim_time = run_time_buoyancy
    solver.stop_wall_time = run_time * 3600.
    solver.stop_iteration = run_time_iter

    # Analysis
    analysis_tasks = []
    snapshots = solver.evaluator.add_file_handler(data_dir + 'slices',
                                                  sim_dt=0.1,
                                                  max_writes=max_slice_writes)
    snapshots.add_task("T + T0", name='T')
    snapshots.add_task("enstrophy")
    snapshots.add_task("vorticity")
    analysis_tasks.append(snapshots)

    snapshots_small = solver.evaluator.add_file_handler(
        data_dir + 'slices_small', sim_dt=0.1, max_writes=max_slice_writes)
    snapshots_small.add_task("T + T0", name='T', scales=0.25)
    snapshots_small.add_task("enstrophy", scales=0.25)
    snapshots_small.add_task("vorticity", scales=0.25)
    analysis_tasks.append(snapshots_small)

    if coeff_output:
        coeffs = solver.evaluator.add_file_handler(data_dir + 'coeffs',
                                                   sim_dt=0.1,
                                                   max_writes=max_slice_writes)
        coeffs.add_task("T+T0", name="T", layout='c')
        coeffs.add_task("T - plane_avg(T)", name="T'", layout='c')
        coeffs.add_task("w", layout='c')
        coeffs.add_task("u", layout='c')
        coeffs.add_task("enstrophy", layout='c')
        coeffs.add_task("vorticity", layout='c')
        analysis_tasks.append(coeffs)

    profiles = solver.evaluator.add_file_handler(data_dir + 'profiles',
                                                 sim_dt=0.1,
                                                 max_writes=max_writes)
    profiles.add_task("plane_avg(T+T0)", name="T")
    profiles.add_task("plane_avg(T)", name="T'")
    profiles.add_task("plane_avg(u)", name="u")
    profiles.add_task("plane_avg(w)", name="w")
    profiles.add_task("plane_avg(enstrophy)", name="enstrophy")
    # This may have an error:
    profiles.add_task("plane_avg(conv_flux_z)/plane_avg(kappa_flux_z) + 1",
                      name="Nu")
    profiles.add_task("plane_avg(conv_flux_z) + plane_avg(kappa_flux_z)",
                      name="Nu_2")

    analysis_tasks.append(profiles)

    scalar = solver.evaluator.add_file_handler(data_dir + 'scalar',
                                               sim_dt=0.1,
                                               max_writes=max_writes)
    scalar.add_task("vol_avg(T)", name="IE")
    scalar.add_task("vol_avg(KE)", name="KE")
    scalar.add_task("vol_avg(-T*z)", name="PE_fluc")
    scalar.add_task("vol_avg(T) + vol_avg(KE) + vol_avg(-T*z)", name="TE")
    scalar.add_task("0.5*vol_avg(u_fluc*u_fluc+w_fluc*w_fluc)", name="KE_fluc")
    scalar.add_task("0.5*vol_avg(u*u)", name="KE_x")
    scalar.add_task("0.5*vol_avg(w*w)", name="KE_z")
    scalar.add_task("0.5*vol_avg(u_fluc*u_fluc)", name="KE_x_fluc")
    scalar.add_task("0.5*vol_avg(w_fluc*w_fluc)", name="KE_z_fluc")
    scalar.add_task("vol_avg(plane_avg(u)**2)", name="u_avg")
    scalar.add_task("vol_avg((u - plane_avg(u))**2)", name="u1")
    scalar.add_task("vol_avg(conv_flux_z) + 1.", name="Nu")
    analysis_tasks.append(scalar)

    # workaround for issue #29
    problem.namespace['enstrophy'].store_last = True

    # CFL
    CFL = flow_tools.CFL(solver,
                         initial_dt=0.1,
                         cadence=1,
                         safety=cfl_safety,
                         max_change=1.5,
                         min_change=0.5,
                         max_dt=0.1,
                         threshold=0.1)
    CFL.add_velocities(('u', 'w'))

    # Flow properties
    flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
    flow.add_property("sqrt(u*u + w*w) / R", name='Re')

    first_step = True
    # Main loop
    try:
        logger.info('Starting loop')
        Re_avg = 0
        while solver.ok and np.isfinite(Re_avg):
            dt = CFL.compute_dt()
            solver.step(dt)  #, trim=True)
            Re_avg = flow.grid_average('Re')
            log_string = 'Iteration: {:5d}, '.format(solver.iteration)
            log_string += 'Time: {:8.3e}, dt: {:8.3e}, '.format(
                solver.sim_time, dt)
            log_string += 'Re: {:8.3e}/{:8.3e}'.format(Re_avg, flow.max('Re'))
            logger.info(log_string)

            if first_step:
                if verbose:
                    import matplotlib
                    matplotlib.use('Agg')
                    import matplotlib.pyplot as plt
                    fig = plt.figure()
                    ax = fig.add_subplot(1, 1, 1)
                    ax.spy(solver.pencils[0].L,
                           markersize=1,
                           markeredgewidth=0.0)
                    fig.savefig(data_dir + "sparsity_pattern.png", dpi=1200)

                    import scipy.sparse.linalg as sla
                    LU = sla.splu(solver.pencils[0].LHS.tocsc(),
                                  permc_spec='NATURAL')
                    fig = plt.figure()
                    ax = fig.add_subplot(1, 2, 1)
                    ax.spy(LU.L.A, markersize=1, markeredgewidth=0.0)
                    ax = fig.add_subplot(1, 2, 2)
                    ax.spy(LU.U.A, markersize=1, markeredgewidth=0.0)
                    fig.savefig(data_dir + "sparsity_pattern_LU.png", dpi=1200)

                    logger.info("{} nonzero entries in LU".format(LU.nnz))
                    logger.info("{} nonzero entries in LHS".format(
                        solver.pencils[0].LHS.tocsc().nnz))
                    logger.info("{} fill in factor".format(
                        LU.nnz / solver.pencils[0].LHS.tocsc().nnz))
                first_step = False
                start_time = time.time()
    except:
        logger.error('Exception raised, triggering end of main loop.')
        raise
    finally:
        end_time = time.time()
        main_loop_time = end_time - start_time
        n_iter_loop = solver.iteration - 1
        logger.info('Iterations: {:d}'.format(n_iter_loop))
        logger.info('Sim end time: {:f}'.format(solver.sim_time))
        logger.info('Run time: {:f} sec'.format(main_loop_time))
        logger.info('Run time: {:f} cpu-hr'.format(main_loop_time / 60 / 60 *
                                                   domain.dist.comm_cart.size))
        logger.info('iter/sec: {:f} (main loop only)'.format(n_iter_loop /
                                                             main_loop_time))

        final_checkpoint = solver.evaluator.add_file_handler(
            data_dir + 'final_checkpoint', iter=1)
        final_checkpoint.add_system(solver.state, layout='c')
        solver.step(dt)  #clean this up in the future...works for now.
        post.merge_analysis(data_dir + 'final_checkpoint')

        if not no_join:
            logger.info('beginning join operation')
            post.merge_analysis(data_dir + 'checkpoints')

            for task in analysis_tasks:
                logger.info(task.base_path)
                post.merge_analysis(task.base_path)

        logger.info(40 * "=")
        logger.info('Iterations: {:d}'.format(n_iter_loop))
        logger.info('Sim end time: {:f}'.format(solver.sim_time))
        logger.info('Run time: {:f} sec'.format(main_loop_time))
        logger.info('Run time: {:f} cpu-hr'.format(main_loop_time / 60 / 60 *
                                                   domain.dist.comm_cart.size))
        logger.info('iter/sec: {:f} (main loop only)'.format(n_iter_loop /
                                                             main_loop_time))
timeseries = IVP.evaluator.add_file_handler(os.path.join(
    data_dir, 'timeseries'),
                                            iter=10,
                                            max_writes=np.inf)
timeseries.add_task("vol_avg(sqrt(u*u))", name='urms')
timeseries.add_task("vol_avg(sqrt(v*v))", name='vrms')
timeseries.add_task("vol_avg(0.5*(u*u+v*v))", name='Ekin')
timeseries.add_task("vol_avg((u*u)/(u*u+v*v))", name='Ekin_x_ratio')
if mhd:
    timeseries.add_task("vol_avg(0.5*(Bx*Bx+By*By))", name='Mag_en')
    timeseries.add_task("vol_avg(sqrt(Bx*Bx))", name='Bx')
    timeseries.add_task("vol_avg(sqrt(By*By))", name='By')
analysis_tasks = [snapshots, slices, timeseries]

flow = flow_tools.GlobalFlowProperty(IVP, cadence=10)
flow.add_property("0.5*(u*u + v*v)", name="Ekin")
flow.add_property("u", name="u")
flow.add_property("v", name="v")
# Main loop
try:
    logger.info('Starting loop')
    start_time = time.time()
    while IVP.ok:
        dt = CFL.compute_dt()
        dt = IVP.step(dt)
        if (IVP.iteration - 1) % 10 == 0:
            logger.info(
                'Iteration: %i, Time: %e, dt: %e, Kinetic Energy: %e, u max: %e, v max: %e'
                % (IVP.iteration, IVP.sim_time, dt,
                   flow.volume_average('Ekin'), flow.max('u'), flow.max('v')))
Example #5
0
def main():
    args = get_args()
    # Parameters
    Lx, Lz = (args.lx, args.lz)
    Prandtl = args.prandtl
    Rayleigh = args.rayleigh
    seed = args.seed

    # Create bases and domain
    x_basis = de.Fourier('x', args.res_x, interval=(0, Lx), dealias=3/2)  # 256
    z_basis = de.Chebyshev('z', args.res_z, interval=(-Lz/2, Lz/2), dealias=3/2)  # 64
    domain = de.Domain([x_basis, z_basis], grid_dtype=np.float64)

    # 2D Boussinesq hydrodynamics
    problem = de.IVP(domain, variables=['p','b','u','w','bz','uz','wz'])
    problem.meta['p','b','u','w']['z']['dirichlet'] = True
    problem.parameters['P'] = (Rayleigh * Prandtl)**(-1/2)
    problem.parameters['R'] = (Rayleigh / Prandtl)**(-1/2)
    problem.parameters['F'] = F = 1
    problem.add_equation("dx(u) + wz = 0")
    # problem.add_equation("dt(b) - P*(dx(dx(b)) + dz(bz)) - F*w       = -(u*dx(b) + w*bz)")
    problem.add_equation("dt(b) - P*(dx(dx(b)) + dz(bz))             = -(u*dx(b) + w*bz)")
    problem.add_equation("dt(u) - R*(dx(dx(u)) + dz(uz)) + dx(p)     = -(u*dx(u) + w*uz)")
    problem.add_equation("dt(w) - R*(dx(dx(w)) + dz(wz)) + dz(p) - b = -(u*dx(w) + w*wz)")
    problem.add_equation("bz - dz(b) = 0")
    problem.add_equation("uz - dz(u) = 0")
    problem.add_equation("wz - dz(w) = 0")
    problem.add_bc("left(b) = 0.5")
    problem.add_bc("left(u) = 0")
    problem.add_bc("left(w) = 0")
    problem.add_bc("right(b) = -0.5")
    problem.add_bc("right(u) = 0")
    problem.add_bc("right(w) = 0", condition="(nx != 0)")
    problem.add_bc("right(p) = 0", condition="(nx == 0)")

    # Build solver
    solver = problem.build_solver(de.timesteppers.RK222)
    logger.info('Solver built')

    # Initial conditions or restart
    if not pathlib.Path('restart.h5').exists():

        # Initial conditions
        x = domain.grid(0)
        z = domain.grid(1)
        b = solver.state['b']
        bz = solver.state['bz']

        # Random perturbations, initialized globally for same results in parallel
        gshape = domain.dist.grid_layout.global_shape(scales=1)
        slices = domain.dist.grid_layout.slices(scales=1)
        rand = np.random.RandomState(seed=seed)
        noise = rand.standard_normal(gshape)[slices]

        # Linear background + perturbations damped at walls
        zb, zt = z_basis.interval
        pert =  1e-3 * noise * (zt - z) * (z - zb)
        b['g'] += F * pert
        b.differentiate('z', out=bz)

        # Timestepping and output
        dt = args.dt
        stop_sim_time = args.stop_sim_time
        fh_mode = 'overwrite'

    else:
        # Restart
        write, last_dt = solver.load_state('restart.h5', -1)

        # Timestepping and output
        dt = last_dt
        stop_sim_time = args.stop_sim_time
        fh_mode = 'append'

    # Integration parameters
    solver.stop_sim_time = stop_sim_time

    # Analysis
    snapshots = solver.evaluator.add_file_handler('snapshots', sim_dt=0.05, max_writes=250, mode=fh_mode)
    snapshots.add_system(solver.state)

    # CFL
    CFL = flow_tools.CFL(solver, initial_dt=dt, cadence=1, safety=1,
                         max_change=1.5, min_change=0., max_dt=0.125, threshold=0.05)
    CFL.add_velocities(('u', 'w'))

    # Flow properties
    flow = flow_tools.GlobalFlowProperty(solver, cadence=10)
    flow.add_property("sqrt(u*u + w*w) / R", name='Re')

    # Main loop
    try:
        logger.info('Starting loop')
        start_time = time.time()
        while solver.proceed:
            dt = CFL.compute_dt()
            dt = 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 Re = %f' %flow.max('Re'))
    except:
        logger.error('Exception raised, triggering end of main loop.')
        raise
    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)/60/60*domain.dist.comm_cart.size))
traces = solver.evaluator.add_file_handler(data_dir + '/traces',
                                           sim_dt=0.1,
                                           max_writes=np.inf)
traces.add_task(avg(0.5 * ρ * dot(u, u)), name='KE')
traces.add_task(avg(IE), name='IE')
traces.add_task(avg(Re), name='Re')
traces.add_task(avg(ω**2), name='enstrophy')
traces.add_task(x_avg(np.sqrt(dot(τu1, τu1))), name='τu1')
traces.add_task(x_avg(np.sqrt(dot(τu2, τu2))), name='τu2')
traces.add_task(x_avg(np.sqrt(τs1**2)), name='τs1')
traces.add_task(x_avg(np.sqrt(τs2**2)), name='τs2')

report_cadence = 10
good_solution = True

flow = flow_tools.GlobalFlowProperty(solver, cadence=report_cadence)
flow.add_property(Re, name='Re')
flow.add_property(KE, name='KE')
flow.add_property(IE, name='IE')
flow.add_property(τu1, name='τu1')
flow.add_property(τu2, name='τu2')
flow.add_property(τs1, name='τs1')
flow.add_property(τs2, name='τs2')

KE_avg = 0
while solver.proceed and good_solution:
    # advance
    solver.step(Δt)
    if solver.iteration % report_cadence == 0:
        KE_avg = flow.grid_average('KE')
        IE_avg = flow.grid_average('IE')
Example #7
0
def FC_convection(Rayleigh=1e6, Prandtl=1, stiffness=1e4, m_rz=3, gamma=5/3,
                      n_rho_cz=3.5, n_rho_rz=1, 
                      nz_cz=128, nz_rz=128,
                      nx = None,
                      width=None,
                      single_chebyshev=False,
                      rk222=False,
                      superstep=False,
                      dense=False, nz_dense=64,
                      oz=False,
                      fixed_flux=False,
                      run_time=23.5, run_time_buoyancies=np.inf, run_time_iter=np.inf,
                      dynamic_diffusivities=False,
                      max_writes=20,out_cadence=0.1, no_coeffs=False, no_join=False,
                      restart=None, data_dir='./', verbose=False, label=None):

    def format_number(number, no_format_min=0.1, no_format_max=10):
        if number > no_format_max or number < no_format_min:
            try:
                mantissa = "{:e}".format(number).split("+")[0].split("e")[0].rstrip("0") or "0"
                power    = "{:e}".format(number).split("+")[1].lstrip("0") or "0"
            except:
                mantissa = "{:e}".format(number).split("-")[0].split("e")[0].rstrip("0") or "0"
                power    = "{:e}".format(number).split("-")[1].lstrip("0") or "0"
                power    = "-"+power
            if mantissa[-1]==".":
                mantissa = mantissa[:-1]
            mantissa += "e"
        else:
            mantissa = "{:f}".format(number).rstrip("0") or "0"
            if mantissa[-1]==".":
                mantissa = mantissa[:-1]
            power = ""
        number_string = mantissa+power
        return number_string
     
    # save data in directory named after script
    if data_dir[-1] != '/':
        data_dir += '/'
    data_dir += sys.argv[0].split('.py')[0]
    if fixed_flux:
        data_dir += '_flux'
    if dynamic_diffusivities:
        data_dir += '_dynamic'
    if oz:
        data_dir += '_oz'
    data_dir += "_nrhocz{}_Ra{}_S{}".format(format_number(n_rho_cz), format_number(Rayleigh), format_number(stiffness))
    if width:
        data_dir += "_erf{}".format(format_number(width))
    if label:
        data_dir += "_{}".format(label)
    data_dir += '/'

    from dedalus.tools.config import config
    
    config['logging']['filename'] = os.path.join(data_dir,'logs/dedalus_log')
    config['logging']['file_level'] = 'DEBUG'

    import mpi4py.MPI
    if mpi4py.MPI.COMM_WORLD.rank == 0:
        if not os.path.exists('{:s}/'.format(data_dir)):
            os.makedirs('{:s}/'.format(data_dir))
        logdir = os.path.join(data_dir,'logs')
        if not os.path.exists(logdir):
            os.mkdir(logdir)
    logger = logging.getLogger(__name__)
    logger.info("saving run in: {}".format(data_dir))
    
    import dedalus.public as de
    from dedalus.tools  import post
    from dedalus.extras import flow_tools


    from dedalus.core.future import FutureField
    from stratified_dynamics import multitropes
    from tools.checkpointing import Checkpoint
    
    checkpoint_min = 30
    
    initial_time = time.time()

    logger.info("Starting Dedalus script {:s}".format(sys.argv[0]))

    constant_Prandtl=True
    mixed_temperature_flux=None
    if oz:
        stable_top=True
        if not fixed_flux:
            mixed_temperature_flux=True
    else:
        stable_top=False
        
    # Set domain
    if nx is None:
        nx = nz_cz*4
        
    if single_chebyshev:
        nz = nz_cz
        nz_list = [nz_cz]
    else:
        if dense:
            nz = nz_rz+nz_dense+nz_cz
            #nz_list = [nz_rz, int(nz_dense/2), int(nz_dense/2), nz_cz]
            nz_list = [nz_rz, nz_dense, nz_cz]
        else:
            nz = nz_rz+nz_cz
            nz_list = [nz_rz, nz_cz]
    
    if dynamic_diffusivities:
        atmosphere = multitropes.FC_multitrope_2d_kappa_mu(nx=nx, nz=nz_list, stiffness=stiffness, m_rz=m_rz, gamma=gamma,
                                         n_rho_cz=n_rho_cz, n_rho_rz=n_rho_rz, 
                                         verbose=verbose, width=width,
                                         constant_Prandtl=constant_Prandtl,
                                         stable_top=stable_top)
    else:
        atmosphere = multitropes.FC_multitrope(nx=nx, nz=nz_list, stiffness=stiffness, m_rz=m_rz, gamma=gamma,
                                         n_rho_cz=n_rho_cz, n_rho_rz=n_rho_rz, 
                                         verbose=verbose, width=width,
                                         constant_Prandtl=constant_Prandtl,
                                         stable_top=stable_top)
    
    atmosphere.set_IVP_problem(Rayleigh, Prandtl)
        
    atmosphere.set_BC(mixed_temperature_flux=mixed_temperature_flux, fixed_flux=fixed_flux)
    problem = atmosphere.get_problem()

        
    if atmosphere.domain.distributor.rank == 0:
        if not os.path.exists('{:s}/'.format(data_dir)):
            os.makedirs('{:s}/'.format(data_dir))

    if rk222:
        logger.info("timestepping using RK222")
        ts = de.timesteppers.RK222
        cfl_safety_factor = 0.2*2
    else:
        logger.info("timestepping using RK443")
        ts = de.timesteppers.RK443
        cfl_safety_factor = 0.2*4

    # Build solver
    solver = problem.build_solver(ts)

    # initial conditions
    if restart is None:
        mode = "overwrite"
    else:
        mode = "append"
        
    logger.info("checkpointing in {}".format(data_dir))
    checkpoint = Checkpoint(data_dir)

    if restart is None:
        atmosphere.set_IC(solver)
        dt = None
    else:
        logger.info("restarting from {}".format(restart))
        dt = checkpoint.restart(restart, solver)
        
    checkpoint.set_checkpoint(solver, wall_dt=checkpoint_min*60, mode=mode)
    
    logger.info("thermal_time = {:g}, top_thermal_time = {:g}".format(atmosphere.thermal_time, atmosphere.top_thermal_time))
    
    max_dt = atmosphere.min_BV_time 
    max_dt = atmosphere.buoyancy_time*out_cadence
    if dt is None: dt = max_dt/5
        
    report_cadence = 1
    output_time_cadence = out_cadence*atmosphere.buoyancy_time
    solver.stop_sim_time  = solver.sim_time + run_time_buoyancies*atmosphere.buoyancy_time
    solver.stop_iteration = solver.iteration + run_time_iter
    solver.stop_wall_time = run_time*3600

    logger.info("output cadence = {:g}".format(output_time_cadence))

    analysis_tasks = atmosphere.initialize_output(solver, data_dir, coeffs_output=not(no_coeffs), sim_dt=output_time_cadence, max_writes=max_writes, mode=mode)

    
    cfl_cadence = 1
    CFL = flow_tools.CFL(solver, initial_dt=dt, cadence=cfl_cadence, safety=cfl_safety_factor,
                         max_change=1.5, min_change=0.5, max_dt=max_dt, threshold=0.1)

    if superstep:
        CFL_traditional = flow_tools.CFL(solver, initial_dt=max_dt, cadence=cfl_cadence, safety=cfl_safety_factor,
                                         max_change=1.5, min_change=0.5, max_dt=max_dt, threshold=0.1)

        CFL_traditional.add_velocities(('u', 'w'))
    
        vel_u = FutureField.parse('u', CFL.solver.evaluator.vars, CFL.solver.domain)
        delta_x = atmosphere.Lx/nx
        CFL.add_frequency(vel_u/delta_x)
        vel_w = FutureField.parse('w', CFL.solver.evaluator.vars, CFL.solver.domain)
        mean_delta_z_cz = atmosphere.Lz_cz/nz_cz
        CFL.add_frequency(vel_w/mean_delta_z_cz)
    else:
        CFL.add_velocities(('u', 'w'))


    
    # Flow properties
    flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
    flow.add_property("Re_rms", name='Re')

    try:
        logger.info("starting main loop")
        start_time = time.time()
        start_iter = solver.iteration
        good_solution = True
        first_step = True
        while solver.ok and good_solution:

            dt = CFL.compute_dt()
            # advance
            solver.step(dt)

            effective_iter = solver.iteration - start_iter

            # update lists
            if effective_iter % report_cadence == 0:
                Re_avg = flow.grid_average('Re')
                log_string = 'Iteration: {:5d}, Time: {:8.3e} ({:8.3e}), '.format(solver.iteration, solver.sim_time, solver.sim_time/atmosphere.buoyancy_time)
                log_string += 'dt: {:8.3e}'.format(dt)
                if superstep:
                    dt_traditional = CFL_traditional.compute_dt()
                    log_string += ' (vs {:8.3e})'.format(dt_traditional)
                log_string += ', '
                log_string += 'Re: {:8.3e}/{:8.3e}'.format(Re_avg, flow.max('Re'))
                logger.info(log_string)

            if not np.isfinite(Re_avg):
                good_solution = False
                logger.info("Terminating run.  Trapped on Reynolds = {}".format(Re_avg))
                
            if first_step:
                if verbose:
                    import matplotlib
                    matplotlib.use('Agg')
                    import matplotlib.pyplot as plt
                    fig = plt.figure()
                    ax = fig.add_subplot(1,1,1)
                    ax.spy(solver.pencils[0].L, markersize=0.5, markeredgewidth=0.0)
                    fig.savefig(data_dir+"sparsity_pattern.png", dpi=2400)
                    #fig.savefig(data_dir+"sparsity_pattern.svg", format="svg")

                    import scipy.sparse.linalg as sla
                    LU = sla.splu(solver.pencils[0].LHS.tocsc(), permc_spec='NATURAL')
                    fig = plt.figure()
                    ax = fig.add_subplot(1,2,1)
                    ax.spy(LU.L.A, markersize=1, markeredgewidth=0.0)
                    ax = fig.add_subplot(1,2,2)
                    ax.spy(LU.U.A, markersize=1, markeredgewidth=0.0)
                    fig.savefig(data_dir+"sparsity_pattern_LU.png", dpi=1200)
                    #fig.savefig(data_dir+"sparsity_pattern_LU.svg", format="svg")

                    logger.info("{} nonzero entries in LU".format(LU.nnz))
                    logger.info("{} nonzero entries in LHS".format(solver.pencils[0].LHS.tocsc().nnz))
                    logger.info("{} fill in factor".format(LU.nnz/solver.pencils[0].LHS.tocsc().nnz))
                first_step = False
                start_time = time.time()

    except:
        logger.error('Exception raised, triggering end of main loop.')
        raise
    finally:
        end_time = time.time()

        # Print statistics
        elapsed_time = end_time - start_time
        elapsed_sim_time = solver.sim_time
        N_iterations = solver.iteration - 1
        logger.info('main loop time: {:e}'.format(elapsed_time))
        logger.info('Iterations: {:d}'.format(N_iterations))
        logger.info('iter/sec: {:g}'.format(N_iterations/(elapsed_time)))
        if N_iterations > 0:
            logger.info('Average timestep: {:e}'.format(elapsed_sim_time / N_iterations))
        
        logger.info('beginning join operation')
        try:
            final_checkpoint = Checkpoint(data_dir, checkpoint_name='final_checkpoint')
            final_checkpoint.set_checkpoint(solver, wall_dt=1, mode="append")
            solver.step(dt) #clean this up in the future...works for now.
            post.merge_process_files(data_dir+'/final_checkpoint/')
        except:
            print('cannot save final checkpoint')

        if not(no_join):
            logger.info(data_dir+'/checkpoint/')
            post.merge_process_files(data_dir+'/checkpoint/')

            for task in analysis_tasks:
                logger.info(analysis_tasks[task].base_path)
                post.merge_process_files(analysis_tasks[task].base_path)

        if (atmosphere.domain.distributor.rank==0):
 
            N_TOTAL_CPU = atmosphere.domain.distributor.comm_cart.size

            # Print statistics
            print('-' * 40)
            total_time = end_time-initial_time
            main_loop_time = end_time - start_time
            startup_time = start_time-initial_time
            n_steps = solver.iteration-1
            print('  startup time:', startup_time)
            print('main loop time:', main_loop_time)
            print('    total time:', total_time)
            if n_steps > 0:
                print('    iterations:', n_steps)
                print(' loop sec/iter:', main_loop_time/n_steps)
                print('    average dt:', solver.sim_time/n_steps)
                print("          N_cores, Nx, Nz, startup     main loop,   main loop/iter, main loop/iter/grid, n_cores*main loop/iter/grid")
                print('scaling:',
                    ' {:d} {:d} {:d}'.format(N_TOTAL_CPU,nx,nz),
                    ' {:8.3g} {:8.3g} {:8.3g} {:8.3g} {:8.3g}'.format(startup_time,
                                                                    main_loop_time, 
                                                                    main_loop_time/n_steps, 
                                                                    main_loop_time/n_steps/(nx*nz), 
                                                                    N_TOTAL_CPU*main_loop_time/n_steps/(nx*nz)))
            print('-' * 40)
    return data_dir
Example #8
0
def FC_polytrope(Rayleigh=1e4,
                 Prandtl=1,
                 aspect_ratio=4,
                 Taylor=None,
                 theta=0,
                 nz=128,
                 nx=None,
                 ny=None,
                 threeD=False,
                 mesh=None,
                 n_rho_cz=3,
                 epsilon=1e-4,
                 gamma=5 / 3,
                 run_time=23.5,
                 run_time_buoyancies=None,
                 run_time_iter=np.inf,
                 fixed_T=False,
                 fixed_flux=False,
                 mixed_flux_T=False,
                 const_mu=True,
                 const_kappa=True,
                 dynamic_diffusivities=False,
                 split_diffusivities=False,
                 restart=None,
                 start_new_files=False,
                 rk222=False,
                 safety_factor=0.2,
                 max_writes=20,
                 no_slip=False,
                 data_dir='./',
                 out_cadence=0.1,
                 no_coeffs=False,
                 no_volumes=False,
                 no_join=False,
                 verbose=False):

    import dedalus.public as de
    from dedalus.tools import post
    from dedalus.extras import flow_tools

    import time
    import os
    import sys
    from stratified_dynamics import polytropes
    from tools.checkpointing import Checkpoint

    checkpoint_min = 30

    initial_time = time.time()

    logger.info("Starting Dedalus script {:s}".format(sys.argv[0]))

    if nx is None:
        nx = int(np.round(nz * aspect_ratio))
    if threeD and ny is None:
        ny = nx

    if threeD:
        atmosphere = polytropes.FC_polytrope_3d(nx=nx, ny=ny, nz=nz, mesh=mesh, constant_kappa=const_kappa, constant_mu=const_mu,\
                                        epsilon=epsilon, gamma=gamma, n_rho_cz=n_rho_cz, aspect_ratio=aspect_ratio,\
                                        fig_dir=data_dir)
    else:
        if dynamic_diffusivities:
            atmosphere = polytropes.FC_polytrope_2d_kappa_mu(nx=nx, nz=nz, constant_kappa=const_kappa, constant_mu=const_mu,\
                                        epsilon=epsilon, gamma=gamma, n_rho_cz=n_rho_cz, aspect_ratio=aspect_ratio,\
                                        fig_dir=data_dir)
        else:
            atmosphere = polytropes.FC_polytrope_2d(nx=nx, nz=nz, constant_kappa=const_kappa, constant_mu=const_mu,\
                                        epsilon=epsilon, gamma=gamma, n_rho_cz=n_rho_cz, aspect_ratio=aspect_ratio,\
                                        fig_dir=data_dir)
    if epsilon < 1e-4:
        ncc_cutoff = 1e-14
    elif epsilon > 1e-1:
        ncc_cutoff = 1e-6
    else:
        ncc_cutoff = 1e-10

    if threeD:
        atmosphere.set_IVP_problem(Rayleigh,
                                   Prandtl,
                                   Taylor=Taylor,
                                   theta=theta,
                                   ncc_cutoff=ncc_cutoff,
                                   split_diffusivities=split_diffusivities)
    else:
        atmosphere.set_IVP_problem(Rayleigh,
                                   Prandtl,
                                   ncc_cutoff=ncc_cutoff,
                                   split_diffusivities=split_diffusivities)

    bc_dict = {
        'stress_free': False,
        'no_slip': False,
        'fixed_flux': False,
        'mixed_flux_temperature': False,
        'fixed_temperature': False
    }
    if no_slip:
        bc_dict['no_slip'] = True
    else:
        bc_dict['stress_free'] = True

    if fixed_flux:
        bc_dict['fixed_flux'] = True
    elif mixed_flux_T:
        bc_dict['mixed_flux_temperature'] = True
    else:
        bc_dict['fixed_temperature'] = True
    atmosphere.set_BC(**bc_dict)

    problem = atmosphere.get_problem()

    if atmosphere.domain.distributor.rank == 0:
        if not os.path.exists('{:s}/'.format(data_dir)):
            os.mkdir('{:s}/'.format(data_dir))

    if rk222:
        logger.info("timestepping using RK222")
        ts = de.timesteppers.RK222
        cfl_safety_factor = safety_factor * 2
    else:
        logger.info("timestepping using RK443")
        ts = de.timesteppers.RK443
        cfl_safety_factor = safety_factor * 4

    # Build solver
    solver = problem.build_solver(ts)

    #Check atmosphere
    logger.info("thermal_time = {:g}, top_thermal_time = {:g}".format(atmosphere.thermal_time,\
                                                                    atmosphere.top_thermal_time))
    logger.info("full atm HS check")
    atmosphere.check_atmosphere(make_plots=False,
                                rho=atmosphere.get_full_rho(solver),
                                T=atmosphere.get_full_T(solver))

    if restart is None or start_new_files:
        mode = "overwrite"
    else:
        mode = "append"

    logger.info('checkpointing in {}'.format(data_dir))
    checkpoint = Checkpoint(data_dir)

    if restart is None:
        atmosphere.set_IC(solver)
        dt = None
    else:
        logger.info("restarting from {}".format(restart))
        dt = checkpoint.restart(restart, solver)

    checkpoint.set_checkpoint(solver, wall_dt=checkpoint_min * 60, mode=mode)

    if run_time_buoyancies != None:
        solver.stop_sim_time = solver.sim_time + run_time_buoyancies * atmosphere.buoyancy_time
    else:
        solver.stop_sim_time = 100 * atmosphere.thermal_time

    solver.stop_iteration = solver.iteration + run_time_iter
    solver.stop_wall_time = run_time * 3600
    report_cadence = 1
    output_time_cadence = out_cadence * atmosphere.buoyancy_time
    Hermitian_cadence = 100

    logger.info("stopping after {:g} time units".format(solver.stop_sim_time))
    logger.info("output cadence = {:g}".format(output_time_cadence))

    if threeD:
        analysis_tasks = atmosphere.initialize_output(
            solver,
            data_dir,
            sim_dt=output_time_cadence,
            coeffs_output=not (no_coeffs),
            mode=mode,
            max_writes=max_writes,
            volumes_output=not (no_volumes))
    else:
        analysis_tasks = atmosphere.initialize_output(
            solver,
            data_dir,
            sim_dt=output_time_cadence,
            coeffs_output=not (no_coeffs),
            mode=mode,
            max_writes=max_writes)

    #Set up timestep defaults
    max_dt = output_time_cadence
    if dt is None: dt = max_dt

    cfl_cadence = 1
    cfl_threshold = 0.1
    CFL = flow_tools.CFL(solver,
                         initial_dt=dt,
                         cadence=cfl_cadence,
                         safety=cfl_safety_factor,
                         max_change=1.5,
                         min_change=0.5,
                         max_dt=max_dt,
                         threshold=cfl_threshold)
    if threeD:
        CFL.add_velocities(('u', 'v', 'w'))
    else:
        CFL.add_velocities(('u', 'w'))

    # Flow properties
    flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
    flow.add_property("Re_rms", name='Re')
    if verbose:
        flow.add_property("Pe_rms", name='Pe')
        flow.add_property("Nusselt_AB17", name='Nusselt')

    start_iter = solver.iteration
    start_sim_time = solver.sim_time

    try:
        start_time = time.time()
        start_iter = solver.iteration
        logger.info('starting main loop')
        good_solution = True
        first_step = True
        while solver.ok and good_solution:
            dt = CFL.compute_dt()
            # advance
            solver.step(dt)

            effective_iter = solver.iteration - start_iter
            Re_avg = flow.grid_average('Re')

            if threeD and effective_iter % Hermitian_cadence == 0:
                for field in solver.state.fields:
                    field.require_grid_space()

            # update lists
            if effective_iter % report_cadence == 0:
                log_string = 'Iteration: {:5d}, Time: {:8.3e} ({:8.3e}), dt: {:8.3e}, '.format(
                    solver.iteration - start_iter, solver.sim_time,
                    (solver.sim_time - start_sim_time) /
                    atmosphere.buoyancy_time, dt)
                if verbose:
                    log_string += '\n\t\tRe: {:8.5e}/{:8.5e}'.format(
                        Re_avg, flow.max('Re'))
                    log_string += '; Pe: {:8.5e}/{:8.5e}'.format(
                        flow.grid_average('Pe'), flow.max('Pe'))
                    log_string += '; Nu: {:8.5e}/{:8.5e}'.format(
                        flow.grid_average('Nusselt'), flow.max('Nusselt'))
                else:
                    log_string += 'Re: {:8.3e}/{:8.3e}'.format(
                        Re_avg, flow.max('Re'))
                logger.info(log_string)

            if not np.isfinite(Re_avg):
                good_solution = False
                logger.info(
                    "Terminating run.  Trapped on Reynolds = {}".format(
                        Re_avg))

            if first_step:
                if verbose:
                    import matplotlib
                    matplotlib.use('Agg')
                    import matplotlib.pyplot as plt
                    fig = plt.figure()
                    ax = fig.add_subplot(1, 1, 1)
                    ax.spy(solver.pencils[0].L,
                           markersize=1,
                           markeredgewidth=0.0)
                    fig.savefig(data_dir + "sparsity_pattern.png", dpi=1200)

                    import scipy.sparse.linalg as sla
                    LU = sla.splu(solver.pencils[0].LHS.tocsc(),
                                  permc_spec='NATURAL')
                    fig = plt.figure()
                    ax = fig.add_subplot(1, 2, 1)
                    ax.spy(LU.L.A, markersize=1, markeredgewidth=0.0)
                    ax = fig.add_subplot(1, 2, 2)
                    ax.spy(LU.U.A, markersize=1, markeredgewidth=0.0)
                    fig.savefig(data_dir + "sparsity_pattern_LU.png", dpi=1200)

                    logger.info("{} nonzero entries in LU".format(LU.nnz))
                    logger.info("{} nonzero entries in LHS".format(
                        solver.pencils[0].LHS.tocsc().nnz))
                    logger.info("{} fill in factor".format(
                        LU.nnz / solver.pencils[0].LHS.tocsc().nnz))
                first_step = False
                start_time = time.time()
    except:
        logger.error('Exception raised, triggering end of main loop.')
    finally:
        end_time = time.time()

        # Print statistics
        elapsed_time = end_time - start_time
        elapsed_sim_time = solver.sim_time
        N_iterations = solver.iteration - 1
        logger.info('main loop time: {:e}'.format(elapsed_time))
        logger.info('Iterations: {:d}'.format(N_iterations))
        logger.info('iter/sec: {:g}'.format(N_iterations / (elapsed_time)))
        if N_iterations > 0:
            logger.info('Average timestep: {:e}'.format(elapsed_sim_time /
                                                        N_iterations))

        if not no_join:
            logger.info('beginning join operation')
            try:
                final_checkpoint = Checkpoint(
                    data_dir, checkpoint_name='final_checkpoint')
                final_checkpoint.set_checkpoint(solver,
                                                wall_dt=1,
                                                mode="append")
                solver.step(dt)  #clean this up in the future...works for now.
                post.merge_process_files(data_dir + '/final_checkpoint/',
                                         cleanup=False)
            except:
                print('cannot save final checkpoint')

            logger.info(data_dir + '/checkpoint/')
            post.merge_process_files(data_dir + '/checkpoint/', cleanup=False)

            for task in analysis_tasks.keys():
                logger.info(analysis_tasks[task].base_path)
                post.merge_process_files(analysis_tasks[task].base_path,
                                         cleanup=False)

        if (atmosphere.domain.distributor.rank == 0):

            logger.info('main loop time: {:e}'.format(elapsed_time))
            if start_iter > 1:
                logger.info('Iterations (this run): {:d}'.format(N_iterations -
                                                                 start_iter))
                logger.info('Iterations (total): {:d}'.format(N_iterations -
                                                              start_iter))
            logger.info('iter/sec: {:g}'.format(N_iterations / (elapsed_time)))
            if N_iterations > 0:
                logger.info('Average timestep: {:e}'.format(elapsed_sim_time /
                                                            N_iterations))

            N_TOTAL_CPU = atmosphere.domain.distributor.comm_cart.size

            # Print statistics
            print('-' * 40)
            total_time = end_time - initial_time
            main_loop_time = end_time - start_time
            startup_time = start_time - initial_time
            n_steps = solver.iteration - 1
            print('  startup time:', startup_time)
            print('main loop time:', main_loop_time)
            print('    total time:', total_time)
            if n_steps > 0:
                print('    iterations:', n_steps)
                print(' loop sec/iter:', main_loop_time / n_steps)
                print('    average dt:', solver.sim_time / n_steps)
                print(
                    "          N_cores, Nx, Nz, startup     main loop,   main loop/iter, main loop/iter/grid, n_cores*main loop/iter/grid"
                )
                print(
                    'scaling:', ' {:d} {:d} {:d}'.format(N_TOTAL_CPU, nx, nz),
                    ' {:8.3g} {:8.3g} {:8.3g} {:8.3g} {:8.3g}'.format(
                        startup_time, main_loop_time, main_loop_time / n_steps,
                        main_loop_time / n_steps / (nx * nz),
                        N_TOTAL_CPU * main_loop_time / n_steps / (nx * nz)))
            print('-' * 40)
Example #9
0
def FC_convection(Rayleigh=1e6, Prandtl=1,
                  ChemicalPrandtl=1, ChemicalReynolds=10, stiffness=1e4,
                      n_rho_cz=3.5, n_rho_rz=1, 
                      nz_cz=128, nz_rz=128,
                      nx = None,
                      width=None,
                      single_chebyshev=False,
                      rk222=False,
                      superstep=False,
                      dense=False, nz_dense=64,
                      oz=False,
                      restart=None, data_dir='./', verbose=False):
    import numpy as np
    import time
    from stratified_dynamics import multitropes
    import os
    from dedalus.core.future import FutureField
    
    initial_time = time.time()

    logger.info("Starting Dedalus script {:s}".format(sys.argv[0]))

    if oz:
        constant_Prandtl=False
        stable_top=True
        mixed_temperature_flux=True
    else:
        constant_Prandtl=True
        stable_top=False
        mixed_temperature_flux=None
        
    # Set domain
    if nx is None:
        nx = nz_cz*4
        
    if single_chebyshev:
        nz = nz_cz
        nz_list = [nz_cz]
    else:
        if dense:
            nz = nz_rz+nz_dense+nz_cz
            #nz_list = [nz_rz, int(nz_dense/2), int(nz_dense/2), nz_cz]
            nz_list = [nz_rz, nz_dense, nz_cz]
        else:
            nz = nz_rz+nz_cz
            nz_list = [nz_rz, nz_cz]

    atmosphere = multitropes.FC_multitrope_rxn(nx=nx, nz=nz_list, stiffness=stiffness, 
                                         n_rho_cz=n_rho_cz, n_rho_rz=n_rho_rz, 
                                         verbose=verbose, width=width,
                                         constant_Prandtl=constant_Prandtl,
                                         stable_top=stable_top)
    
    atmosphere.set_IVP_problem(Rayleigh, Prandtl, ChemicalPrandtl, ChemicalReynolds)
        
    atmosphere.set_BC(mixed_temperature_flux=mixed_temperature_flux)
    problem = atmosphere.get_problem()

    #atmosphere.plot_atmosphere()
    #atmosphere.plot_scaled_atmosphere()

        
    if atmosphere.domain.distributor.rank == 0:
        if not os.path.exists('{:s}/'.format(data_dir)):
            os.mkdir('{:s}/'.format(data_dir))

    if rk222:
        logger.info("timestepping using RK222")
        ts = de.timesteppers.RK222
        cfl_safety_factor = 0.2*2
    else:
        logger.info("timestepping using RK443")
        ts = de.timesteppers.RK443
        cfl_safety_factor = 0.2*4

    # Build solver
    solver = problem.build_solver(ts)

    if do_checkpointing:
        checkpoint = Checkpoint(data_dir)
        checkpoint.set_checkpoint(solver, wall_dt=1800)

    # initial conditions
    if restart is None:
        atmosphere.set_IC(solver)
    else:
        if do_checkpointing:
            logger.info("restarting from {}".format(restart))
            checkpoint.restart(restart, solver)
        else:
            logger.error("No checkpointing capability in this branch of Dedalus.  Aborting.")
            raise

    logger.info("thermal_time = {:g}, top_thermal_time = {:g}".format(atmosphere.thermal_time, atmosphere.top_thermal_time))
    
    max_dt = atmosphere.min_BV_time 
    max_dt = atmosphere.buoyancy_time*0.25

    report_cadence = 1
    output_time_cadence = 0.1*atmosphere.buoyancy_time
    solver.stop_sim_time = np.inf
    solver.stop_iteration= np.inf
    solver.stop_wall_time = 23.5*3600

    logger.info("output cadence = {:g}".format(output_time_cadence))

    analysis_tasks = atmosphere.initialize_output(solver, data_dir, sim_dt=output_time_cadence)

    
    cfl_cadence = 1
    CFL = flow_tools.CFL(solver, initial_dt=max_dt, cadence=cfl_cadence, safety=cfl_safety_factor,
                         max_change=1.5, min_change=0.5, max_dt=max_dt, threshold=0.1)

    if superstep:
        CFL_traditional = flow_tools.CFL(solver, initial_dt=max_dt, cadence=cfl_cadence, safety=cfl_safety_factor,
                                         max_change=1.5, min_change=0.5, max_dt=max_dt, threshold=0.1)

        CFL_traditional.add_velocities(('u', 'w'))
    
        vel_u = FutureField.parse('u', CFL.solver.evaluator.vars, CFL.solver.domain)
        delta_x = atmosphere.Lx/nx
        CFL.add_frequency(vel_u/delta_x)
        vel_w = FutureField.parse('w', CFL.solver.evaluator.vars, CFL.solver.domain)
        mean_delta_z_cz = atmosphere.Lz_cz/nz_cz
        CFL.add_frequency(vel_w/mean_delta_z_cz)
    else:
        CFL.add_velocities(('u', 'w'))


    
    # Flow properties
    flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
    flow.add_property("Re_rms", name='Re')

    try:
        start_time = time.time()
        while solver.ok:

            dt = CFL.compute_dt()
            # advance
            solver.step(dt)

            # update lists
            if solver.iteration % report_cadence == 0:
                log_string = 'Iteration: {:5d}, Time: {:8.3e} ({:8.3e}), '.format(solver.iteration, solver.sim_time, solver.sim_time/atmosphere.buoyancy_time)
                log_string += 'dt: {:8.3e}'.format(dt)
                if superstep:
                    dt_traditional = CFL_traditional.compute_dt()
                    log_string += ' (vs {:8.3e})'.format(dt_traditional)
                log_string += ', '
                log_string += 'Re: {:8.3e}/{:8.3e}'.format(flow.grid_average('Re'), flow.max('Re'))
                logger.info(log_string)
    except:
        logger.error('Exception raised, triggering end of main loop.')
        raise
    finally:
        end_time = time.time()

        # Print statistics
        elapsed_time = end_time - start_time
        elapsed_sim_time = solver.sim_time
        N_iterations = solver.iteration 
        logger.info('main loop time: {:e}'.format(elapsed_time))
        logger.info('Iterations: {:d}'.format(N_iterations))
        logger.info('iter/sec: {:g}'.format(N_iterations/(elapsed_time)))
        logger.info('Average timestep: {:e}'.format(elapsed_sim_time / N_iterations))
        
        logger.info('beginning join operation')
        if do_checkpointing:
            logger.info(data_dir+'/checkpoint/')
            post.merge_analysis(data_dir+'/checkpoint/')

        for task in analysis_tasks:
            logger.info(analysis_tasks[task].base_path)
            post.merge_analysis(analysis_tasks[task].base_path)

        if (atmosphere.domain.distributor.rank==0):

            logger.info('main loop time: {:e}'.format(elapsed_time))
            logger.info('Iterations: {:d}'.format(N_iterations))
            logger.info('iter/sec: {:g}'.format(N_iterations/(elapsed_time)))
            logger.info('Average timestep: {:e}'.format(elapsed_sim_time / N_iterations))
 
            N_TOTAL_CPU = atmosphere.domain.distributor.comm_cart.size

            # Print statistics
            print('-' * 40)
            total_time = end_time-initial_time
            main_loop_time = end_time - start_time
            startup_time = start_time-initial_time
            n_steps = solver.iteration-1
            print('  startup time:', startup_time)
            print('main loop time:', main_loop_time)
            print('    total time:', total_time)
            print('Iterations:', solver.iteration)
            print('Average timestep:', solver.sim_time / n_steps)
            print("          N_cores, Nx, Nz, startup     main loop,   main loop/iter, main loop/iter/grid, n_cores*main loop/iter/grid")
            print('scaling:',
                  ' {:d} {:d} {:d}'.format(N_TOTAL_CPU,nx,nz),
                  ' {:8.3g} {:8.3g} {:8.3g} {:8.3g} {:8.3g}'.format(startup_time,
                                                                    main_loop_time, 
                                                                    main_loop_time/n_steps, 
                                                                    main_loop_time/n_steps/(nx*nz), 
                                                                    N_TOTAL_CPU*main_loop_time/n_steps/(nx*nz)))
            print('-' * 40)
def FC_polytrope(dynamics_file,
                 Rayleigh=1e4,
                 Prandtl=1,
                 aspect_ratio=4,
                 Taylor=None,
                 theta=0,
                 nz=128,
                 nx=None,
                 ny=None,
                 threeD=False,
                 mesh=None,
                 n_rho_cz=3,
                 epsilon=1e-4,
                 gamma=5 / 3,
                 run_time=23.5,
                 run_time_buoyancies=None,
                 run_time_iter=np.inf,
                 fixed_T=False,
                 fixed_flux=False,
                 mixed_flux_T=False,
                 const_mu=True,
                 const_kappa=True,
                 dynamic_diffusivities=False,
                 split_diffusivities=False,
                 chemistry=True,
                 ChemicalPrandtl=1,
                 Qu_0=5e-8,
                 phi_0=10,
                 restart=None,
                 start_new_files=False,
                 scalar_file=None,
                 rk222=False,
                 safety_factor=0.2,
                 max_writes=20,
                 data_dir='./',
                 out_cadence=0.1,
                 no_coeffs=False,
                 no_join=False,
                 verbose=False):

    import dedalus.public as de
    from dedalus.tools import post
    from dedalus.extras import flow_tools

    import time
    import os
    import sys
    from stratified_dynamics import polytropes
    from tools.checkpointing import Checkpoint

    checkpoint_min = 30

    initial_time = time.time()

    logger.info("Starting Dedalus script {:s}".format(sys.argv[0]))

    if nx is None:
        nx = int(np.round(nz * aspect_ratio))
    if threeD and ny is None:
        ny = nx

    eqn_dict = {
        'nx': nx,
        'nz': nz,
        'constant_kappa': const_kappa,
        'constant_mu': const_mu,
        'epsilon': epsilon,
        'gamma': gamma,
        'n_rho_cz': n_rho_cz,
        'aspect_ratio': aspect_ratio,
        'fig_dir': data_dir
    }
    if threeD:
        eqn_dict['mesh'] = mesh
        eqn_dict['nz'] = nz
        atmosphere = polytropes.FC_polytrope_rxn_3d(**eqn_dict)

    else:
        if dynamic_diffusivities:
            atmosphere = polytropes.FC_polytrope_2d_kappa(**eqn_dict)
        else:
            if chemistry:
                atmosphere = polytropes.FC_polytrope_rxn_2d(**eqn_dict)
            else:
                atmosphere = polytropes.FC_polytrope_2d(**eqn_dict)

    if epsilon < 1e-4:
        ncc_cutoff = 1e-14
    elif epsilon > 1e-1:
        ncc_cutoff = 1e-6
    else:
        ncc_cutoff = 1e-10

    problem_dict = {
        'ncc_cutoff': ncc_cutoff,
        'split_diffusivities': split_diffusivities
    }
    if threeD:
        problem_dict['Taylor'] = Taylor
        problem_dict['theta'] = theta
    if chemistry:
        problem_dict['ChemicalPrandtl'] = ChemicalPrandtl
        problem_dict['Qu_0'] = Qu_0
        problem_dict['phi_0'] = phi_0

    atmosphere.set_IVP_problem(Rayleigh, Prandtl, **problem_dict)

    if fixed_flux:
        atmosphere.set_BC(fixed_flux=True, stress_free=True)
    elif mixed_flux_T:
        atmosphere.set_BC(mixed_flux_temperature=True, stress_free=True)
    else:
        atmosphere.set_BC(fixed_temperature=True, stress_free=True)

    problem = atmosphere.get_problem()

    if atmosphere.domain.distributor.rank == 0:
        if not os.path.exists('{:s}/'.format(data_dir)):
            os.mkdir('{:s}/'.format(data_dir))

    if rk222:
        logger.info("timestepping using RK222")
        ts = de.timesteppers.RK222
        cfl_safety_factor = safety_factor * 2
    else:
        logger.info("timestepping using RK443")
        ts = de.timesteppers.RK443
        cfl_safety_factor = safety_factor * 4

    # Build solver
    solver = problem.build_solver(ts)

    #Check atmosphere
    logger.info("thermal_time = {:g}, top_thermal_time = {:g}".format(atmosphere.thermal_time,\
                                                                    atmosphere.top_thermal_time))
    logger.info("full atm HS check")
    atmosphere.check_atmosphere(make_plots=False,
                                rho=atmosphere.get_full_rho(solver),
                                T=atmosphere.get_full_T(solver))

    if restart is None:
        mode = "overwrite"
    else:
        mode = "append"

    logger.info('checkpointing in {}'.format(data_dir))
    checkpoint = Checkpoint(data_dir)

    if restart is None:
        atmosphere.set_IC(solver)
        dt = None
    else:
        logger.info("restarting from {}".format(restart))
        dt = checkpoint.restart(restart, solver)

    checkpoint.set_checkpoint(solver, wall_dt=checkpoint_min * 60, mode=mode)

    if run_time_buoyancies != None:
        solver.stop_sim_time = solver.sim_time + run_time_buoyancies * atmosphere.buoyancy_time
    else:
        solver.stop_sim_time = 100 * atmosphere.thermal_time

    solver.stop_iteration = solver.iteration + run_time_iter
    solver.stop_wall_time = run_time * 3600
    report_cadence = 1
    output_time_cadence = out_cadence * atmosphere.buoyancy_time
    Hermitian_cadence = 100

    logger.info("stopping after {:g} time units".format(solver.stop_sim_time))
    logger.info("output cadence = {:g}".format(output_time_cadence))
    analysis_tasks = atmosphere.initialize_output(
        solver,
        data_dir,
        sim_dt=output_time_cadence,
        coeffs_output=not (no_coeffs),
        mode=mode,
        max_writes=max_writes)

    # Reinjecting dynamics and tracers if desired
    logger.info("Re-injecting scalars")

    def reset_variable(key, h5_val, grid=False, grad=False):
        # Get variable
        k = solver.state[key]
        k.set_scales(1, keep_data=True)

        # Set grid or coefficient space
        if grid:
            gc = 'g'
            slice = solver.domain.dist.grid_layout.slices(k.meta[:]['scale'])
        else:
            gc = 'c'
            slice = solver.domain.dist.coeff_layout.slices(k.meta[:]['scale'])

        # Set initial value of variable
        k[gc] = h5_val[slice]

        if grad:  # Set gradient if called for
            k.differentiate('z', out=k)

        logger.info("Re-injecting: {}".format(key))
        return k

    objects = {}
    # Set all the dynamic variables

    if threeD:
        keys = ['u', 'u_z', 'v', 'v_z', 'w', 'w_z', 'T1', 'T1_z', 'ln_rho1']
        grads = [
            False, True, False, True, False, True, False, True, False, False,
            True, False, True
        ]
        h5_keys = ['u', 'u', 'v', 'v', 'w', 'w', 'T', 'T', 'ln_rho']
    else:
        keys = ['u', 'u_z', 'w', 'w_z', 'T1', 'T1_z', 'ln_rho1']
        grads = [
            False, True, False, True, False, True, False, False, True, False,
            True
        ]
        h5_keys = ['u', 'u', 'w', 'w', 'T', 'T', 'ln_rho']
    h5File_c = h5py.File(dynamics_file, 'r')
    dt = h5File_c['scales']['timestep'][-1]
    h5File_c = h5File_c['tasks']
    for i, K in enumerate(keys):
        # Restarting with final profile of run
        objects[K] = reset_variable(K, h5File_c[h5_keys[i]][-1], grad=grads[i])

    if scalar_file != 'None':
        keys = ['f', 'f_z', 'C', 'C_z', 'G', 'G_z']
        grads = [False, True, False, True]
        h5_keys = ['f', 'f', 'C', 'C', 'G', 'G']
        h5File_g = h5py.File(scalar_file, 'r')['tasks']
        for i, K in enumerate(keys):
            # Restarting with initial profile
            objects[K] = reset_variable(K,
                                        h5File_g[h5_keys[i]][0],
                                        grid=True,
                                        grad=grads[i])

    #Set up timestep defaults
    max_dt = output_time_cadence / 2
    if dt is None: dt = max_dt

    cfl_cadence = 1
    cfl_threshold = 0.1
    CFL = flow_tools.CFL(solver,
                         initial_dt=dt,
                         cadence=cfl_cadence,
                         safety=cfl_safety_factor,
                         max_change=1.5,
                         min_change=0.5,
                         max_dt=max_dt,
                         threshold=cfl_threshold)
    if threeD:
        CFL.add_velocities(('u', 'v', 'w'))
    else:
        CFL.add_velocities(('u', 'w'))

    # Flow properties
    flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
    flow.add_property("Re_rms", name='Re')
    if verbose:
        flow.add_property("Pe_rms", name='Pe')
        flow.add_property("Nusselt_AB17", name='Nusselt')

    start_iter = solver.iteration
    start_sim_time = solver.sim_time

    try:
        start_time = time.time()
        start_iter = solver.iteration
        logger.info('starting main loop')
        good_solution = True
        first_step = True
        while solver.ok and good_solution:
            dt = CFL.compute_dt()
            # advance
            solver.step(dt)

            effective_iter = solver.iteration - start_iter

            if threeD and effective_iter % Hermitian_cadence == 0:
                for field in solver.state.fields:
                    field.require_grid_space()

            # update lists
            if effective_iter % report_cadence == 0:
                Re_avg = flow.grid_average('Re')
                log_string = 'Iteration: {:5d}, Time: {:8.3e} ({:8.3e}), dt: {:8.3e}, '.format(
                    solver.iteration - start_iter, solver.sim_time,
                    (solver.sim_time - start_sim_time) /
                    atmosphere.buoyancy_time, dt)
                if verbose:
                    log_string += '\n\t\tRe: {:8.5e}/{:8.5e}'.format(
                        Re_avg, flow.max('Re'))
                    log_string += '; Pe: {:8.5e}/{:8.5e}'.format(
                        flow.grid_average('Pe'), flow.max('Pe'))
                    log_string += '; Nu: {:8.5e}/{:8.5e}'.format(
                        flow.grid_average('Nusselt'), flow.max('Nusselt'))
                else:
                    log_string += 'Re: {:8.3e}/{:8.3e}'.format(
                        Re_avg, flow.max('Re'))
                logger.info(log_string)

            if not np.isfinite(Re_avg):
                good_solution = False
                logger.info(
                    "Terminating run.  Trapped on Reynolds = {}".format(
                        Re_avg))

            if first_step:
                if verbose:
                    import matplotlib
                    matplotlib.use('Agg')
                    import matplotlib.pyplot as plt
                    fig = plt.figure()
                    ax = fig.add_subplot(1, 1, 1)
                    ax.spy(solver.pencils[0].L,
                           markersize=1,
                           markeredgewidth=0.0)
                    fig.savefig(data_dir + "sparsity_pattern.png", dpi=1200)

                    import scipy.sparse.linalg as sla
                    LU = sla.splu(solver.pencils[0].LHS.tocsc(),
                                  permc_spec='NATURAL')
                    fig = plt.figure()
                    ax = fig.add_subplot(1, 2, 1)
                    ax.spy(LU.L.A, markersize=1, markeredgewidth=0.0)
                    ax = fig.add_subplot(1, 2, 2)
                    ax.spy(LU.U.A, markersize=1, markeredgewidth=0.0)
                    fig.savefig(data_dir + "sparsity_pattern_LU.png", dpi=1200)

                    logger.info("{} nonzero entries in LU".format(LU.nnz))
                    logger.info("{} nonzero entries in LHS".format(
                        solver.pencils[0].LHS.tocsc().nnz))
                    logger.info("{} fill in factor".format(
                        LU.nnz / solver.pencils[0].LHS.tocsc().nnz))
                first_step = False
                start_time = time.time()
    except:
        logger.error('Exception raised, triggering end of main loop.')
        raise
    finally:
        end_time = time.time()

        # Print statistics
        elapsed_time = end_time - start_time
        elapsed_sim_time = solver.sim_time
        N_iterations = solver.iteration - 1
        logger.info('main loop time: {:e}'.format(elapsed_time))
        logger.info('Iterations: {:d}'.format(N_iterations))
        logger.info('iter/sec: {:g}'.format(N_iterations / (elapsed_time)))
        if N_iterations > 0:
            logger.info('Average timestep: {:e}'.format(elapsed_sim_time /
                                                        N_iterations))

        if not no_join:
            logger.info('beginning join operation')
            try:
                final_checkpoint = Checkpoint(
                    data_dir, checkpoint_name='final_checkpoint')
                final_checkpoint.set_checkpoint(solver,
                                                wall_dt=1,
                                                mode="append")
                solver.step(dt)  #clean this up in the future...works for now.
                post.merge_process_files(data_dir + '/final_checkpoint/',
                                         cleanup=True)
            except:
                print('cannot save final checkpoint')

            logger.info(data_dir + '/checkpoint/')
            post.merge_process_files(data_dir + '/checkpoint/', cleanup=True)

            for task in analysis_tasks.keys():
                logger.info(analysis_tasks[task].base_path)
                post.merge_process_files(analysis_tasks[task].base_path,
                                         cleanup=True)

        if (atmosphere.domain.distributor.rank == 0):

            logger.info('main loop time: {:e}'.format(elapsed_time))
            if start_iter > 1:
                logger.info('Iterations (this run): {:d}'.format(N_iterations -
                                                                 start_iter))
                logger.info('Iterations (total): {:d}'.format(N_iterations -
                                                              start_iter))
            logger.info('iter/sec: {:g}'.format(N_iterations / (elapsed_time)))
            if N_iterations > 0:
                logger.info('Average timestep: {:e}'.format(elapsed_sim_time /
                                                            N_iterations))

            N_TOTAL_CPU = atmosphere.domain.distributor.comm_cart.size

            # Print statistics
            print('-' * 40)
            total_time = end_time - initial_time
            main_loop_time = end_time - start_time
            startup_time = start_time - initial_time
            n_steps = solver.iteration - 1
            print('  startup time:', startup_time)
            print('main loop time:', main_loop_time)
            print('    total time:', total_time)
            if n_steps > 0:
                print('    iterations:', n_steps)
                print(' loop sec/iter:', main_loop_time / n_steps)
                print('    average dt:', solver.sim_time / n_steps)
                print(
                    "          N_cores, Nx, Nz, startup     main loop,   main loop/iter, main loop/iter/grid, n_cores*main loop/iter/grid"
                )
                print(
                    'scaling:', ' {:d} {:d} {:d}'.format(N_TOTAL_CPU, nx, nz),
                    ' {:8.3g} {:8.3g} {:8.3g} {:8.3g} {:8.3g}'.format(
                        startup_time, main_loop_time, main_loop_time / n_steps,
                        main_loop_time / n_steps / (nx * nz),
                        N_TOTAL_CPU * main_loop_time / n_steps / (nx * nz)))
            print('-' * 40)
 def RUN(self,
         scheme=de.timesteppers.RK443,
         adding=False,
         sim_time=2,
         wall_time=np.inf,
         tight=False,
         save=20):
     self.solver = self.problem.build_solver(scheme)
     if adding:
         t = load_last_value('times.txt')
         temp = load_last_array('temperatures.txt', shape=self.saving_shape)
         Variables = load_arrays('variables.txt',
                                 frequency=1,
                                 shape=self.saving_shape)
         T = self.solver.state['T']
         Psi = self.solver.state['psi']
         Curl = self.solver.state['curl']
         T['g'] = temp
         T.differentiate('z', out=self.solver.state['Tz'])
         Psi['g'] = Variables[0]
         Psi.differentiate('z', out=self.solver.state['psiz'])
         Curl['g'] = Variables[1]
         Curl.differentiate('z', out=self.solver.state['curlz'])
     else:
         t = 0
         print("Clearing old data ...")
         clear_file('temperatures.txt')
         clear_file('times.txt')
         clear_file('variables.txt')
         # Initial conditions ----------------------------------------------------------
         print("Initializing Values ...")
         eps = 1e-4
         k = 3.117
         x, z = self.problem.domain.grids(scales=1)
         T = self.solver.state['T']
         T['g'] = 1 - z + eps * np.sin(k * x) * np.sin(2 * np.pi * z)
         T.differentiate('z', out=self.solver.state['Tz'])
     # Stopping Parameters ---------------------------------------------------------
     self.solver.stop_sim_time = sim_time  # Length of simulation.
     self.solver.stop_wall_time = wall_time  # Real time allowed to compute.
     self.solver.stop_iteration = np.inf  # Maximum iterations allowed.
     # Control Flow ----------------------------------------------------------------
     dt = 1e-4
     if tight:
         cfl = flow_tools.CFL(self.solver,
                              initial_dt=dt,
                              cadence=1,
                              safety=1,
                              max_change=1.5,
                              min_change=0.01,
                              max_dt=0.01,
                              min_dt=1e-10)
     else:
         cfl = flow_tools.CFL(self.solver,
                              initial_dt=dt,
                              cadence=10,
                              safety=1,
                              max_change=1.5,
                              min_change=0.5,
                              max_dt=0.01,
                              min_dt=1e-6)
     cfl.add_velocities(('u', 'v'))
     # Flow properties (print during run; not recorded in the records files)
     flow = flow_tools.GlobalFlowProperty(self.solver, cadence=1)
     flow.add_property("sqrt(u **2 + v **2) / Ra", name='Re')
     # MAIN COMPUTATION LOOP -------------------------------------------------------
     try:
         print("####### BEGINNING CALCULATIONS #######")
         print("Starting main loop")
         start_time = time.time()
         while self.solver.ok:
             # Recompute time step and iterate.
             dt = self.solver.step(cfl.compute_dt())
             t = t + dt
             if self.solver.iteration % 10 == 0:
                 info = "Iteration {:>5d}, Time: {:.7f}, dt: {:.2e}".format(
                     self.solver.iteration, self.solver.sim_time, dt)
                 Re = flow.max("Re")
                 info += ", Max Re = {:f}".format(Re)
                 print(info)
                 if np.isnan(Re):
                     raise ValueError("Reynolds number went to infinity!!"
                                      "\nRe = {}".format(Re))
             if save:
                 if self.solver.iteration % save == 0:
                     T = self.solver.state['T']
                     append_unique_array('temperatures.txt', T['g'])
                     append_unique_value('times.txt', t)
     except BaseException as e:
         print("Exception raised, triggering end of main loop.")
         raise
     finally:
         print("####### CALCULATIONS FINISHED #######")
         total_time = time.time() - start_time
         print("Iterations: {:d}".format(self.solver.iteration))
         print("Sim end time: {:.3e}".format(self.solver.sim_time))
         print("Run time: {:.3e} sec".format(total_time))
         print("END OF SIMULATION\n")
         T = self.solver.state['T']
         Psi = self.solver.state['psi']
         Curl = self.solver.state['curl']
         append_unique_array('temperatures.txt', T['g'])
         append_unique_value('times.txt', t)
         append_unique_array('variables.txt', Psi['g'])
         append_unique_array('variables.txt', Curl['g'])
# Analysis
snap = model.solver.evaluator.add_file_handler('snapshots', sim_dt=0.2, 
                                               max_writes=10)
snap.add_task("interp(b, z=0)", scales=1, name='b midplane')
snap.add_task("interp(u, z=0)", scales=1, name='u midplane')
snap.add_task("interp(v, z=0)", scales=1, name='v midplane')
snap.add_task("interp(w, z=0)", scales=1, name='w midplane')

# CFL
CFL = flow_tools.CFL(model.solver, initial_dt=1e-4, cadence=5, 
                     safety=1.5, max_change=1.5, min_change=0.5, max_dt=0.05)
CFL.add_velocities(('u', 'v', 'w'))

# Flow properties
flow = flow_tools.GlobalFlowProperty(model.solver, cadence=10)
flow.add_property("sqrt(u*u + v*v + w*w) / ν", name='Re')

# Main loop
try:
    logger.info('Starting loop')
    start_run_time = time.time()
    while model.solver.ok:
        dt = CFL.compute_dt()
        model.solver.step(dt)
        if (model.solver.iteration-1) % 100 == 0:
            logger.info('Iteration: %i, Time: %e, dt: %e' %(
                        model.solver.iteration, model.solver.sim_time, dt))
            logger.info('Max Re = %f' %flow.max('Re'))
except:
    logger.error('Exception raised, triggering end of main loop.')
Example #13
0
def FC_convection(Rayleigh=1e6,
                  Prandtl=1,
                  stiffness=3,
                  m_rz=3,
                  gamma=5 / 3,
                  MHD=False,
                  MagneticPrandtl=1,
                  B0_amplitude=1,
                  n_rho_cz=1,
                  n_rho_rz=5,
                  nz_cz=128,
                  nz_rz=128,
                  nx=None,
                  width=None,
                  single_chebyshev=False,
                  rk222=False,
                  superstep=False,
                  dense=False,
                  nz_dense=64,
                  oz=False,
                  fixed_flux=False,
                  run_time=23.5,
                  run_time_buoyancies=np.inf,
                  run_time_iter=np.inf,
                  dynamic_diffusivities=False,
                  max_writes=20,
                  out_cadence=0.1,
                  no_coeffs=False,
                  no_join=False,
                  restart=None,
                  data_dir='./',
                  verbose=False,
                  label=None):
    def format_number(number, no_format_min=0.1, no_format_max=10):
        if number > no_format_max or number < no_format_min:
            try:
                mantissa = "{:e}".format(number).split("+")[0].split(
                    "e")[0].rstrip("0") or "0"
                power = "{:e}".format(number).split("+")[1].lstrip("0") or "0"
            except:
                mantissa = "{:e}".format(number).split("-")[0].split(
                    "e")[0].rstrip("0") or "0"
                power = "{:e}".format(number).split("-")[1].lstrip("0") or "0"
                power = "-" + power
            if mantissa[-1] == ".":
                mantissa = mantissa[:-1]
            mantissa += "e"
        else:
            mantissa = "{:f}".format(number).rstrip("0") or "0"
            if mantissa[-1] == ".":
                mantissa = mantissa[:-1]
            power = ""
        number_string = mantissa + power
        return number_string

    data_dir = './'
    # save data in directory named after script
    if data_dir[-1] != '/':
        data_dir += '/'
    data_dir += sys.argv[0].split('.py')[0]
    data_dir += "_nrhocz{}_Ra{}_S{}".format(format_number(n_rho_cz),
                                            format_number(Rayleigh),
                                            format_number(stiffness))
    if width:
        data_dir += "_erf{}".format(format_number(width))
    if args['--MHD']:
        data_dir += '_MHD'
    if label:
        data_dir += "_{}".format(label)
    data_dir += '/'

    from dedalus.tools.config import config

    config['logging']['filename'] = os.path.join(data_dir, 'logs/dedalus_log')
    config['logging']['file_level'] = 'DEBUG'

    import mpi4py.MPI
    if mpi4py.MPI.COMM_WORLD.rank == 0:
        if not os.path.exists('{:s}/'.format(data_dir)):
            os.makedirs('{:s}/'.format(data_dir))
        logdir = os.path.join(data_dir, 'logs')
        if not os.path.exists(logdir):
            os.mkdir(logdir)
    logger = logging.getLogger(__name__)
    logger.info("saving run in: {}".format(data_dir))

    import dedalus.public as de
    from dedalus.tools import post
    from dedalus.extras import flow_tools

    from dedalus.core.future import FutureField
    from stratified_dynamics import multitropes
    from tools.checkpointing import Checkpoint

    checkpoint_min = 30

    initial_time = time.time()

    logger.info("Starting Dedalus script {:s}".format(sys.argv[0]))

    constant_Prandtl = True
    stable_top = True
    mixed_temperature_flux = True

    # Set domain
    if nx is None:
        nx = nz_cz * 4

    if single_chebyshev:
        nz = nz_cz
        nz_list = [nz_cz]
    else:
        nz = nz_rz + nz_cz
        nz_list = [nz_rz, nz_cz]

    eqns_dict = {
        'stiffness': stiffness,
        'nx': nx,
        'nz': nz_list,
        'n_rho_cz': n_rho_cz,
        'n_rho_rz': n_rho_rz,
        'verbose': verbose,
        'width': width,
        'constant_Prandtl': constant_Prandtl,
        'stable_top': stable_top,
        'gamma': gamma,
        'm_rz': m_rz
    }
    if MHD:
        atmosphere = multitropes.FC_MHD_multitrope_guidefield_2d(**eqns_dict)

        atmosphere.set_IVP_problem(Rayleigh,
                                   Prandtl,
                                   MagneticPrandtl,
                                   guidefield_amplitude=B0_amplitude)
    else:
        atmosphere = multitropes.FC_multitrope(**eqns_dict)
        atmosphere.set_IVP_problem(Rayleigh, Prandtl)

    atmosphere.set_BC()
    problem = atmosphere.get_problem()

    if atmosphere.domain.distributor.rank == 0:
        if not os.path.exists('{:s}/'.format(data_dir)):
            os.mkdir('{:s}/'.format(data_dir))

    if rk222:
        logger.info("timestepping using RK222")
        ts = de.timesteppers.RK222
        cfl_safety_factor = 0.2 * 2
    else:
        logger.info("timestepping using RK443")
        ts = de.timesteppers.RK443
        cfl_safety_factor = 0.2 * 4

    # Build solver
    solver = problem.build_solver(ts)

    if restart is None:
        mode = "overwrite"
    else:
        mode = "append"

    checkpoint = Checkpoint(data_dir)

    # initial conditions
    if restart is None:
        atmosphere.set_IC(solver)
        dt = None
    else:
        logger.info("restarting from {}".format(restart))
        dt = checkpoint.restart(restart, solver)

    checkpoint.set_checkpoint(solver, wall_dt=checkpoint_min * 60, mode=mode)

    logger.info("thermal_time = {:g}, top_thermal_time = {:g}".format(
        atmosphere.thermal_time, atmosphere.top_thermal_time))

    max_dt = atmosphere.min_BV_time
    max_dt = atmosphere.buoyancy_time * out_cadence
    if dt is None: dt = max_dt

    report_cadence = 1
    output_time_cadence = out_cadence * atmosphere.buoyancy_time
    solver.stop_sim_time = solver.sim_time + run_time_buoyancies * atmosphere.buoyancy_time
    solver.stop_iteration = solver.iteration + run_time_iter
    solver.stop_wall_time = run_time * 3600

    logger.info("output cadence = {:g}".format(output_time_cadence))

    analysis_tasks = atmosphere.initialize_output(
        solver,
        data_dir,
        coeffs_output=not (no_coeffs),
        sim_dt=output_time_cadence,
        max_writes=max_writes,
        mode=mode)

    cfl_cadence = 1
    CFL = flow_tools.CFL(solver,
                         initial_dt=dt,
                         cadence=cfl_cadence,
                         safety=cfl_safety_factor,
                         max_change=1.5,
                         min_change=0.5,
                         max_dt=max_dt,
                         threshold=0.1)

    CFL.add_velocities(('u', 'w'))
    if MHD:
        CFL.add_velocities(
            ('Bx/sqrt(4*pi*rho_full)', 'Bz/sqrt(4*pi*rho_full)'))

    # Flow properties
    flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
    flow.add_property("Re_rms", name='Re')
    if MHD:
        flow.add_property("abs(dx(Bx) + dz(Bz))", name='divB')

    try:
        start_time = time.time()
        while solver.ok:

            dt = CFL.compute_dt()
            # advance
            solver.step(dt)

            # update lists
            if solver.iteration % report_cadence == 0:
                Re_avg = flow.grid_average('Re')
                if not np.isfinite(Re_avg):
                    solver.ok = False
                log_string = 'Iteration: {:5d}, Time: {:8.3e} ({:8.3e}), dt: {:8.3e}, '.format(
                    solver.iteration, solver.sim_time,
                    solver.sim_time / atmosphere.buoyancy_time, dt)
                log_string += 'Re: {:8.3e}/{:8.3e}'.format(
                    Re_avg, flow.max('Re'))
                if MHD:
                    log_string += ', divB: {:8.3e}/{:8.3e}'.format(
                        flow.grid_average('divB'), flow.max('divB'))

                logger.info(log_string)
    except:
        logger.error('Exception raised, triggering end of main loop.')
        raise
    finally:
        end_time = time.time()

        # Print statistics
        elapsed_time = end_time - start_time
        elapsed_sim_time = solver.sim_time
        N_iterations = solver.iteration
        logger.info('main loop time: {:e}'.format(elapsed_time))
        logger.info('Iterations: {:d}'.format(N_iterations))
        logger.info('iter/sec: {:g}'.format(N_iterations / (elapsed_time)))
        logger.info('Average timestep: {:e}'.format(elapsed_sim_time /
                                                    N_iterations))

        if not no_join:
            logger.info('beginning join operation')
            logger.info(data_dir + '/checkpoint/')
            post.merge_process_files(data_dir + '/checkpoint/', cleanup=False)

            for task in analysis_tasks:
                logger.info(analysis_tasks[task].base_path)
                post.merge_process_files(analysis_tasks[task].base_path,
                                         cleanup=False)

        if (atmosphere.domain.distributor.rank == 0):

            logger.info('main loop time: {:e}'.format(elapsed_time))
            logger.info('Iterations: {:d}'.format(N_iterations))
            logger.info('iter/sec: {:g}'.format(N_iterations / (elapsed_time)))
            logger.info('Average timestep: {:e}'.format(elapsed_sim_time /
                                                        N_iterations))

            N_TOTAL_CPU = atmosphere.domain.distributor.comm_cart.size

            # Print statistics
            print('-' * 40)
            total_time = end_time - initial_time
            main_loop_time = end_time - start_time
            startup_time = start_time - initial_time
            n_steps = solver.iteration - 1
            print('  startup time:', startup_time)
            print('main loop time:', main_loop_time)
            print('    total time:', total_time)
            print('Iterations:', solver.iteration)
            print('Average timestep:', solver.sim_time / n_steps)
            print(
                "          N_cores, Nx, Nz, startup     main loop,   main loop/iter, main loop/iter/grid, n_cores*main loop/iter/grid"
            )
            print(
                'scaling:', ' {:d} {:d} {:d}'.format(N_TOTAL_CPU, nx, nz),
                ' {:8.3g} {:8.3g} {:8.3g} {:8.3g} {:8.3g}'.format(
                    startup_time, main_loop_time, main_loop_time / n_steps,
                    main_loop_time / n_steps / (nx * nz),
                    N_TOTAL_CPU * main_loop_time / n_steps / (nx * nz)))
            print('-' * 40)
Example #14
0
snapshots = solver.evaluator.add_file_handler('snapshots',
                                              sim_dt=snapshotStep,
                                              max_writes=600,
                                              mode=fh_mode)
snapshots.add_system(solver.state)

output_cadence = 5

# CFL
#CFL = flow_tools.CFL(solver, initial_dt=dt, cadence=10, safety=10,
#                     max_change=1.5, min_change=0.1, max_dt=max_timestep, threshold=0.05)
#CFL.add_velocities(('vx', 'vy'))
#CFL.add_velocities(('vx2', 'vy2'))

# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=output_cadence)
flow.add_property("dw(n)*dw(n) + dw(T)*dw(T)", name='dw_enstrophy')

curr_time = time.time()

# Main loop
try:
    logger.info('Starting loop')
    start_time = time.time()
    while solver.proceed:
        #dt = CFL.compute_dt()
        dt = solver.step(dt)
        if (solver.iteration - 1) % output_cadence == 0:
            next_time = time.time()
            logger.info('Iteration: %i, Time: %e, dt: %e' %
                        (solver.iteration, solver.sim_time, dt))
Example #15
0
spinup_model.set_fields(u=noise,
                        v=noise,
                        w=noise,
                        b=spinup_model.z / Lz - noise)
spinup_model.stop_at(sim_time=spinup_time)

# Make CFL
spinup_CFL = flow_tools.CFL(spinup_model.solver,
                            initial_dt=dt0,
                            cadence=10,
                            max_change=1.5,
                            safety=0.5)
spinup_CFL.add_velocities(('u', 'v', 'w'))

# Flow properties: Reynolds number and three Nusselt numbers
spinup_stats = flow_tools.GlobalFlowProperty(spinup_model.solver,
                                             cadence=spinup_log_cadence)
add_reynolds_number(spinup_stats)
add_nusselt_numbers(spinup_stats)

# Spinup loop with coarse model
try:
    logger.info("Starting coarse spinup. Ra: {}, closure: {}".format(
        Ra, closure_name))
    start_run_time = time.time()
    log_time = time.time()
    while spinup_model.solver.ok:
        dt = spinup_CFL.compute_dt()
        spinup_model.solver.step(dt)
        if (spinup_model.solver.iteration - 1) % spinup_log_cadence == 0:
            compute_time = time.time() - log_time
            log_time = time.time()
Example #16
0
an1.add_task('p1+p2', name='diff_p1', layout='g')
an1.add_task('u1+u2', name='diff_u', layout='g')
an1.add_task('w1+w2', name='diff_w', layout='g')
an2 = solver.evaluator.add_file_handler('data_scalars',
                                        sim_dt=param.scalar_sim_dt,
                                        max_writes=100)
an2.add_task('integ(((u1+u2)**2 + (w1+w2)**2) /(a0+a1+a2)/2)',
             name='KE',
             layout='g')
an2.add_task('integ((u2*u2+w2*w2)/(a0+a1+a2)/2)', name='KE_pert', layout='g')
an2.add_task('-integ(u2*dx(t2xx) + u2*dz(t2xz) + w2*dx(t2xz) + w2*dz(t2zz))',
             name='D',
             layout='g')

# Monitoring
flow = flow_tools.GlobalFlowProperty(solver, cadence=param.CFL['cadence'])
flow.add_property('integ((u2*u2+w2*w2)/(a0+a1+a2)/2)', name='KE_pert')

# CFL calculator
CFL = flow_tools.CFL(solver, **param.CFL)
CFL.add_velocities(('u1+u2', 'w1+w2'))

# Main loop
dt_floor = 2**(np.log2(param.CFL['min_dt']) // 1)
try:
    logger.info('Starting loop')
    start_time = time.time()
    while solver.ok:
        dt = CFL.compute_dt()
        dt = dt_floor * (dt // dt_floor)
        dt = solver.step(dt, trim=param.trim)
Example #17
0
                                      ny=ny,
                                      nz=nz,
                                      ν=1,
                                      κ=1,
                                      closure=closure)

model.set_bc("no penetration", "top", "bottom")
model.set_bc("no slip", "top", "bottom")
model.set_bc("no flux", "top", "bottom")
model.build_solver()

# Initial condition
fields = {}
for field in ['u', 'v', 'w', 'b']:
    noise = dedaLES.random_noise(model.domain)
    pert = a * noise * model.z * (1 - model.z)
    fields[field] = pert

model.set_fields(**fields)
model.stop_at(iteration=3)

# Flow properties
flow = flow_tools.GlobalFlowProperty(model.solver)
flow.add_property("sqrt(u*u + v*v + w*w) / ν", name='Re_domain')

model.solver.step(dt)
logger.info('Max domain Re = %e' % flow.max("Re_domain"))

model.solver.step(dt)
logger.info('Max domain Re = %e' % flow.max("Re_domain"))
initial_b = initial_N2*model.z
model.set_fields(
    u = noise,
    v = noise,
    b = initial_b + np.sqrt(initial_N2) * noise
    #w = noise,
)

model.stop_at(sim_time=run_time)

dt_gizmo = dedaLES.TimeStepGizmo(model.solver, initial_dt=initial_dt, cadence=dt_cadence, max_change=1.5, safety=dt_safety)
dt_gizmo.add_velocities(('u', 'v', 'w'))
dt_gizmo.add_diffusivity('ν_sgs')
dt_gizmo.add_diffusivity('κb_sgs')

stats = flow_tools.GlobalFlowProperty(model.solver, cadence=stats_cadence)

stats.add_property("w*b", name="wb")
stats.add_property("ε + ε_sgs", name="epsilon")
stats.add_property("ν_sgs", name="eddy viscosity")
stats.add_property("κb_sgs", name="eddy diffusivity")
stats.add_property("χ + χ_sgs", name="chi")
stats.add_property("w*w", name="w_sq")
stats.add_property("sqrt(u*u + v*v + w*w) / ν", name='Re')

analysis = model.solver.evaluator.add_file_handler("mixed_layer_analysis_{}".format(identifier(model, closure=closure)), 
                                                   iter=analysis_cadence, max_writes=max_writes)
analysis.add_system(model.solver.state, layout='g')
analysis.add_task("interp(b, y=0)", scales=1, name='b midplane')
analysis.add_task("interp(u, y=0)", scales=1, name='u midplane')
analysis.add_task("interp(v, y=0)", scales=1, name='v midplane')
    def solve_IVP(self,
                  dt,
                  CFL,
                  data_dir,
                  analysis_tasks,
                  task_args=(),
                  pre_loop_args=(),
                  task_kwargs={},
                  pre_loop_kwargs={},
                  time_div=None,
                  track_fields=['Pe'],
                  threeD=False,
                  Hermitian_cadence=100,
                  no_join=False,
                  mode='append'):
        """Logic for a while-loop that solves an initial value problem.

        Parameters
        ----------
        dt                  : float
            The initial timestep of the simulation
        CFL                 : a Dedalus CFL object
            A CFL object that calculates the timestep of the simulation on the fly
        data_dir            : string
            The parent directory of output files
        analysis_tasks      : OrderedDict()
            An OrderedDict of dedalus FileHandler objects
        task_args, task_kwargs : list, dict, optional
            arguments & keyword arguments to the self._special_tasks() function
        pre_loop_args, pre_loop_kwargs: list, dict, optional
            arguments & keyword arguments to the self.pre_loop_setup() function
        time_div            : float, optional
            A siulation time to divide the normal time by for easier output tracking
        threeD              : bool, optional
            If True, occasionally force the solution to grid space to remove Hermitian errors
        Hermitian_cadence   : int, optional
            The number of timesteps between grid space forcings in 3D.
        no_join             : bool, optional
            If True, do not join files at the end of the simulation run.
        mode                : string, optional
            Dedalus output mode for final checkpoint. "append" or "overwrite"
        args, kwargs        : list and dictionary
            Additional arguments and keyword arguments to be passed to the self.special_tasks() function
        """

        # Flow properties
        self.flow = flow_tools.GlobalFlowProperty(self.solver, cadence=1)
        for f in track_fields:
            self.flow.add_property(f, name=f)

        self.pre_loop_setup(*pre_loop_args, **pre_loop_kwargs)

        start_time = time.time()
        # Main loop
        count = 0
        try:
            logger.info('Starting loop')
            init_time = self.solver.sim_time
            start_iter = self.solver.iteration
            while (self.solver.ok):
                dt = CFL.compute_dt()
                self.solver.step(dt)  #, trim=True)

                # prevents blow-up over long timescales in 3D due to hermitian-ness
                effective_iter = self.solver.iteration - start_iter
                if threeD and effective_iter % Hermitian_cadence == 0:
                    for field in self.solver.state.fields:
                        field.require_grid_space()

                self.special_tasks(*task_args, **task_kwargs)

                #reporting string
                self.iteration_report(dt, track_fields, time_div=time_div)

                if not np.isfinite(self.flow.grid_average(track_fields[0])):
                    break
        except:
            raise
            logger.error('Exception raised, triggering end of main loop.')
        finally:
            end_time = time.time()
            main_loop_time = end_time - start_time
            n_iter_loop = self.solver.iteration - 1
            logger.info('Iterations: {:d}'.format(n_iter_loop))
            logger.info('Sim end time: {:f}'.format(self.solver.sim_time))
            logger.info('Run time: {:f} sec'.format(main_loop_time))
            logger.info('Run time: {:f} cpu-hr'.format(
                main_loop_time / 60 / 60 *
                self.de_domain.domain.dist.comm_cart.size))
            logger.info('iter/sec: {:f} (main loop only)'.format(
                n_iter_loop / main_loop_time))
            try:
                final_checkpoint = Checkpoint(
                    data_dir, checkpoint_name='final_checkpoint')
                final_checkpoint.set_checkpoint(self.solver,
                                                wall_dt=1,
                                                mode=mode)
                self.solver.step(
                    dt)  #clean this up in the future...works for now.
                post.merge_process_files(data_dir + '/final_checkpoint/',
                                         cleanup=False)
            except:
                raise
                print('cannot save final checkpoint')
            finally:
                if not no_join:
                    logger.info('beginning join operation')
                    post.merge_analysis(data_dir + 'checkpoint')

                    for key, task in analysis_tasks.items():
                        logger.info(task.base_path)
                        post.merge_analysis(task.base_path)

                logger.info(40 * "=")
                logger.info('Iterations: {:d}'.format(n_iter_loop))
                logger.info('Sim end time: {:f}'.format(self.solver.sim_time))
                logger.info('Run time: {:f} sec'.format(main_loop_time))
                logger.info('Run time: {:f} cpu-hr'.format(
                    main_loop_time / 60 / 60 *
                    self.de_domain.domain.dist.comm_cart.size))
                logger.info('iter/sec: {:f} (main loop only)'.format(
                    n_iter_loop / main_loop_time))
Example #20
0
def Rayleigh_Benard(Rayleigh=1e6,
                    Prandtl=1,
                    nz=64,
                    nx=None,
                    ny=None,
                    aspect=4,
                    fixed_flux=False,
                    fixed_T=False,
                    mixed_flux_T=True,
                    stress_free=False,
                    no_slip=True,
                    restart=None,
                    run_time=23.5,
                    run_time_buoyancy=None,
                    run_time_iter=np.inf,
                    run_time_therm=1,
                    max_writes=20,
                    max_slice_writes=20,
                    output_dt=0.2,
                    data_dir='./',
                    coeff_output=True,
                    verbose=False,
                    no_join=False,
                    do_bvp=False,
                    num_bvps=10,
                    bvp_convergence_factor=1e-3,
                    bvp_equil_time=10,
                    bvp_resolution_factor=1,
                    bvp_transient_time=30,
                    bvp_final_equil_time=None,
                    min_bvp_time=50,
                    first_bvp_time=20,
                    first_bvp_convergence_factor=1e-2,
                    threeD=False,
                    seed=42,
                    mesh=None,
                    overwrite=False):
    import os
    from dedalus.tools.config import config

    config['logging']['filename'] = os.path.join(data_dir, 'logs/dedalus_log')
    config['logging']['file_level'] = 'DEBUG'

    import mpi4py.MPI
    if mpi4py.MPI.COMM_WORLD.rank == 0:
        if not os.path.exists('{:s}/'.format(data_dir)):
            os.makedirs('{:s}/'.format(data_dir))
        logdir = os.path.join(data_dir, 'logs')
        if not os.path.exists(logdir):
            os.mkdir(logdir)
    logger = logging.getLogger(__name__)
    logger.info("saving run in: {}".format(data_dir))

    import time
    from dedalus import public as de
    from dedalus.extras import flow_tools
    from dedalus.tools import post

    # input parameters
    logger.info("Ra = {}, Pr = {}".format(Rayleigh, Prandtl))

    # Parameters
    Lz = 1.
    Lx = aspect * Lz
    Ly = aspect * Lz
    if nx is None:
        nx = int(nz * aspect)
    if ny is None:
        ny = int(nz * aspect)

    if threeD:
        logger.info("resolution: [{}x{}x{}]".format(nx, ny, nz))
        equations = BoussinesqEquations3D(nx=nx,
                                          ny=ny,
                                          nz=nz,
                                          Lx=Lx,
                                          Ly=Ly,
                                          Lz=Lz,
                                          mesh=mesh)
    else:
        logger.info("resolution: [{}x{}]".format(nx, nz))
        equations = BoussinesqEquations2D(nx=nx, nz=nz, Lx=Lx, Lz=Lz)
    equations.set_IVP(Rayleigh, Prandtl)

    bc_dict = {
        'fixed_flux': None,
        'fixed_temperature': None,
        'mixed_flux_temperature': None,
        'mixed_temperature_flux': None,
        'stress_free': None,
        'no_slip': None
    }
    if mixed_flux_T:
        bc_dict['mixed_flux_temperature'] = True
    elif fixed_T:
        bc_dict['fixed_temperature'] = True
    elif fixed_flux:
        bc_dict['fixed_flux'] = True

    if stress_free:
        bc_dict['stress_free'] = True
    elif no_slip:
        bc_dict['no_slip'] = True

    supercrit = Rayleigh / RA_CRIT

    equations.set_BC(**bc_dict)

    # Build solver
    ts = de.timesteppers.RK443
    cfl_safety = 0.8

    solver = equations.problem.build_solver(ts)
    logger.info('Solver built')

    checkpoint = Checkpoint(data_dir)
    if isinstance(restart, type(None)):
        equations.set_IC(solver, seed=seed)
        dt = None
        mode = 'overwrite'
    else:
        logger.info("restarting from {}".format(restart))
        checkpoint.restart(restart, solver)
        if overwrite:
            mode = 'overwrite'
        else:
            mode = 'append'
    checkpoint.set_checkpoint(solver, wall_dt=checkpoint_min * 60, mode=mode)

    # Integration parameters
    #    if not isinstance(run_time_therm, type(None)):
    #        solver.stop_sim_time = run_time_therm*equations.thermal_time + solver.sim_time
    #    elif not isinstance(run_time_buoyancy, type(None)):
    #        solver.stop_sim_time  = run_time_buoyancy + solver.sim_time
    #    else:
    solver.stop_sim_time = np.inf
    solver.stop_wall_time = run_time * 3600.
    solver.stop_iteration = run_time_iter
    Hermitian_cadence = 100

    # Analysis
    max_dt = output_dt
    analysis_tasks = equations.initialize_output(solver,
                                                 data_dir,
                                                 coeff_output=coeff_output,
                                                 output_dt=output_dt,
                                                 mode=mode)

    # CFL
    CFL = flow_tools.CFL(solver,
                         initial_dt=0.1,
                         cadence=1,
                         safety=cfl_safety,
                         max_change=1.5,
                         min_change=0.5,
                         max_dt=max_dt,
                         threshold=0.1)
    if threeD:
        CFL.add_velocities(('u', 'v', 'w'))
    else:
        CFL.add_velocities(('u', 'w'))

    # Flow properties
    flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
    flow.add_property("Re", name='Re')
    #    flow.add_property("interp(w, z=0.95)", name='w near top')

    #    u, v, w = solver.state['u'], solver.state['v'], solver.state['w']

    if do_bvp:
        if not threeD:
            ny = 0
            atmo_class = BoussinesqEquations2D
        else:
            atmo_class = BoussinesqEquations3D
        bvp_solver = BoussinesqBVPSolver(atmo_class, nx, ny, nz, \
                                   flow, equations.domain.dist.comm_cart, \
                                   solver, num_bvps, bvp_equil_time, \
                                   threeD=threeD,
                                   bvp_transient_time=bvp_transient_time, \
                                   bvp_run_threshold=bvp_convergence_factor, \
                                   bvp_l2_check_time=1, mesh=mesh,\
                                   first_bvp_time=first_bvp_time,
                                   first_run_threshold=first_bvp_convergence_factor,\
                                   plot_dir='{}/bvp_plots/'.format(data_dir),\
                                   min_avg_dt=1e-10, final_equil_time=bvp_final_equil_time,
                                   min_bvp_time=min_bvp_time)
        bc_dict.pop('stress_free')
        bc_dict.pop('no_slip')

#   print(equations.domain.grid(0), equations.domain.grid(1), equations.domain.grid(2))

    first_step = True
    # Main loop
    try:
        logger.info('Starting loop')
        Re_avg = 0
        continue_bvps = True
        not_corrected_times = True
        init_time = solver.sim_time
        start_iter = solver.iteration
        while (solver.ok and np.isfinite(Re_avg)) and continue_bvps:
            dt = CFL.compute_dt()
            solver.step(dt)  #, trim=True)

            # Solve for blow-up over long timescales in 3D due to hermitian-ness
            effective_iter = solver.iteration - start_iter
            if threeD and effective_iter % Hermitian_cadence == 0:
                for field in solver.state.fields:
                    field.require_grid_space()

            Re_avg = flow.grid_average('Re')
            log_string = 'Iteration: {:5d}, '.format(solver.iteration)
            log_string += 'Time: {:8.3e} ({:8.3e} therm), dt: {:8.3e}, '.format(
                solver.sim_time, solver.sim_time / equations.thermal_time, dt)
            log_string += 'Re: {:8.3e}/{:8.3e}'.format(Re_avg, flow.max('Re'))
            logger.info(log_string)

            if not_corrected_times and Re_avg > 1:
                if not isinstance(run_time_therm, type(None)):
                    solver.stop_sim_time = run_time_therm * equations.thermal_time + solver.sim_time
                elif not isinstance(run_time_buoyancy, type(None)):
                    solver.stop_sim_time = run_time_buoyancy + solver.sim_time
                not_corrected_times = False

            if do_bvp:
                bvp_solver.update_avgs(dt, Re_avg, min_Re=np.sqrt(supercrit))
                if bvp_solver.check_if_solve():
                    atmo_kwargs = {'nz': nz * bvp_resolution_factor, 'Lz': Lz}
                    diff_args = [Rayleigh, Prandtl]
                    bvp_solver.solve_BVP(atmo_kwargs, diff_args, bc_dict)
                if bvp_solver.terminate_IVP():
                    continue_bvps = False

            if first_step:
                if verbose:
                    import matplotlib
                    matplotlib.use('Agg')
                    import matplotlib.pyplot as plt
                    fig = plt.figure()
                    ax = fig.add_subplot(1, 1, 1)
                    ax.spy(solver.pencils[0].L,
                           markersize=1,
                           markeredgewidth=0.0)
                    fig.savefig(data_dir + "sparsity_pattern.png", dpi=1200)

                    import scipy.sparse.linalg as sla
                    LU = sla.splu(solver.pencils[0].LHS.tocsc(),
                                  permc_spec='NATURAL')
                    fig = plt.figure()
                    ax = fig.add_subplot(1, 2, 1)
                    ax.spy(LU.L.A, markersize=1, markeredgewidth=0.0)
                    ax = fig.add_subplot(1, 2, 2)
                    ax.spy(LU.U.A, markersize=1, markeredgewidth=0.0)
                    fig.savefig(data_dir + "sparsity_pattern_LU.png", dpi=1200)

                    logger.info("{} nonzero entries in LU".format(LU.nnz))
                    logger.info("{} nonzero entries in LHS".format(
                        solver.pencils[0].LHS.tocsc().nnz))
                    logger.info("{} fill in factor".format(
                        LU.nnz / solver.pencils[0].LHS.tocsc().nnz))
                first_step = False
                start_time = time.time()
    except:
        raise
        logger.error('Exception raised, triggering end of main loop.')
    finally:
        end_time = time.time()
        main_loop_time = end_time - start_time
        n_iter_loop = solver.iteration - 1
        logger.info('Iterations: {:d}'.format(n_iter_loop))
        logger.info('Sim end time: {:f}'.format(solver.sim_time))
        logger.info('Run time: {:f} sec'.format(main_loop_time))
        logger.info('Run time: {:f} cpu-hr'.format(
            main_loop_time / 60 / 60 * equations.domain.dist.comm_cart.size))
        logger.info('iter/sec: {:f} (main loop only)'.format(n_iter_loop /
                                                             main_loop_time))
        try:
            final_checkpoint = Checkpoint(data_dir,
                                          checkpoint_name='final_checkpoint')
            final_checkpoint.set_checkpoint(solver, wall_dt=1, mode=mode)
            solver.step(dt)  #clean this up in the future...works for now.
            post.merge_process_files(data_dir + '/final_checkpoint/',
                                     cleanup=False)
        except:
            raise
            print('cannot save final checkpoint')
        finally:
            if not no_join:
                logger.info('beginning join operation')
                post.merge_analysis(data_dir + 'checkpoints')

                for task in analysis_tasks:
                    logger.info(task.base_path)
                    post.merge_analysis(task.base_path)

            logger.info(40 * "=")
            logger.info('Iterations: {:d}'.format(n_iter_loop))
            logger.info('Sim end time: {:f}'.format(solver.sim_time))
            logger.info('Run time: {:f} sec'.format(main_loop_time))
            logger.info('Run time: {:f} cpu-hr'.format(
                main_loop_time / 60 / 60 *
                equations.domain.dist.comm_cart.size))
            logger.info('iter/sec: {:f} (main loop only)'.format(
                n_iter_loop / main_loop_time))
    def simulate(self,
                 initial_conditions=None,
                 scheme=de.timesteppers.RK443,
                 sim_time=2,
                 wall_time=np.inf,
                 tight=False):
        """Load initial conditions, run the simulation, and merge results.

        Parameters:
            initial_conditions (None, str): determines from what source to
                draw the initial conditions. Valid options are as follows:
                - None: use trivial conditions (T_ = 1 - z, T = 1 - z + eps).
                - 'resume': use the most recent state file in the
                    records directory (load both model and DA system).
                - An .h5 filename: load state variables for the model and
                    reset the data assimilation state variables to zero.
            scheme (de.timesteppers): The kind of solver to use. Options are
                RK443 (de.timesteppers.RK443), RK111, RK222, RKSMR, etc.
            sim_time (float): The maximum amount of simulation time allowed
                (in seconds) before ending the simulation.
            wall_time (float): The maximum amound of computing time allowed
                (in seconds) before ending the simulation.
            tight (bool): If True, set a low cadence and min_dt for refined
                simulation. If False, set a higher cadence and min_dt for a
                more coarse (but faster) simulation.
        """
        if not self.problem:
            raise TypeError("problem not initialized (run setup())")

        self.logger.debug("\n")
        self.logger.debug("NEW SIMULATION")
        solver = self.problem.build_solver(scheme)
        self.logger.info("Solver built")

        N = int(self.problem.parameters['N'])

        # Initial conditions --------------------------------------------------
        if initial_conditions is None:  # "Trivial" conditions.
            dt = 1e-4
            eps = 1e-4
            k = 3.117

            x, z = self.problem.domain.grids(scales=1)
            T, T_ = solver.state['T'], solver.state['T_']
            # Start T from rest plus a small perturbation.
            T['g'] = 1 - z + eps * np.sin(k * x) * np.sin(2 * np.pi * z)
            T.differentiate('z', out=solver.state['Tz'])
            # Start T_ from rest.
            T_['g'] = 1 - z
            T_.differentiate('z', out=solver.state['Tz_'])
            self.logger.info("Using trivial initial conditions")

        elif isinstance(initial_conditions, str):  # Load data from a file.
            # Resume: load the state of the last (merged) state file.
            resume = initial_conditions == "resume"
            if resume:
                initial_conditions = self._get_merged_file("states")
            if not initial_conditions.endswith(".h5"):
                raise ValueError(
                    "'{}' is not an h5 file".format(initial_conditions))
            # Load the data from the specified h5 file into the system.
            self.logger.info("Loading initial conditions from {}".format(
                initial_conditions))

            with h5py.File(initial_conditions, 'r') as infile:
                dt = infile["scales/timestep"][-1] * .01  # initial dt
                errs = []
                tasks = ["T", "Tz", "psi", "psiz", "zeta", "zetaz"]
                if resume:  # Only load assimilating variables to resume.
                    tasks += ["T_", "Tz_", "psi_", "psiz_", "zeta_", "zetaz_"]
                    solver.sim_time = infile["scales/sim_time"][-1]
                    niters = infile["scales/iteration"][-1]
                    solver.initial_iteration = niters
                    solver.iteration = niters
                for name in tasks:
                    # Get task data from the h5 file (recording failures).
                    try:
                        data = infile["tasks/" + name][-1, :, :]
                    except KeyError as e:
                        errs.append("tasks/" + name)
                        continue
                    # Determine the chunk belonging to this process.
                    chunk = data.shape[1] // SIZE
                    subset = data[:, RANK * chunk:(RANK + 1) * chunk]
                    # Change the corresponding state variable.
                    scale = solver.state[name]['g'].shape[0] / \
                                        self.problem.parameters["xsize"]
                    solver.state[name].set_scales(1)
                    solver.state[name]['g'] = subset
                    solver.state[name].set_scales(scale)
                if errs:
                    raise KeyError("Missing keys in '{}': '{}'".format(
                        initial_conditions, "', '".join(errs)))
            # Initial conditions for assimilating system: T_0 = P_4(T0).
            if not resume:
                G = self.problem.domain.new_field()
                G['c'] = solver.state['T']['c'].copy()
                solver.state['T_']['g'] = BoussinesqDataAssimilation2D.P_N(
                    G, 4, True)
                solver.state['T_'].differentiate('z', out=solver.state['Tz_'])

        # Driving / projection function arguments -----------------------------

        dT = solver.state['T_'] - solver.state['T']
        self.problem.parameters["driving"].args = [dT, N]
        self.problem.parameters["driving"].original_args = [dT, N]

        # Stopping Parameters -------------------------------------------------

        solver.stop_sim_time = sim_time  # Length of simulation.
        solver.stop_wall_time = wall_time  # Real time allowed to compute.
        solver.stop_iteration = np.inf  # Maximum iterations allowed.

        # State snapshots -----------------------------------------------------

        # Save the entire state in states/ files. USE iter, NOT sim_dt.
        # NOTE: This is where BoussinesqDataAssimilation2Dmovie differs.
        snaps = solver.evaluator.add_file_handler(os.path.join(
            self.records_dir, "states"),
                                                  iter=1,
                                                  max_writes=5000,
                                                  mode="append")
        snaps.add_task("T")
        snaps.add_task("T_")
        snaps.add_task("driving", name="P_N")

        # Control Flow --------------------------------------------------------
        if tight:  # Tighter control flow (slower but safer).
            cfl = flow_tools.CFL(solver,
                                 initial_dt=dt,
                                 cadence=1,
                                 safety=1,
                                 max_change=1.5,
                                 min_change=0.01,
                                 max_dt=0.01,
                                 min_dt=1e-10)
        else:  # Looser control flow (faster but risky).
            cfl = flow_tools.CFL(solver,
                                 initial_dt=dt,
                                 cadence=10,
                                 safety=1,
                                 max_change=1.5,
                                 min_change=0.5,
                                 max_dt=0.01,
                                 min_dt=1e-6)
        cfl.add_velocities(('v', 'w'))
        cfl.add_velocities(('v_', 'w_'))

        # Flow properties (print during run; not recorded in the records files)
        flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
        flow.add_property("sqrt(v **2 + w **2) / Ra", name='Re')
        flow.add_property("sqrt(v_**2 + w_**2) / Ra", name='Re_')

        # MAIN COMPUTATION LOOP -----------------------------------------------
        try:
            self.logger.info("Starting main loop")
            start_time = time.time()
            while solver.ok:
                # Recompute time step and iterate.
                dt = cfl.compute_dt()
                dt = solver.step(dt)

                # Print info to the screen every 10 iterations.
                if solver.iteration % 10 == 0:
                    info = "Iteration {:>5d}, Time: {:.7f}, dt: {:.2e}".format(
                        solver.iteration, solver.sim_time, dt)
                    Re = flow.max("Re")
                    Re_ = flow.max("Re_")
                    info += ", Max Re = {:f}".format(Re)
                    info += ", Max Re_ = {:f}".format(Re_)
                    self.logger.info(info)
                    # Make sure the simulation hasn't blown up.
                    if np.isnan(Re) or np.isnan(Re_):
                        raise ValueError("Reynolds number went to infinity!!"
                                         "\nRe = {}, Re_ = {}".format(Re, Re_))
        except BaseException as e:
            self.logger.error("Exception raised, triggering end of main loop.")
            raise
        finally:
            total_time = time.time() - start_time
            cpu_hr = total_time / 60 / 60 * SIZE
            self.logger.info("Iterations: {:d}".format(solver.iteration))
            self.logger.info("Sim end time: {:.3e}".format(solver.sim_time))
            self.logger.info("Run time: {:.3e} sec".format(total_time))
            self.logger.info("Run time: {:.3e} cpu-hr".format(cpu_hr))
            self.logger.debug("END OF SIMULATION\n")
Example #22
0
                                          max_writes=200)
integ.add_task("Avg_x(vx)", name='<vx>_x', scales=1)
integ.add_task("Avg_x(vy)", name='<vy>_x', scales=1)
integ.add_task("Avg_x(Bx)", name='<Bx>_x', scales=1)
integ.add_task("Avg_x(By)", name='<By>_x', scales=1)
analysis_tasks.append(integ)

timeseries = solver.evaluator.add_file_handler(os.path.join(
    data_dir, 'timeseries'),
                                               sim_dt=1e-2)
timeseries.add_task("0.5*integ(vx**2 + vy**2)", name='Ekin')
timeseries.add_task("0.5*integ(Bx**2 + By**2)", name='Emag')
analysis_tasks.append(timeseries)

# Flow properties
flow = flow_tools.GlobalFlowProperty(solver, cadence=1)
flow.add_property("0.5*(vx**2 + vy**2)", name='Ekin')

try:
    logger.info('Starting loop')
    start_run_time = time.time()

    while solver.ok:
        solver.step(dt)
        if (solver.iteration - 1) % 1 == 0:
            logger.info('Iteration: %i, Time: %e, dt: %e' %
                        (solver.iteration, solver.sim_time, dt))
            logger.info('Max E_kin = %17.12e' % flow.max('Ekin'))
except:
    logger.error('Exception raised, triggering end of main loop.')
    raise
def Dedalus_Inelastic_S(Ra,Pr,Np,Ta,end_time,snaps=True):
    ''' Builds a problem with its equations, parameters and boundary conditions and solves it, saving values in a h5py '''
    m=1.5
    theta = 1-np.exp(-Np/m)
    Ly, Lz = 2,1
    Ny, Nz = 192,96
    Lat=np.pi/4
    initial_timestep = 1.5e-4           # Initial timestep
    snapshot_skip=10
    analysis_freq = 1.5e-3              # Frequency analysis files are outputted
    end_sim_time = end_time                  # Stop time in simulations units
    end_wall_time = np.inf              # Stop time in wall time
    end_iterations = np.inf             # Stop time in iterations
    max_dt=0.005
    save_direc = "raw_data/"
    pathlib.Path(save_direc).mkdir(parents=True, exist_ok=True)
    # Create bases and domain
    y_basis = de.Fourier('y', Ny, interval=(0, Ly), dealias=3/2)   # Fourier basis in the x
    z_basis = de.Chebyshev('z', Nz, interval=(0, Lz), dealias=3/2) # Chebyshev basis in the z
    domain = de.Domain([y_basis, z_basis], grid_dtype=np.float64)  # Defining our domain
    z = domain.grid(1, scales=1)                                   # accessing the z values
    # 2D Anelastic hydrodynamics
    problem = de.IVP(domain, variables=['p', 's', 'u', 'v', 'w', 'sz', 'uz', 'vz', 'wz'])
    problem.meta['p','s','u','w']['z']['dirichlet'] = True
    # Defining model parameters
    problem.parameters['Ly'] = Ly
    problem.parameters['Lz'] = Lz
    problem.parameters['Ra'] = Ra
    problem.parameters['Pr'] = Pr
    problem.parameters['Ta'] = Ta
    problem.parameters['Lat'] = Lat
    problem.parameters['m'] = m
    problem.parameters['theta'] = theta
    problem.parameters['X'] = Ra/Pr
    problem.parameters['Y'] = (Pr*Pr*theta) / Ra
    problem.parameters['T'] = Ta**(1/2)
    Sfilename='S-Ra:'+str(Ra)+'-Pr:'+str(Pr)+'-Ta:'+str(Ta)+'-Np:'+str(Np)
    KEfilename='KE-Ra:'+str(Ra)+'-Pr:'+str(Pr)+'-Ta:'+str(Ta)+'-Np:'+str(Np)
    # Non-constant coeffiecents
    rho_ref = domain.new_field(name='rho_ref')
    rho_ref['g'] = (1-theta*z)**m
    rho_ref.meta['y']['constant'] = True
    problem.parameters['rho_ref'] = rho_ref         # Background state for rho
    T_ref = domain.new_field(name='T_ref')
    T_ref['g'] = 1-theta*z
    T_ref.meta['y']['constant'] = True
    problem.parameters['T_ref'] = T_ref             # Background state for T
    dz_rho_ref = domain.new_field(name='dz_rho_ref')
    dz_rho_ref['g'] = -theta*m*((1-theta*z)**(m-1))
    dz_rho_ref.meta['y']['constant'] = True
    problem.parameters['dz_rho_ref'] = dz_rho_ref   # z-derivative of rho_ref
    # Defining d/dz of s, u, and w for reducing our equations to first order
    problem.add_equation("sz - dz(s) = 0")
    problem.add_equation("uz - dz(u) = 0")
    problem.add_equation("vz - dz(v) = 0")
    problem.add_equation("wz - dz(w) = 0")
    # mass continuity with rho_ref and dz(rho_ref) expanded analytically
    problem.add_equation("  (1-theta*z)*(dy(v) + wz) - theta*m*w = 0 ")
    # x-component of the momentum equation
    problem.add_equation("  rho_ref*( dt(u) - dy(dy(u)) - dz(uz) + T*(w*cos(Lat) - v*sin(Lat)) ) - dz_rho_ref*uz = -rho_ref*( v*dy(u) + w*uz ) ")
    # y-component of the momentum equation
    problem.add_equation("  rho_ref*( dt(v) - (4/3)*dy(dy(v)) - dz(vz) - (1/3)*dy(wz) + T*u*sin(Lat) ) + dy(p) - dz_rho_ref*(vz + dy(w)) = -rho_ref*( v*dy(v) + w*vz )")
    # z-component of the momentum equation
    problem.add_equation("  rho_ref*T_ref*( dt(w) - X*s - dy(dy(w)) - (4/3)*dz(wz) - (1/3)*dy(vz) - T*u*cos(Lat) ) + T_ref*dz(p) + theta*m*p + (2/3)*theta*m*rho_ref*( 2*wz - dy(v) ) = -rho_ref*T_ref*( v*dy(w) + w*wz )")
    # entropy diffusion equation
    problem.add_equation("  T_ref*( Pr*dt(s) - dy(dy(s)) - dz(sz) ) + theta*(m+1)*sz = -Pr*T_ref*( v*dy(s) + w*sz ) + 2*Y*( dy(v)*dy(v) + wz*wz + vz*dy(w) - (1/3)*(dy(v) + wz)*(dy(v) + wz) + (1/2)*(dy(u)*dy(u) + uz*uz + vz*vz + dy(w)*dy(w)) )")
    # Flux equations for use in analysis outputs
    problem.add_bc("left(w) = 0")            # Impermeable bottom boundary
    problem.add_bc("right(w) = 0", condition="(ny != 0)")   # Impermeable top boundary
    problem.add_bc("right(p) = 0", condition="(ny == 0)")   # Required for equations to be well-posed - see https://bit.ly/2nPVWIg for a related discussion
    problem.add_bc("left(uz) = 0")           # Stress-free bottom boundary
    problem.add_bc("right(uz) = 0")          # Stress-free top boundary
    problem.add_bc("left(vz) = 0")
    problem.add_bc("right(vz) = 0")
    problem.add_bc("right(s) = 0")           # Fixed entropy at upper boundary, arbitarily set to 0
    problem.add_bc("left(sz) = -1")          # Fixed flux at bottom boundary, F = F_cond
    # Build solver
    solver = problem.build_solver(de.timesteppers.RK222)
    logger.info('Solver built')
    # Initial conditions
    x = domain.grid(0)
    z = domain.grid(1)
    s = solver.state['s']
    w = solver.state['w']
    sz = solver.state['sz']
    # Random perturbations, initialized globally for same results in parallel
    gshape = domain.dist.grid_layout.global_shape(scales=1)
    slices = domain.dist.grid_layout.slices(scales=1)
    rand = np.random.RandomState(seed=42)
    noise = rand.standard_normal(gshape)[slices]
    # Linear background + perturbations damped at walls
    zb, zt = z_basis.interval
    pert =  1e-5*noise*(zt - z)*(z - zb)
    s['g'] = pert
    s.differentiate('z', out=sz)
    dt = initial_timestep # Initial timestep
    # Integration parameters --- Note if these are all set to np.inf, simulation will perpetually run.
    solver.stop_sim_time = end_sim_time
    solver.stop_wall_time = end_wall_time
    solver.stop_iteration = end_iterations
    # CFL criterion
    CFL = flow_tools.CFL(solver,initial_dt=dt,cadence=1,safety=1,max_change=1.5,min_change=0.01,max_dt=0.01,min_dt=1e-10)
    CFL.add_velocities(('v', 'w'))
    # Flow properties
    flow = flow_tools.GlobalFlowProperty(solver,cadence=10)
    flow.add_property("sqrt(u*u + v*v + w*w)",name='Re')
    # Analysis tasks
    analysis = solver.evaluator.add_file_handler(save_direc +'analysis',sim_dt=analysis_freq,max_writes=5000)
    analysis.add_task("s", layout='g', name='entropy')
    # Flux decomposition - Internal energy equation
    analysis.add_task("integ( integ( sqrt(u*u + v*v + w*w) , 'y')/Ly, 'z')/Lz", layout='g', name='Re') # Mean Reynolds number
    analysis.add_task(" integ( (integ(0.5*(u*u + v*v + w*w)*rho_ref,'y')/Ly), 'z')/Lz", layout='g', name='KE') # Mean KE
    # Creating a parameter file
    run_parameters = solver.evaluator.add_file_handler(save_direc+'run_parameters',wall_dt=1e20,max_writes=1)
    run_parameters.add_task(Ly,name="Ly")
    run_parameters.add_task(Lz,name="Lz")
    run_parameters.add_task(Ra,name="Ra")
    run_parameters.add_task(Pr,name="Pr")
    run_parameters.add_task(Np,name="Np")
    run_parameters.add_task(m,name="m")
    run_parameters.add_task(Ny,name="Ny")
    run_parameters.add_task(Nz,name="Nz")
    run_parameters.add_task("z",layout='g',name="z_grid")
    run_parameters.add_task(analysis_freq,name="ana_freq")
    run_parameters.add_task(max_dt,name="max_dt")
    try: # Main loop
        logger.info('Starting loop')
        start_time = time.time()
        while solver.ok:
            dt = CFL.compute_dt()
            dt = solver.step(dt)
            time.sleep(0.02)
            if (solver.iteration) == 1:
                # Prints various parameters to terminal upon starting the simulation
                logger.info('Parameter values imported form run_param_file.py:')
                logger.info('Ly = {}, Lz = {}; (Resolution of {},{})'.format(Ly, Lz, Ny, Nz))
                logger.info('Ra = {}, Pr = {}, Np = {}'.format(Ra, Pr, Np))
                logger.info('Analysis files outputted every {}'.format(analysis_freq))
                if end_sim_time != np.inf:
                    logger.info('Simulation finishes at sim_time = {}'.format(end_sim_time))
                elif end_wall_time != np.inf:
                    logger.info('Simulation finishes at wall_time = {}'.format(end_wall_time))
                elif end_iterations != np.inf:
                    logger.info('Simulation finishes at iteration {}'.format(end_iterations))
                else:
                    logger.info('No clear end point defined. Simulation may run perpetually.')
            if (solver.iteration-1) % 10 == 0:
                # Prints progress information include maximum Reynolds number every 10 iterations
                logger.info('Iteration: %i, Time: %e, dt: %e' %(solver.iteration, solver.sim_time, dt))
                logger.info('Max Re = %f' %flow.max('Re'))
    except:
        logger.error('Exception raised, triggering end of main loop.')
        raise
    finally:
        # Prints concluding information upon reaching the end of the simulation.
        end_time = time.time()
        if snaps==True:
            with open(Sfilename,'w') as file:
                    file.write(str(gshape[0])+' '+str(gshape[1]))
        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)/60/60*domain.dist.comm_cart.size))
        with h5py.File("raw_data/analysis/analysis_s1/analysis_s1_p0.h5", mode='r') as file:
            times=file['scales']['sim_time'][:]
            data=file['tasks']['entropy'][:,:,:]
            ke=file['tasks']['KE'][:]
            ke=np.squeeze(ke)
            times=np.squeeze(times)
            n=times.shape[0]
            if snaps==True:
                for i in range(n):
                    if i%snapshot_skip==0:
                        append_unique_value(Sfilename,times[i])
                        append_unique_array(Sfilename,data[i])
            save_fct_txt(times,ke,KEfilename)