import matplotlib.pyplot as plt
import numpy as np

import pyro.lm_atm.LM_atm_interface as lm_interface
import pyro.mesh.array_indexer as ai
import pyro.mesh.boundary as bnd
import pyro.multigrid.variable_coeff_MG as vcMG
from pyro.mesh import patch, reconstruction
from pyro.simulation_null import NullSimulation, bc_setup, grid_setup



[docs]
class Basestate:
    def __init__(self, ny, *, ng=0):
        self.ny = ny
        self.ng = ng
        self.qy = ny + 2*ng

        self.d = np.zeros((self.qy), dtype=np.float64)

        self.jlo = ng
        self.jhi = ng+ny-1


[docs]
    def v(self, buf=0):
        return self.d[self.jlo-buf:self.jhi+1+buf]



[docs]
    def v2d(self, buf=0):
        return self.d[np.newaxis, self.jlo-buf:self.jhi+1+buf]



[docs]
    def v2dp(self, shift, buf=0):
        return self.d[np.newaxis, self.jlo+shift-buf:self.jhi+1+shift+buf]



[docs]
    def jp(self, shift, buf=0):
        return self.d[self.jlo-buf+shift:self.jhi+1+buf+shift]





[docs]
class Simulation(NullSimulation):

    def __init__(self, solver_name, problem_name, problem_func, rp, *,
                 problem_finalize_func=None, problem_source_func=None,
                 timers=None):

        super().__init__(solver_name, problem_name, problem_func, rp,
                         problem_finalize_func=problem_finalize_func,
                         problem_source_func=problem_source_func,
                         timers=timers)

        self.base = {}
        self.aux_data = None
        self.in_preevolve = False


[docs]
    def initialize(self):
        """
        Initialize the grid and variables for low Mach atmospheric flow
        and set the initial conditions for the chosen problem.
        """

        myg = grid_setup(self.rp, ng=4)

        bc_dens, bc_xodd, bc_yodd = bc_setup(self.rp)

        my_data = patch.CellCenterData2d(myg)

        my_data.register_var("density", bc_dens)
        my_data.register_var("x-velocity", bc_xodd)
        my_data.register_var("y-velocity", bc_yodd)

        # we'll keep the internal energy around just as a diagnostic
        my_data.register_var("eint", bc_dens)

        # phi -- used for the projections.  The boundary conditions
        # here depend on velocity.  At a wall or inflow, we already
        # have the velocity we want on the boundary, so we want
        # Neumann (dphi/dn = 0).  For outflow, we want Dirichlet (phi
        # = 0) -- this ensures that we do not introduce any tangental
        # acceleration.
        bcs = []
        for bc in [self.rp.get_param("mesh.xlboundary"),
                   self.rp.get_param("mesh.xrboundary"),
                   self.rp.get_param("mesh.ylboundary"),
                   self.rp.get_param("mesh.yrboundary")]:
            bctype = None
            if bc == "periodic":
                bctype = "periodic"
            elif bc in ["reflect", "slipwall"]:
                bctype = "neumann"
            elif bc in ["outflow"]:
                bctype = "dirichlet"
            bcs.append(bctype)

        bc_phi = bnd.BC(xlb=bcs[0], xrb=bcs[1], ylb=bcs[2], yrb=bcs[3])

        my_data.register_var("phi-MAC", bc_phi)
        my_data.register_var("phi", bc_phi)

        # gradp -- used in the projection and interface states.  We'll do the
        # same BCs as density
        my_data.register_var("gradp_x", bc_dens)
        my_data.register_var("gradp_y", bc_dens)

        my_data.create()

        self.cc_data = my_data

        # some auxiliary data that we'll need to fill GC in, but isn't
        # really part of the main solution
        aux_data = patch.CellCenterData2d(myg)

        aux_data.register_var("coeff", bc_dens)
        aux_data.register_var("source_y", bc_yodd)

        aux_data.create()
        self.aux_data = aux_data

        # we also need storage for the 1-d base state -- we'll store this
        # in the main class directly.
        self.base["rho0"] = Basestate(myg.ny, ng=myg.ng)
        self.base["p0"] = Basestate(myg.ny, ng=myg.ng)

        # now set the initial conditions for the problem
        self.problem_func(self.cc_data, self.base, self.rp)

        # Construct beta_0
        gamma = self.rp.get_param("eos.gamma")
        self.base["beta0"] = Basestate(myg.ny, ng=myg.ng)
        self.base["beta0"].d[:] = self.base["p0"].d**(1.0/gamma)

        # we'll also need beta_0 on vertical edges -- on the domain edges,
        # just do piecewise constant
        self.base["beta0-edges"] = Basestate(myg.ny, ng=myg.ng)
        self.base["beta0-edges"].jp(1)[:] = \
            0.5*(self.base["beta0"].v() + self.base["beta0"].jp(1))
        self.base["beta0-edges"].d[myg.jlo] = self.base["beta0"].d[myg.jlo]
        self.base["beta0-edges"].d[myg.jhi+1] = self.base["beta0"].d[myg.jhi]



[docs]
    def make_prime(self, a, a0):
        return a - a0.v2d(buf=a0.ng)



[docs]
    def method_compute_timestep(self):
        """
        The timestep() function computes the advective timestep
        (CFL) constraint.  The CFL constraint says that information
        cannot propagate further than one zone per timestep.

        We use the driver.cfl parameter to control what fraction of the CFL
        step we actually take.
        """

        myg = self.cc_data.grid

        cfl = self.rp.get_param("driver.cfl")

        u = self.cc_data.get_var("x-velocity")
        v = self.cc_data.get_var("y-velocity")

        # the timestep is min(dx/|u|, dy|v|)
        xtmp = ytmp = 1.e33
        if not abs(u).max() == 0:
            xtmp = myg.dx/abs(u.v()).max()
        if not abs(v).max() == 0:
            ytmp = myg.dy/abs(v.v()).max()

        dt = cfl*min(xtmp, ytmp)

        # We need an alternate timestep that accounts for buoyancy, to
        # handle the case where the velocity is initially zero.
        rho = self.cc_data.get_var("density")
        rho0 = self.base["rho0"]
        rhoprime = self.make_prime(rho, rho0)

        g = self.rp.get_param("lm-atmosphere.grav")

        F_buoy = (abs(rhoprime*g).v()/rho.v()).max()

        dt_buoy = np.sqrt(2.0*myg.dx/F_buoy)

        self.dt = min(dt, dt_buoy)
        if self.verbose > 0:
            print(f"timestep is {dt}")



[docs]
    def preevolve(self):
        """
        preevolve is called before we being the timestepping loop.  For
        the low Mach solver, this does an initial projection on the
        velocity field and then goes through the full evolution to get the
        value of phi.  The fluid state (rho, u, v) is then reset to values
        before this evolve.
        """

        self.in_preevolve = True

        myg = self.cc_data.grid

        rho = self.cc_data.get_var("density")
        u = self.cc_data.get_var("x-velocity")
        v = self.cc_data.get_var("y-velocity")

        self.cc_data.fill_BC("density")
        self.cc_data.fill_BC("x-velocity")
        self.cc_data.fill_BC("y-velocity")

        # 1. do the initial projection.  This makes sure that our original
        # velocity field satisfies div U = 0

        # the coefficient for the elliptic equation is beta_0^2/rho
        coeff = 1/rho
        beta0 = self.base["beta0"]
        coeff.v()[:, :] = coeff.v()*beta0.v2d()**2

        # next create the multigrid object.  We defined phi with
        # the right BCs previously
        mg = vcMG.VarCoeffCCMG2d(myg.nx, myg.ny,
                                 xl_BC_type=self.cc_data.BCs["phi"].xlb,
                                 xr_BC_type=self.cc_data.BCs["phi"].xrb,
                                 yl_BC_type=self.cc_data.BCs["phi"].ylb,
                                 yr_BC_type=self.cc_data.BCs["phi"].yrb,
                                 xmin=myg.xmin, xmax=myg.xmax,
                                 ymin=myg.ymin, ymax=myg.ymax,
                                 coeffs=coeff,
                                 coeffs_bc=self.cc_data.BCs["density"],
                                 verbose=0)

        # first compute div{beta_0 U}
        div_beta_U = mg.soln_grid.scratch_array()

        # u/v are cell-centered, divU is cell-centered
        div_beta_U.v()[:, :] = \
            0.5*beta0.v2d()*(u.ip(1) - u.ip(-1))/myg.dx + \
            0.5*(beta0.v2dp(1)*v.jp(1) - beta0.v2dp(-1)*v.jp(-1))/myg.dy

        # solve D (beta_0^2/rho) G (phi/beta_0) = D( beta_0 U )

        # set the RHS to divU and solve
        mg.init_RHS(div_beta_U)
        mg.solve(rtol=1.e-10)

        # store the solution in our self.cc_data object -- include a single
        # ghostcell
        phi = self.cc_data.get_var("phi")
        phi[:, :] = mg.get_solution(grid=myg)

        # get the cell-centered gradient of phi and update the
        # velocities
        # FIXME: this update only needs to be done on the interior
        # cells -- not ghost cells
        gradp_x, gradp_y = mg.get_solution_gradient(grid=myg)

        coeff = 1.0/rho
        coeff.v()[:, :] = coeff.v()*beta0.v2d()

        u.v()[:, :] -= coeff.v()*gradp_x.v()
        v.v()[:, :] -= coeff.v()*gradp_y.v()

        # fill the ghostcells
        self.cc_data.fill_BC("x-velocity")
        self.cc_data.fill_BC("y-velocity")

        # 2. now get an approximation to gradp at n-1/2 by going through the
        # evolution.

        # store the current solution -- we'll restore it in a bit
        orig_data = patch.cell_center_data_clone(self.cc_data)

        # get the timestep
        self.method_compute_timestep()

        # evolve
        self.evolve()

        # update gradp_x and gradp_y in our main data object
        new_gp_x = self.cc_data.get_var("gradp_x")
        new_gp_y = self.cc_data.get_var("gradp_y")

        orig_gp_x = orig_data.get_var("gradp_x")
        orig_gp_y = orig_data.get_var("gradp_y")

        orig_gp_x[:, :] = new_gp_x[:, :]
        orig_gp_y[:, :] = new_gp_y[:, :]

        self.cc_data = orig_data

        if self.verbose > 0:
            print("done with the pre-evolution")

        self.in_preevolve = False



[docs]
    def evolve(self):
        """
        Evolve the low Mach system through one timestep.
        """

        rho = self.cc_data.get_var("density")
        u = self.cc_data.get_var("x-velocity")
        v = self.cc_data.get_var("y-velocity")

        gradp_x = self.cc_data.get_var("gradp_x")
        gradp_y = self.cc_data.get_var("gradp_y")

        # note: the base state quantities do not have valid ghost cells
        beta0 = self.base["beta0"]
        beta0_edges = self.base["beta0-edges"]

        rho0 = self.base["rho0"]

        phi = self.cc_data.get_var("phi")

        myg = self.cc_data.grid

        # ---------------------------------------------------------------------
        # create the limited slopes of rho, u and v (in both directions)
        # ---------------------------------------------------------------------
        limiter = self.rp.get_param("lm-atmosphere.limiter")

        ldelta_rx = reconstruction.limit(rho, myg, 1, limiter)
        ldelta_ux = reconstruction.limit(u, myg, 1, limiter)
        ldelta_vx = reconstruction.limit(v, myg, 1, limiter)

        ldelta_ry = reconstruction.limit(rho, myg, 2, limiter)
        ldelta_uy = reconstruction.limit(u, myg, 2, limiter)
        ldelta_vy = reconstruction.limit(v, myg, 2, limiter)

        # ---------------------------------------------------------------------
        # get the advective velocities
        # ---------------------------------------------------------------------

        """
        the advective velocities are the normal velocity through each cell
        interface, and are defined on the cell edges, in a MAC type
        staggered form

                         n+1/2
                        v
                         i,j+1/2
                    +------+------+
                    |             |
            n+1/2   |             |   n+1/2
           u        +     U       +  u
            i-1/2,j |      i,j    |   i+1/2,j
                    |             |
                    +------+------+
                         n+1/2
                        v
                         i,j-1/2

        """

        # this returns u on x-interfaces and v on y-interfaces.  These
        # constitute the MAC grid
        if self.verbose > 0:
            print("  making MAC velocities")

        # create the coefficient to the grad (pi/beta) term
        coeff = self.aux_data.get_var("coeff")
        coeff.v()[:, :] = 1.0/rho.v()
        coeff.v()[:, :] = coeff.v()*beta0.v2d()
        self.aux_data.fill_BC("coeff")

        # create the source term
        source = self.aux_data.get_var("source_y")

        g = self.rp.get_param("lm-atmosphere.grav")
        rhoprime = self.make_prime(rho, rho0)
        source.v()[:, :] = rhoprime.v()*g/rho.v()
        self.aux_data.fill_BC("source_y")

        _um, _vm = lm_interface.mac_vels(myg.ng, myg.dx, myg.dy, self.dt,
                                           u, v,
                                           ldelta_ux, ldelta_vx,
                                           ldelta_uy, ldelta_vy,
                                           coeff*gradp_x, coeff*gradp_y,
                                           source)

        u_MAC = ai.ArrayIndexer(d=_um, grid=myg)
        v_MAC = ai.ArrayIndexer(d=_vm, grid=myg)

        # ---------------------------------------------------------------------
        # do a MAC projection to make the advective velocities divergence
        # free
        # ---------------------------------------------------------------------

        # we will solve D (beta_0^2/rho) G phi = D (beta_0 U^MAC), where
        # phi is cell centered, and U^MAC is the MAC-type staggered
        # grid of the advective velocities.

        if self.verbose > 0:
            print("  MAC projection")

        # create the coefficient array: beta0**2/rho
        # MZ!!!! probably don't need the buf here
        coeff.v(buf=1)[:, :] = 1.0/rho.v(buf=1)
        coeff.v(buf=1)[:, :] = coeff.v(buf=1)*beta0.v2d(buf=1)**2

        # create the multigrid object
        mg = vcMG.VarCoeffCCMG2d(myg.nx, myg.ny,
                                 xl_BC_type=self.cc_data.BCs["phi-MAC"].xlb,
                                 xr_BC_type=self.cc_data.BCs["phi-MAC"].xrb,
                                 yl_BC_type=self.cc_data.BCs["phi-MAC"].ylb,
                                 yr_BC_type=self.cc_data.BCs["phi-MAC"].yrb,
                                 xmin=myg.xmin, xmax=myg.xmax,
                                 ymin=myg.ymin, ymax=myg.ymax,
                                 coeffs=coeff,
                                 coeffs_bc=self.cc_data.BCs["density"],
                                 verbose=0)

        # first compute div{beta_0 U}
        div_beta_U = mg.soln_grid.scratch_array()

        # MAC velocities are edge-centered.  div{beta_0 U} is cell-centered.
        div_beta_U.v()[:, :] = \
            beta0.v2d()*(u_MAC.ip(1) - u_MAC.v())/myg.dx + \
            (beta0_edges.v2dp(1)*v_MAC.jp(1) -
             beta0_edges.v2d()*v_MAC.v())/myg.dy

        # solve the Poisson problem
        mg.init_RHS(div_beta_U)
        mg.solve(rtol=1.e-12)

        # update the normal velocities with the pressure gradient -- these
        # constitute our advective velocities.  Note that what we actually
        # solved for here is phi/beta_0
        phi_MAC = self.cc_data.get_var("phi-MAC")
        phi_MAC[:, :] = mg.get_solution(grid=myg)

        coeff = self.aux_data.get_var("coeff")
        coeff.v()[:, :] = 1.0/rho.v()
        coeff.v()[:, :] = coeff.v()*beta0.v2d()
        self.aux_data.fill_BC("coeff")

        coeff_x = myg.scratch_array()
        b = (3, 1, 0, 0)  # this seems more than we need
        coeff_x.v(buf=b)[:, :] = 0.5*(coeff.ip(-1, buf=b) + coeff.v(buf=b))

        coeff_y = myg.scratch_array()
        b = (0, 0, 3, 1)
        coeff_y.v(buf=b)[:, :] = 0.5*(coeff.jp(-1, buf=b) + coeff.v(buf=b))

        # we need the MAC velocities on all edges of the computational domain
        # here we do U = U - (beta_0/rho) grad (phi/beta_0)
        b = (0, 1, 0, 0)
        u_MAC.v(buf=b)[:, :] -= \
                coeff_x.v(buf=b)*(phi_MAC.v(buf=b) - phi_MAC.ip(-1, buf=b))/myg.dx

        b = (0, 0, 0, 1)
        v_MAC.v(buf=b)[:, :] -= \
                coeff_y.v(buf=b)*(phi_MAC.v(buf=b) - phi_MAC.jp(-1, buf=b))/myg.dy

        # ---------------------------------------------------------------------
        # predict rho to the edges and do its conservative update
        # ---------------------------------------------------------------------
        _rx, _ry = lm_interface.rho_states(myg.ng, myg.dx, myg.dy, self.dt,
                                             rho, u_MAC, v_MAC,
                                             ldelta_rx, ldelta_ry)

        rho_xint = ai.ArrayIndexer(d=_rx, grid=myg)
        rho_yint = ai.ArrayIndexer(d=_ry, grid=myg)

        rho_old = rho.copy()

        rho.v()[:, :] -= self.dt*(
            #  (rho u)_x
            (rho_xint.ip(1)*u_MAC.ip(1) - rho_xint.v()*u_MAC.v())/myg.dx +
            #  (rho v)_y
            (rho_yint.jp(1)*v_MAC.jp(1) - rho_yint.v()*v_MAC.v())/myg.dy)

        self.cc_data.fill_BC("density")

        # update eint as a diagnostic
        eint = self.cc_data.get_var("eint")
        gamma = self.rp.get_param("eos.gamma")
        eint.v()[:, :] = self.base["p0"].v2d()/(gamma - 1.0)/rho.v()

        # ---------------------------------------------------------------------
        # recompute the interface states, using the advective velocity
        # from above
        # ---------------------------------------------------------------------
        if self.verbose > 0:
            print("  making u, v edge states")

        coeff = self.aux_data.get_var("coeff")
        coeff.v()[:, :] = 2.0/(rho.v() + rho_old.v())
        coeff.v()[:, :] = coeff.v()*beta0.v2d()
        self.aux_data.fill_BC("coeff")

        _ux, _vx, _uy, _vy = \
               lm_interface.states(myg.ng, myg.dx, myg.dy, self.dt,
                                     u, v,
                                     ldelta_ux, ldelta_vx,
                                     ldelta_uy, ldelta_vy,
                                     coeff*gradp_x, coeff*gradp_y,
                                     source,
                                     u_MAC, v_MAC)

        u_xint = ai.ArrayIndexer(d=_ux, grid=myg)
        v_xint = ai.ArrayIndexer(d=_vx, grid=myg)
        u_yint = ai.ArrayIndexer(d=_uy, grid=myg)
        v_yint = ai.ArrayIndexer(d=_vy, grid=myg)

        # ---------------------------------------------------------------------
        # update U to get the provisional velocity field
        # ---------------------------------------------------------------------
        if self.verbose > 0:
            print("  doing provisional update of u, v")

        # compute (U.grad)U

        # we want u_MAC U_x + v_MAC U_y
        advect_x = myg.scratch_array()
        advect_y = myg.scratch_array()

        advect_x.v()[:, :] = \
            0.5*(u_MAC.v() + u_MAC.ip(1))*(u_xint.ip(1) - u_xint.v())/myg.dx +\
            0.5*(v_MAC.v() + v_MAC.jp(1))*(u_yint.jp(1) - u_yint.v())/myg.dy

        advect_y.v()[:, :] = \
            0.5*(u_MAC.v() + u_MAC.ip(1))*(v_xint.ip(1) - v_xint.v())/myg.dx +\
            0.5*(v_MAC.v() + v_MAC.jp(1))*(v_yint.jp(1) - v_yint.v())/myg.dy

        proj_type = self.rp.get_param("lm-atmosphere.proj_type")

        if proj_type == 1:
            u.v()[:, :] -= (self.dt*advect_x.v() + self.dt*gradp_x.v())
            v.v()[:, :] -= (self.dt*advect_y.v() + self.dt*gradp_y.v())

        elif proj_type == 2:
            u.v()[:, :] -= self.dt*advect_x.v()
            v.v()[:, :] -= self.dt*advect_y.v()

        # add the gravitational source
        rho_half = 0.5*(rho + rho_old)
        rhoprime = self.make_prime(rho_half, rho0)
        source[:, :] = rhoprime*g/rho_half
        self.aux_data.fill_BC("source_y")

        v[:, :] += self.dt*source

        self.cc_data.fill_BC("x-velocity")
        self.cc_data.fill_BC("y-velocity")

        if self.verbose > 0:
            print("min/max rho = {}, {}".format(self.cc_data.min("density"), self.cc_data.max("density")))
            print("min/max u   = {}, {}".format(self.cc_data.min("x-velocity"), self.cc_data.max("x-velocity")))
            print("min/max v   = {}, {}".format(self.cc_data.min("y-velocity"), self.cc_data.max("y-velocity")))

        # ---------------------------------------------------------------------
        # project the final velocity
        # ---------------------------------------------------------------------

        # now we solve L phi = D (U* /dt)
        if self.verbose > 0:
            print("  final projection")

        # create the coefficient array: beta0**2/rho
        coeff = 1.0/rho
        coeff.v()[:, :] = coeff.v()*beta0.v2d()**2

        # create the multigrid object
        mg = vcMG.VarCoeffCCMG2d(myg.nx, myg.ny,
                                 xl_BC_type=self.cc_data.BCs["phi"].xlb,
                                 xr_BC_type=self.cc_data.BCs["phi"].xrb,
                                 yl_BC_type=self.cc_data.BCs["phi"].ylb,
                                 yr_BC_type=self.cc_data.BCs["phi"].yrb,
                                 xmin=myg.xmin, xmax=myg.xmax,
                                 ymin=myg.ymin, ymax=myg.ymax,
                                 coeffs=coeff,
                                 coeffs_bc=self.cc_data.BCs["density"],
                                 verbose=0)

        # first compute div{beta_0 U}

        # u/v are cell-centered, divU is cell-centered
        div_beta_U.v()[:, :] = \
            0.5*beta0.v2d()*(u.ip(1) - u.ip(-1))/myg.dx + \
            0.5*(beta0.v2dp(1)*v.jp(1) - beta0.v2dp(-1)*v.jp(-1))/myg.dy

        mg.init_RHS(div_beta_U/self.dt)

        # use the old phi as our initial guess
        phiGuess = mg.soln_grid.scratch_array()
        phiGuess.v(buf=1)[:, :] = phi.v(buf=1)
        mg.init_solution(phiGuess)

        # solve
        mg.solve(rtol=1.e-12)

        # store the solution in our self.cc_data object -- include a single
        # ghostcell
        phi[:, :] = mg.get_solution(grid=myg)

        # get the cell-centered gradient of p and update the velocities
        # this differs depending on what we projected.
        gradphi_x, gradphi_y = mg.get_solution_gradient(grid=myg)

        # U = U - (beta_0/rho) grad (phi/beta_0)
        coeff = 1.0/rho
        coeff.v()[:, :] = coeff.v()*beta0.v2d()

        u.v()[:, :] -= self.dt*coeff.v()*gradphi_x.v()
        v.v()[:, :] -= self.dt*coeff.v()*gradphi_y.v()

        # store gradp for the next step

        if proj_type == 1:
            gradp_x.v()[:, :] += gradphi_x.v()
            gradp_y.v()[:, :] += gradphi_y.v()

        elif proj_type == 2:
            gradp_x.v()[:, :] = gradphi_x.v()
            gradp_y.v()[:, :] = gradphi_y.v()

        self.cc_data.fill_BC("x-velocity")
        self.cc_data.fill_BC("y-velocity")

        self.cc_data.fill_BC("gradp_x")
        self.cc_data.fill_BC("gradp_y")

        # increment the time
        if not self.in_preevolve:
            self.cc_data.t += self.dt
            self.n += 1



[docs]
    def dovis(self):
        """
        Do runtime visualization
        """
        plt.clf()

        # plt.rc("font", size=10)

        rho = self.cc_data.get_var("density")
        rho0 = self.base["rho0"]
        rhoprime = self.make_prime(rho, rho0)

        u = self.cc_data.get_var("x-velocity")
        v = self.cc_data.get_var("y-velocity")

        myg = self.cc_data.grid

        magvel = np.sqrt(u**2 + v**2)

        vort = myg.scratch_array()

        dv = 0.5*(v.ip(1) - v.ip(-1))/myg.dx
        du = 0.5*(u.jp(1) - u.jp(-1))/myg.dy

        vort.v()[:, :] = dv - du

        _, axes = plt.subplots(nrows=2, ncols=2, num=1)
        plt.subplots_adjust(hspace=0.25)

        fields = [rho, magvel, vort, rhoprime]
        field_names = [r"$\rho$", r"|U|", r"$\nabla \times U$", r"$\rho'$"]

        for n, f in enumerate(fields):
            ax = axes.flat[n]

            img = ax.imshow(np.transpose(f.v()),
                            interpolation="nearest", origin="lower",
                            extent=[myg.xmin, myg.xmax, myg.ymin, myg.ymax], cmap=self.cm)

            ax.set_xlabel("x")
            ax.set_ylabel("y")
            ax.set_title(field_names[n])

            plt.colorbar(img, ax=ax)

        plt.figtext(0.05, 0.0125, f"t = {self.cc_data.t:10.5f}")

        plt.pause(0.001)
        plt.draw()



[docs]
    def write_extras(self, f):
        """
        Output simulation-specific data to the h5py file f
        """

        # we implement our own version to allow us to store the base
        # state

        gb = f.create_group("base state")

        for name, state in self.base.items():
            gb.create_dataset(name, data=state.d)



[docs]
    def read_extras(self, f):
        """
        read in any simulation-specific data from an h5py file object f
        """

        gb = f["base state"]
        for name in gb:
            self.base[name] = Basestate(self.cc_data.grid.ny, ng=self.cc_data.grid.ng)
            self.base[name].d[:] = gb[name]