source: development/stochastic_study/run_model.py @ 3290

"""Stochastic study of the ANUGA implementation of the
shallow water wave equation.

This script runs the model for one realisation of bathymetry as
given in the file bathymetry.txt and outputs a full simulation is \
sww NetCDF format.

The left boundary condition is a timeseries defined in
NetCDF file: input_wave.tms

Note: This scripts needs create_mesh.py to have been run

Suresh Kumar and Ole Nielsen 2006
"""


#------------------------------------------------------------------------------
# Import necessary modules
#------------------------------------------------------------------------------

# Standard modules
import os
import time
import cPickle

# Related major packages
from pyvolution.shallow_water import Domain
from pyvolution.shallow_water import Reflective_boundary
from pyvolution.shallow_water import Transmissive_Momentum_Set_Stage_boundary
from pyvolution.pmesh2domain import pmesh_to_domain_instance
from pyvolution.data_manager import xya2pts
from pyvolution.util import file_function
from caching.caching import cache

# Application specific imports
import project                 # Definition of file names and polygons


#-----------------------------------------------------------------------------
# Read in processor information
#-----------------------------------------------------------------------------

try:
    import pypar
except:
    print 'Could not import pypar'
    myid = 0
    numprocs = 1
    processor_name = 'local host'
else:    
    myid = pypar.rank()
    numprocs = pypar.size()
    processor_name = pypar.Get_processor_name()

print 'I am process %d of %d running on %s' %(myid, numprocs, processor_name)


#-----------------------------------------------------------------------------
# Setup computational domain
#----------------------------------------------------------------------------- 
#print 'Creating domain from', project.mesh_filename

domain = Domain(project.working_dir + project.mesh_filename,
                use_cache=False,
                verbose=False)                
                

#print 'Number of triangles = ', len(domain)
#print domain.statistics()


domain.set_datadir(project.working_dir)
domain.set_quantities_to_be_stored(['stage', 'xmomentum', 'ymomentum'])


#------------------------------------------------------------------------------
# Setup boundary conditions
#------------------------------------------------------------------------------

function = file_function(project.boundary_filename, domain, verbose = False)
Bts = Transmissive_Momentum_Set_Stage_boundary(domain, function) #Input wave
Br = Reflective_boundary(domain) #Wall

# Bind boundary objects to tags
domain.set_boundary({'wave': Bts, 'wall': Br})


#------------------------------------------------------------------------------
# Setup initial conditions
#------------------------------------------------------------------------------
domain.set_quantity('friction', 0.0)
domain.set_quantity('stage', 0.0)

# Get prefitted realisations

finaltime = 22.5
timestep = 0.05


realisation = 0
for filename in os.listdir(project.working_dir):
    if filename.startswith(project.basename) and filename.endswith('.pck'):
        print 'P%d: Reading %s' %(myid, filename)
        fid = open(project.working_dir + filename)
        V = cPickle.load(fid)
        fid.close()

        #if myid == 0:
        #    print 'V', V[6:7,:]
            
        # For each column (each realisation)
        for i in range(V.shape[1]):

            # Distribute work in round-robin fashion
            if i%numprocs == myid:
                
                name = project.basename + '_P%d' %myid    
                domain.set_name(name)                     #Output name
                print 'V', V.shape
                domain.set_quantity('elevation', V[:,i])  #Assign bathymetry

                print 'P%d: Setting quantity %d: %s' %(myid, i, str(V[:4,i]))
                
                domain.set_time(0.0)                      #Reset time

                #---------------------------------------------------
                # Evolve system through time
                #---------------------------------------------------
                print 'P%d: Running realisation %d of %d in block %s'\
                      %(myid, realisation, V.shape[1], filename)

                t0 = time.time()
                for t in domain.evolve(yieldstep = timestep,
                                       finaltime = finaltime):
                    pass
                    domain.write_time()
                    
        
                print 'P%d: Simulation of realisation %d took %.2f seconds'\
                      %(myid, realisation, time.time()-t0)




                #---------------------------------------------------
                # Now extract the 3 timeseries (Ch 5-7-9) and store them
                # in three files for this realisation

                print 'P%d: Extracting time series for realisation %d from file %s'\
                      %(myid, realisation, project.working_dir + domain.filename + '.sww')
                f = file_function(project.working_dir + domain.filename + '.sww',
                                  quantities='stage',
                                  interpolation_points=project.gauges,
                                  verbose=False)


                simulation_name = project.working_dir + \
                                  project.basename + '_realisation_%d' %realisation

                print 'P%d: Writing to file %s'\
                      %(myid, simulation_name + '_' + name + '.txt')

                for k, name in enumerate(project.gauge_names):
                    fid = open(simulation_name + '_' + name + '.txt', 'w')
                    for t in f.get_time():
                        #For all precomputed timesteps
                        val = f(t, point_id = k)[0]
                        fid.write('%f %f\n' %(t, val))

                    fid.close()

                    

            realisation += 1            


pypar.finalize()