<commit_before>"""
Module for 2D solver class.

<NAME> 2015-10-17
"""
from utility import *
import shallowWaterEq
import timeIntegrator as timeIntegrator
import limiter
import time as timeMod
from mpi4py import MPI
import exporter
import ufl
import weakref
from cofs.fieldDefs import fieldMetadata
from cofs.options import modelOptions


class flowSolver2d(frozenClass):
    """Creates and solves 2D depth averaged equations with RT1-P1DG elements"""
    def __init__(self, mesh2d, bathymetry_2d, order=1, options={}):
        self._initialized = False

        # create 3D mesh
        self.mesh2d = mesh2d

        # Time integrator setup
        self.dt = None

        # 2d model specific default options
        options.setdefault('timeStepperType', 'SSPRK33')
        options.setdefault('timerLabels', ['mode2d'])
        options.setdefault('fieldsToExport', ['elev_2d', 'uv_2d'])

        # override default options
        self.options = modelOptions()
        self.options.update(options)

        # simulation time step bookkeeping
        self.simulation_time = 0
        self.iteration = 0
        self.iExport = 1

        self.visu_spaces = {}
        """Maps function space to a space where fields will be projected to for visualization"""

        self.fields = fieldDict()
        """Holds all functions needed by the solver object."""
        self.function_spaces = AttrDict()
        """Holds all function spaces needed by the solver object."""
        self.fields.bathymetry_2d = bathymetry_2d

        self.bnd_functions = {'shallow_water': {}}
        self._isfrozen = True  # disallow creating new attributes

    def setTimeStep(self):
        self.dt = self.options.dt
        if self.dt is None:
            mesh2d_dt = self.eq_sw.getTimeStep(Umag=self.options.uAdvection)
            dt = self.options.cfl_2d*float(mesh2d_dt.dat.data.min()/20.0)
            dt = comm.allreduce(dt, op=MPI.MIN)
            self.dt = dt
        if commrank == 0:
            print 'dt =', self.dt
            sys.stdout.flush()

    def createFunctionSpaces(self):
        """Creates function spaces"""
        self._isfrozen = False
        # ----- function spaces: elev in H, uv in U, mixed is W
        self.function_spaces.P0_2d = FunctionSpace(self.mesh2d, 'DG', 0)
        self.function_spaces.P1_2d = FunctionSpace(self.mesh2d, 'CG', 1)
        self.function_spaces.P1v_2d = VectorFunctionSpace(self.mesh2d, 'CG', 1)
        # 2D velocity space
        if self.options.mimetic:
            self.function_spaces.U_2d = FunctionSpace(self.mesh2d, 'RT', self.options.order+1)
        else:
            self.function_spaces.U_2d = VectorFunctionSpace(self.mesh2d, 'DG', self.options.order, name='U_2d')
        # TODO can this be omitted?
        #self.function_spaces.U_scalar_2d = FunctionSpace(self.mesh2d, 'DG', self.options.order)
        self.function_spaces.H_2d = FunctionSpace(self.mesh2d, 'DG', self.options.order)
        self.function_spaces.V_2d = MixedFunctionSpace([self.function_spaces.U_2d, self.function_spaces.H_2d])

        self.visu_spaces[self.function_spaces.U_2d] = self.function_spaces.P1v_2d
        self.visu_spaces[self.function_spaces.H_2d] = self.function_spaces.P1_2d
        self._isfrozen = True

    def createEquations(self):
        """Creates functions, equations and time steppers."""
        if not hasattr(self, 'U_2d'):
            self.createFunctionSpaces()
        self._isfrozen = False
        # ----- fields
        self.fields.solution2d = Function(self.function_spaces.V_2d)

        # ----- Equations
        self.eq_sw = shallowWaterEq.shallowWaterEquations(
            self.fields.solution2d,
            self.fields.bathymetry_2d,
            lin_drag=self.options.lin_drag,
            viscosity_h=self.fields.get('hViscosity'),
            uvLaxFriedrichs=self.options.uvLaxFriedrichs,
            coriolis=self.options.coriolis,
            wind_stress=self.options.wind_stress,
            nonlin=self.options.nonlin)

        self.eq_sw.bnd_functions = self.bnd_functions['shallow_water']

        # ----- Time integrators
        self.setTimeStep()
        if self.options.timeStepperType.lower() == 'ssprk33':
            self.timeStepper = timeIntegrator.SSPRK33Stage(self.eq_sw, self.dt,
                                                            self.eq_sw.solver_parameters)
        elif self.options.timeStepperType.lower() == 'ssprk33semi':
            self.timeStepper = timeIntegrator.SSPRK33StageSemiImplicit(self.eq_sw,
                                                            self.dt, self.eq_sw.solver_parameters)
        elif self.options.timeStepperType.lower() == 'forwardeuler':
            self.timeStepper = timeIntegrator.ForwardEuler(self.eq_sw, self.dt,
                                                            self.eq_sw.solver_parameters)
        elif self.options.timeStepperType.lower() == 'cranknicolson':
            self.timeStepper = timeIntegrator.CrankNicolson(self.eq_sw, self.dt,
                                                             self.eq_sw.solver_parameters)
        elif self.options.timeStepperType.lower() == 'sspimex':
            # TODO meaningful solver params
            sp_impl = {
                'ksp_type': 'gmres',
                'pc_type': 'fieldsplit',
                'pc_fieldsplit_type': 'multiplicative',
                }
            sp_expl = {
                'ksp_type': 'gmres',
                'pc_type': 'fieldsplit',
                'pc_fieldsplit_type': 'multiplicative',
                }
            self.timeStepper = timeIntegrator.SSPIMEX(self.eq_sw, self.dt,
                                                      solver_parameters=sp_expl,
                                                      solver_parameters_dirk=sp_impl)
        else:
            raise Exception('Unknown time integrator type: '+str(self.options.timeStepperType))

        # ----- File exporters
        # correct treatment of the split 2d functions
        uv_2d, elev_2d = self.fields.solution2d.split()
        self.fields.uv_2d = uv_2d
        self.fields.elev_2d = elev_2d
        self.visu_spaces[uv_2d.function_space()] = self.function_spaces.P1v_2d
        self.visu_spaces[elev_2d.function_space()] = self.function_spaces.P1_2d
        self.exporters = {}
        e = exporter.exportManager(self.options.outputDir,
                                   self.options.fieldsToExport,
                                   self.fields,
                                   self.visu_spaces,
                                   fieldMetadata,
                                   exportType='vtk',
                                   verbose=self.options.verbose > 0)
        self.exporters['vtk'] = e
        numpyDir = os.path.join(self.options.outputDir, 'numpy')
        e = exporter.exportManager(numpyDir,
                                   self.options.fieldsToExportNumpy,
                                   self.fields,
                                   self.visu_spaces,
                                   fieldMetadata,
                                   exportType='numpy',
                                   verbose=self.options.verbose > 0)
        self.exporters['numpy'] = e
        hdf5Dir = os.path.join(self.options.outputDir, 'hdf5')
        e = exporter.exportManager(hdf5Dir,
                                   self.options.fieldsToExportHDF5,
                                   self.fields,
                                   self.visu_spaces,
                                   fieldMetadata,
                                   exportType='hdf5',
                                   verbose=self.options.verbose > 0)
        self.exporters['hdf5'] = e

        self._initialized = True
        self._isfrozen = True  # disallow creating new attributes

    def assignInitialConditions(self, elev=None, uv_init=None):
        if not self._initialized:
            self.createEquations()
        uv_2d, elev_2d = self.fields.solution2d.split()
        if elev is not None:
            elev_2d.project(elev)
        if uv_init is not None:
            uv_2d.project(uv_init)

        self.timeStepper.initialize(self.fields.solution2d)

    def export(self):
        for key in self.exporters:
            self.exporters[key].export()

    def loadState(self, iExport, t, iteration):
        """Loads simulation state from hdf5 outputs."""
        uv_2d, elev_2d = self.fields.solution2d.split()
        self.exporters['hdf5'].exporters['uv_2d'].load(iExport, uv_2d)
        self.exporters['hdf5'].exporters['elev_2d'].load(iExport, elev_2d)
        self.assignInitialConditions(elev=elev_2d, uv_init=uv_2d)
        self.iExport = iExport
        self.simulation_time = t
        self.iteration = iteration
        self.printState(0.0)
        self.iExport += 1
        for k in self.exporters:
            self.exporters[k].setNextExportIx(self.iExport)

    def printState(self, cputime):
        norm_h = norm(self.fields.solution2d.split()[1])
        norm_u = norm(self.fields.solution2d.split()[0])

        if self.options.checkVolConservation2d:
            Vol2d = compVolume2d(self.fields.solution2d.split()[1],
                                 self.fields.bathymetry_2d)
        if commrank == 0:
            line = ('{iexp:5d} {i:5d} T={t:10.2f} '
                    'eta norm: {e:10.4f} u norm: {u:10.4f} {cpu:5.2f}')
            print(bold(line.format(iexp=self.iExport, i=self.iteration, t=self.simulation_time, e=norm_h,
                                   u=norm_u, cpu=cputime)))
            sys.stdout.flush()

    def iterate(self, updateForcings=None,
                exportFunc=None):
        if not self._initialized:
            self.createEquations()

        T_epsilon = 1.0e-5
        cputimestamp = timeMod.clock()
        next_export_t = self.simulation_time + self.options.TExport

        # initialize conservation checks
        if self.options.checkVolConservation2d:
            eta = self.fields.solution2d.split()[1]
            Vol2d_0 = compVolume2d(eta, self.fields.bathymetry_2d)
            printInfo('Initial volume 2d {0:f}'.format(Vol2d_0))

        # initial export
        self.export()
        if exportFunc is not None:
            exportFunc()
        self.exporters['vtk'].exportBathymetry(self.fields.bathymetry_2d)

        while self.simulation_time <= self.options.T + T_epsilon:

            self.timeStepper.advance(self.simulation_time, self.dt, self.fields.solution2d,
                                     updateForcings)

            # Move to next time step
            self.iteration += 1
            self.simulation_time = self.iteration*self.dt

            # Write the solution to file
            if self.simulation_time >= next_export_t - T_epsilon:
                cputime = timeMod.clock() - cputimestamp
                cputimestamp = timeMod.clock()
                self.printState(cputime)
                if self.options.checkVolConservation2d:
                    Vol2d = compVolume2d(self.fields.solution2d.split()[1],
                                         self.fields.bathymetry_2d)
                if commrank == 0:
                    line = 'Rel. {0:s} error {1:11.4e}'
                    if self.options.checkVolConservation2d:
                        print(line.format('vol 2d', (Vol2d_0 - Vol2d)/Vol2d_0))
                    sys.stdout.flush()

                self.export()
                if exportFunc is not None:
                    exportFunc()

                next_export_t += self.options.TExport
                self.iExport += 1

                if commrank == 0 and len(self.options.timerLabels) > 0:
                    cost = {}
                    relcost = {}
                    totcost = 0
                    for label in self.options.timerLabels:
                        value = timing(label, reset=True)
                        cost[label] = value
                        totcost += value
                    for label in self.options.timerLabels:
                        c = cost[label]
                        relcost = c/max(totcost, 1e-6)
                        print '{0:25s} : {1:11.6f} {2:11.2f}'.format(
                            label, c, relcost)
                        sys.stdout.flush()
<commit_msg>added missing 2d source terms<commit_after>"""
Module for 2D solver class.

<NAME> 2015-10-17
"""
from utility import *
import shallowWaterEq
import timeIntegrator as timeIntegrator
import limiter
import time as timeMod
from mpi4py import MPI
import exporter
import ufl
import weakref
from cofs.fieldDefs import fieldMetadata
from cofs.options import modelOptions


class flowSolver2d(frozenClass):
    """Creates and solves 2D depth averaged equations with RT1-P1DG elements"""
    def __init__(self, mesh2d, bathymetry_2d, order=1, options={}):
        self._initialized = False

        # create 3D mesh
        self.mesh2d = mesh2d

        # Time integrator setup
        self.dt = None

        # 2d model specific default options
        options.setdefault('timeStepperType', 'SSPRK33')
        options.setdefault('timerLabels', ['mode2d'])
        options.setdefault('fieldsToExport', ['elev_2d', 'uv_2d'])

        # override default options
        self.options = modelOptions()
        self.options.update(options)

        # simulation time step bookkeeping
        self.simulation_time = 0
        self.iteration = 0
        self.iExport = 1

        self.visu_spaces = {}
        """Maps function space to a space where fields will be projected to for visualization"""

        self.fields = fieldDict()
        """Holds all functions needed by the solver object."""
        self.function_spaces = AttrDict()
        """Holds all function spaces needed by the solver object."""
        self.fields.bathymetry_2d = bathymetry_2d

        self.bnd_functions = {'shallow_water': {}}
        self._isfrozen = True  # disallow creating new attributes

    def setTimeStep(self):
        self.dt = self.options.dt
        if self.dt is None:
            mesh2d_dt = self.eq_sw.getTimeStep(Umag=self.options.uAdvection)
            dt = self.options.cfl_2d*float(mesh2d_dt.dat.data.min()/20.0)
            dt = comm.allreduce(dt, op=MPI.MIN)
            self.dt = dt
        if commrank == 0:
            print 'dt =', self.dt
            sys.stdout.flush()

    def createFunctionSpaces(self):
        """Creates function spaces"""
        self._isfrozen = False
        # ----- function spaces: elev in H, uv in U, mixed is W
        self.function_spaces.P0_2d = FunctionSpace(self.mesh2d, 'DG', 0)
        self.function_spaces.P1_2d = FunctionSpace(self.mesh2d, 'CG', 1)
        self.function_spaces.P1v_2d = VectorFunctionSpace(self.mesh2d, 'CG', 1)
        # 2D velocity space
        if self.options.mimetic:
            self.function_spaces.U_2d = FunctionSpace(self.mesh2d, 'RT', self.options.order+1)
        else:
            self.function_spaces.U_2d = VectorFunctionSpace(self.mesh2d, 'DG', self.options.order, name='U_2d')
        # TODO can this be omitted?
        #self.function_spaces.U_scalar_2d = FunctionSpace(self.mesh2d, 'DG', self.options.order)
        self.function_spaces.H_2d = FunctionSpace(self.mesh2d, 'DG', self.options.order)
        self.function_spaces.V_2d = MixedFunctionSpace([self.function_spaces.U_2d, self.function_spaces.H_2d])

        self.visu_spaces[self.function_spaces.U_2d] = self.function_spaces.P1v_2d
        self.visu_spaces[self.function_spaces.H_2d] = self.function_spaces.P1_2d
        self._isfrozen = True

    def createEquations(self):
        """Creates functions, equations and time steppers."""
        if not hasattr(self, 'U_2d'):
            self.createFunctionSpaces()
        self._isfrozen = False
        # ----- fields
        self.fields.solution2d = Function(self.function_spaces.V_2d)

        # ----- Equations
        self.eq_sw = shallowWaterEq.shallowWaterEquations(
            self.fields.solution2d,
            self.fields.bathymetry_2d,
            lin_drag=self.options.lin_drag,
            viscosity_h=self.fields.get('hViscosity'),
            uvLaxFriedrichs=self.options.uvLaxFriedrichs,
            coriolis=self.options.coriolis,
            wind_stress=self.options.wind_stress,
            uv_source=self.options.uv_source_2d,
            elev_source=self.options.elev_source_2d,
            nonlin=self.options.nonlin)

        self.eq_sw.bnd_functions = self.bnd_functions['shallow_water']

        # ----- Time integrators
        self.setTimeStep()
        if self.options.timeStepperType.lower() == 'ssprk33':
            self.timeStepper = timeIntegrator.SSPRK33Stage(self.eq_sw, self.dt,
                                                            self.eq_sw.solver_parameters)
        elif self.options.timeStepperType.lower() == 'ssprk33semi':
            self.timeStepper = timeIntegrator.SSPRK33StageSemiImplicit(self.eq_sw,
                                                            self.dt, self.eq_sw.solver_parameters)
        elif self.options.timeStepperType.lower() == 'forwardeuler':
            self.timeStepper = timeIntegrator.ForwardEuler(self.eq_sw, self.dt,
                                                            self.eq_sw.solver_parameters)
        elif self.options.timeStepperType.lower() == 'cranknicolson':
            self.timeStepper = timeIntegrator.CrankNicolson(self.eq_sw, self.dt,
                                                             self.eq_sw.solver_parameters)
        elif self.options.timeStepperType.lower() == 'sspimex':
            # TODO meaningful solver params
            sp_impl = {
                'ksp_type': 'gmres',
                'pc_type': 'fieldsplit',
                'pc_fieldsplit_type': 'multiplicative',
                }
            sp_expl = {
                'ksp_type': 'gmres',
                'pc_type': 'fieldsplit',
                'pc_fieldsplit_type': 'multiplicative',
                }
            self.timeStepper = timeIntegrator.SSPIMEX(self.eq_sw, self.dt,
                                                      solver_parameters=sp_expl,
                                                      solver_parameters_dirk=sp_impl)
        else:
            raise Exception('Unknown time integrator type: '+str(self.options.timeStepperType))

        # ----- File exporters
        # correct treatment of the split 2d functions
        uv_2d, elev_2d = self.fields.solution2d.split()
        self.fields.uv_2d = uv_2d
        self.fields.elev_2d = elev_2d
        self.visu_spaces[uv_2d.function_space()] = self.function_spaces.P1v_2d
        self.visu_spaces[elev_2d.function_space()] = self.function_spaces.P1_2d
        self.exporters = {}
        e = exporter.exportManager(self.options.outputDir,
                                   self.options.fieldsToExport,
                                   self.fields,
                                   self.visu_spaces,
                                   fieldMetadata,
                                   exportType='vtk',
                                   verbose=self.options.verbose > 0)
        self.exporters['vtk'] = e
        numpyDir = os.path.join(self.options.outputDir, 'numpy')
        e = exporter.exportManager(numpyDir,
                                   self.options.fieldsToExportNumpy,
                                   self.fields,
                                   self.visu_spaces,
                                   fieldMetadata,
                                   exportType='numpy',
                                   verbose=self.options.verbose > 0)
        self.exporters['numpy'] = e
        hdf5Dir = os.path.join(self.options.outputDir, 'hdf5')
        e = exporter.exportManager(hdf5Dir,
                                   self.options.fieldsToExportHDF5,
                                   self.fields,
                                   self.visu_spaces,
                                   fieldMetadata,
                                   exportType='hdf5',
                                   verbose=self.options.verbose > 0)
        self.exporters['hdf5'] = e

        self._initialized = True
        self._isfrozen = True  # disallow creating new attributes

    def assignInitialConditions(self, elev=None, uv_init=None):
        if not self._initialized:
            self.createEquations()
        uv_2d, elev_2d = self.fields.solution2d.split()
        if elev is not None:
            elev_2d.project(elev)
        if uv_init is not None:
            uv_2d.project(uv_init)

        self.timeStepper.initialize(self.fields.solution2d)

    def export(self):
        for key in self.exporters:
            self.exporters[key].export()

    def loadState(self, iExport, t, iteration):
        """Loads simulation state from hdf5 outputs."""
        uv_2d, elev_2d = self.fields.solution2d.split()
        self.exporters['hdf5'].exporters['uv_2d'].load(iExport, uv_2d)
        self.exporters['hdf5'].exporters['elev_2d'].load(iExport, elev_2d)
        self.assignInitialConditions(elev=elev_2d, uv_init=uv_2d)
        self.iExport = iExport
        self.simulation_time = t
        self.iteration = iteration
        self.printState(0.0)
        self.iExport += 1
        for k in self.exporters:
            self.exporters[k].setNextExportIx(self.iExport)

    def printState(self, cputime):
        norm_h = norm(self.fields.solution2d.split()[1])
        norm_u = norm(self.fields.solution2d.split()[0])

        if self.options.checkVolConservation2d:
            Vol2d = compVolume2d(self.fields.solution2d.split()[1],
                                 self.fields.bathymetry_2d)
        if commrank == 0:
            line = ('{iexp:5d} {i:5d} T={t:10.2f} '
                    'eta norm: {e:10.4f} u norm: {u:10.4f} {cpu:5.2f}')
            print(bold(line.format(iexp=self.iExport, i=self.iteration, t=self.simulation_time, e=norm_h,
                                   u=norm_u, cpu=cputime)))
            sys.stdout.flush()

    def iterate(self, updateForcings=None,
                exportFunc=None):
        if not self._initialized:
            self.createEquations()

        T_epsilon = 1.0e-5
        cputimestamp = timeMod.clock()
        next_export_t = self.simulation_time + self.options.TExport

        # initialize conservation checks
        if self.options.checkVolConservation2d:
            eta = self.fields.solution2d.split()[1]
            Vol2d_0 = compVolume2d(eta, self.fields.bathymetry_2d)
            printInfo('Initial volume 2d {0:f}'.format(Vol2d_0))

        # initial export
        self.export()
        if exportFunc is not None:
            exportFunc()
        self.exporters['vtk'].exportBathymetry(self.fields.bathymetry_2d)

        while self.simulation_time <= self.options.T + T_epsilon:

            self.timeStepper.advance(self.simulation_time, self.dt, self.fields.solution2d,
                                     updateForcings)

            # Move to next time step
            self.iteration += 1
            self.simulation_time = self.iteration*self.dt

            # Write the solution to file
            if self.simulation_time >= next_export_t - T_epsilon:
                cputime = timeMod.clock() - cputimestamp
                cputimestamp = timeMod.clock()
                self.printState(cputime)
                if self.options.checkVolConservation2d:
                    Vol2d = compVolume2d(self.fields.solution2d.split()[1],
                                         self.fields.bathymetry_2d)
                if commrank == 0:
                    line = 'Rel. {0:s} error {1:11.4e}'
                    if self.options.checkVolConservation2d:
                        print(line.format('vol 2d', (Vol2d_0 - Vol2d)/Vol2d_0))
                    sys.stdout.flush()

                self.export()
                if exportFunc is not None:
                    exportFunc()

                next_export_t += self.options.TExport
                self.iExport += 1

                if commrank == 0 and len(self.options.timerLabels) > 0:
                    cost = {}
                    relcost = {}
                    totcost = 0
                    for label in self.options.timerLabels:
                        value = timing(label, reset=True)
                        cost[label] = value
                        totcost += value
                    for label in self.options.timerLabels:
                        c = cost[label]
                        relcost = c/max(totcost, 1e-6)
                        print '{0:25s} : {1:11.6f} {2:11.2f}'.format(
                            label, c, relcost)
                        sys.stdout.flush()

{'fraction_non_alphanumeric': 0.09750162232316678, 'fraction_numerical': 0.01666937053861129, 'mean_word_length': 1.7033220041662098, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 70, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7172531', 'n_tokens_mistral': 7513, 'n_tokens_neox': 7041, 'n_words': 1280}