Add solution downsampling feature

src/petclaw/geometry.py

@@ -13,11 +13,11 @@ class Patch(pyclaw_geometry.Patch):
 
     __doc__ += pyclaw.util.add_parent_doc(pyclaw_geometry.Patch)
 
-    def __init__(self,dimensions):
+    def __init__(self,dimensions,proc_sizes=None):
 
         super(Patch,self).__init__(dimensions)
 
-        self._da = self._create_DA()
+        self._da = self._create_DA(proc_sizes)
         ranges = self._da.getRanges()
         grid_dimensions = []
         for i,nrange in enumerate(ranges):
@@ -41,7 +41,7 @@ def __init__(self,dimensions):
 
         self.grid = pyclaw_geometry.Grid(grid_dimensions)
 
-    def _create_DA(self):
+    def _create_DA(self,proc_sizes=None):
         r"""Returns a PETSc DA and associated global Vec.
         Note that no local vector is returned.
         """
@@ -62,14 +62,16 @@ def _create_DA(self):
                                           sizes=self.num_cells_global,
                                           periodic_type = periodic_type,
                                           stencil_width=0,
-                                          comm=PETSc.COMM_WORLD)
+                                          comm=PETSc.COMM_WORLD,
+                                          proc_sizes=proc_sizes)
         else:
             DA = PETSc.DA().create(dim=self.num_dim,
                                           dof=1,
                                           sizes=self.num_cells_global,
                                           boundary_type = PETSc.DA.BoundaryType.PERIODIC,
                                           stencil_width=0,
-                                          comm=PETSc.COMM_WORLD)
+                                          comm=PETSc.COMM_WORLD,
+                                          proc_sizes=proc_sizes)
 
         return DA
 
@@ -82,13 +84,13 @@ class Domain(pyclaw_geometry.Domain):
 
     __doc__ += pyclaw.util.add_parent_doc(pyclaw.ClawSolver2D)
 
-    def __init__(self,geom):
+    def __init__(self,geom,proc_sizes=None):
         if not isinstance(geom,list):
             geom = [geom]
         if isinstance(geom[0],Patch):
             self.patches = geom
         elif isinstance(geom[0],pyclaw_geometry.Dimension):
-            self.patches = [Patch(geom)]
+            self.patches = [Patch(geom, proc_sizes)]
 
 class Dimension(pyclaw_geometry.Dimension):
     def __init__(self, lower, upper, num_cells, name='x',

src/petclaw/solution.py

@@ -1,5 +1,6 @@
 from clawpack import pyclaw
 from clawpack.pyclaw.solution import Solution
+import numpy as np
 
 class Solution(Solution):
     """ Parallel Solution class.
@@ -24,4 +25,89 @@ def get_write_func(self, file_format):
         else:
             raise ValueError("File format %s not supported." % file_format)
 
+    def _init_ds_solution(self):
+        """
+        Initializes a downsampled version of the solution
+        """
+        import clawpack.petclaw as pyclaw
+
+        ds_domain = pyclaw.Domain([pyclaw.Dimension(self.domain.patch.lower_global[i],
+                                                    self.domain.patch.upper_global[i],
+                                                    self.domain.patch.num_cells_global[i]/self.downsampling_factors[i],
+                                                    self.domain.grid.dimensions[i].name)
+                                       for i in range(self.domain.num_dim)], proc_sizes=self.state.q_da.getProcSizes())
+        ds_state = self._init_ds_state(self.state)
+        self._ds_solution = pyclaw.Solution(ds_state, ds_domain)
+
+    def _init_ds_state(self, state):
+        """
+        Returns a downsampled version of the given state object
+        """
+        import clawpack.petclaw as pyclaw
+
+        ds_domain = pyclaw.Domain([pyclaw.Dimension(self.domain.patch.lower_global[i],
+                                                    self.domain.patch.upper_global[i],
+                                                    self.domain.patch.num_cells_global[i]/self.downsampling_factors[i],
+                                                    self.domain.grid.dimensions[i].name)
+                                   for i in range(self.domain.num_dim)], proc_sizes=state.q_da.getProcSizes())
+        ds_state = pyclaw.State(ds_domain,state.num_eqn,state.num_aux)
+
+        return ds_state
+
+    def downsample(self, write_aux,write_p):
+        """
+        Returns a downsampled version of the solution by local averaging over the downsampling factors
+        """
+        for i  in range(len(self.states)):
+            state = self.states[i]
+            if i > 0:
+                self.ds_solution.states.append(self._init_ds_state(self.downsample_factors, state))
+            ds_state = self.ds_solution.states[i]
+
+            # Downsample q
+            if write_p:
+                ds_state.p = self._downsample_global_to_local(state.get_q_da_for_ds(np.max(self.downsampling_factors)),
+                                                             state.gqVec, state.num_eqn, ds_state.patch._da.getRanges())
+            else:
+                ds_state.q = self._downsample_global_to_local(state.get_q_da_for_ds(np.max(self.downsampling_factors)),
+                                                             state.gqVec, state.num_eqn, ds_state.patch._da.getRanges())
+            # Dowsample aux
+            if write_aux:
+                ds_state.aux = self._downsample_global_to_local(state.get_aux_da_for_ds(np.max(self.downsampling_factors)),
+                                                               state.gauxVec, state.num_aux, ds_state.patch._da.getRanges())
+
+        return self.ds_solution
+
+    def _downsample_global_to_local(self, da_for_ds, gVec, num_eqn, ds_domain_ranges):
+        """
+        Returns a locally averaged array and handles ranges mapping between the original & downsampled solution objects
+
+        Input:
+            - da_for_ds: the DA object of the original solution (q or aux) with stencil width adjusted as appropriate to
+              handle the boundaries of the array to be averaged because after domain decomposition, there are no guarantees that
+              the downsampling factors will evenly divide the the sub-domain in each direction
+            - gVec: The global vector associated with q or aux
+            - num_eqn: number of components of q or aux
+            - ds_domain_ranges: global ranges of the downsampled solution
+        """
+
+        from skimage.transform import downscale_local_mean
+
+        # Create local array with ghost cells
+        q_local_vec = da_for_ds.createLocalVec()
+        da_for_ds.globalToLocal(gVec, q_local_vec)
+        q_local_array = q_local_vec.array
+        shape = [end-start for start,end in da_for_ds.getGhostRanges()]
+        shape.insert(0, num_eqn)
+        q_local_array = q_local_array.reshape(shape, order='f')
+
+        # Map global ranges in the original DA object to local ranges for local averaging
+        q_da_ghost_ranges = da_for_ds.getGhostRanges()
+        new_global_ranges = tuple([(s * self.downsampling_factors[i], e * self.downsampling_factors[i]) for i, (s,e) in enumerate(ds_domain_ranges)])
+        new_local_slices = (slice(0, q_local_array.shape[0]),)
+        new_local_slices = new_local_slices + tuple([slice(s - gs,q_local_array.shape[i+1] - (ge - e)) for i, ((s,e), (gs, ge)) in enumerate(zip(new_global_ranges, q_da_ghost_ranges))])
+        q_local_array = q_local_array[new_local_slices]
+
+        # Compute local mean
+        return downscale_local_mean(q_local_array, (1,) + self.downsampling_factors)
 

src/petclaw/state.py

@@ -107,7 +107,7 @@ def aux(self):
     def aux(self,val):
         # It would be nice to make this work also for parallel
         # loading from a file.
-        if self.aux_da is None: 
+        if self.aux_da is None:
             num_aux=val.shape[0]
             self._init_aux_da(num_aux)
         self.gauxVec.setArray(val.reshape([-1], order = 'F'))
@@ -136,6 +136,8 @@ def __init__(self,geom,num_eqn,num_aux=0):
 
         self.aux_da = None
         self.q_da = None
+        self._q_da_for_ds = None
+        self._aux_da_for_ds = None
 
         self._p_da = None
         self.gpVec = None
@@ -145,14 +147,14 @@ def __init__(self,geom,num_eqn,num_aux=0):
 
         # ========== Attribute Definitions ===================================
         self.problem_data = {}
-        r"""(dict) - Dictionary of global values for this patch, 
+        r"""(dict) - Dictionary of global values for this patch,
             ``default = {}``"""
         self.t=0.
-        r"""(float) - Current time represented on this patch, 
+        r"""(float) - Current time represented on this patch,
             ``default = 0.0``"""
         self.index_capa = -1
         self.keep_gauges = False
-        r"""(bool) - Keep gauge values in memory for every time step, 
+        r"""(bool) - Keep gauge values in memory for every time step,
         ``default = False``"""
         self.gauge_data = []
         r"""(list) - List of numpy.ndarray objects. Each element of the list
@@ -165,18 +167,18 @@ def __init__(self,geom,num_eqn,num_aux=0):
     def _init_aux_da(self,num_aux,num_ghost=0):
         r"""
         Initializes PETSc DA and global & local Vectors for handling the
-        auxiliary array, aux. 
-        
+        auxiliary array, aux.
+
         Initializes aux_da, gauxVec and _aux_local_vector.
         """
         self.aux_da = self._create_DA(num_aux,num_ghost)
         self.gauxVec = self.aux_da.createGlobalVector()
         self._aux_local_vector = self.aux_da.createLocalVector()
- 
+
     def _init_q_da(self,num_eqn,num_ghost=0):
         r"""
-        Initializes PETSc DA and Vecs for handling the solution, q. 
-        
+        Initializes PETSc DA and Vecs for handling the solution, q.
+
         Initializes q_da, gqVec and _q_local_vector.
         """
         self.q_da = self._create_DA(num_eqn,num_ghost)
@@ -225,7 +227,7 @@ def get_qbc_from_q(self,num_ghost,qbc):
         Returns q with ghost cells attached, by accessing the local vector.
         """
         shape = [n + 2*num_ghost for n in self.grid.num_cells]
-        
+
         self.q_da.globalToLocal(self.gqVec, self._q_local_vector)
         shape.insert(0,self.num_eqn)
         return self._q_local_vector.getArray().reshape(shape, order = 'F')
@@ -235,7 +237,7 @@ def get_auxbc_from_aux(self,num_ghost,auxbc):
         Returns aux with ghost cells attached, by accessing the local vector.
         """
         shape = [n + 2*num_ghost for n in self.grid.num_cells]
-        
+
         self.aux_da.globalToLocal(self.gauxVec, self._aux_local_vector)
         shape.insert(0,self.num_aux)
         return self._aux_local_vector.getArray().reshape(shape, order = 'F')
@@ -316,3 +318,27 @@ def __deepcopy__(self,memo={}):
         result.set_num_ghost(self.q_da.stencil_width) 
 
         return result
+
+    def get_q_da_for_ds(self, ds_stencil_width):
+        """
+        Returns a q_da object with an adjusted width if needed in order to compute local averages for downsampling
+        """
+        if self._q_da_for_ds is None:
+            if self.q_da.getStencilWidth() < ds_stencil_width:
+                self._q_da_for_ds = self.q_da.duplicate(stencil_width=ds_stencil_width)
+            else:
+                self._q_da_for_ds = self.q_da
+
+        return self._q_da_for_ds
+
+    def get_aux_da_for_ds(self, ds_stencil_width):
+        """
+        Returns an aux_da object with an adjusted width if needed in order to compute local averages for downsampling
+        """
+        if self._aux_da_for_ds is None:
+            if self.aux_da.getStencilWidth() < ds_stencil_width:
+                self._aux_da_for_ds = self.aux_da.duplicate(stencil_width=ds_stencil_width)
+            else:
+                self._aux_da_for_ds = self.aux_da
+
+        return self._aux_da_for_ds

src/pyclaw/controller.py

@@ -193,6 +193,7 @@ def __init__(self):
         r"""(function) - Function that computes density of functional F"""
         self.F_file_name = 'F'
         r"""(string) - Name of text file containing functionals"""
+        self.downsampling_factors = None
 
     # ========== Access methods ===============================================
     def __str__(self):        
@@ -326,11 +327,20 @@ def run(self):
             if self.t == self.tfinal:
                 print "Simulation has already reached tfinal."
             return None
-
+
+        self.solution.downsampling_factors = self.downsampling_factors
+
         # Output and save initial frame
         if self.keep_copy:
             self.frames.append(copy.deepcopy(self.solution))
         if self.output_format is not None:
+            if not (self.downsampling_factors is None):
+                if self.solution.domain.num_dim != len(self.downsampling_factors):
+                    raise ValueError("Invalid number of downsampling factors. The number of downsampling factors should match the number of dimensions.")
+                for i, factor in enumerate(self.downsampling_factors):
+                    if self.solution.domain.patch.dimensions[i].num_cells % factor != 0:
+                        raise ValueError("Invalid downsampling factors. The downsampling factors should evenly divide the grid in each dimension.")
+
             if os.path.exists(self.outdir) and self.overwrite==False:
                 raise Exception("Refusing to overwrite existing output data. \
                  \nEither delete/move the directory or set controller.overwrite=True.")
@@ -341,7 +351,7 @@ def run(self):
                                         self.file_prefix_p,
                                         write_aux = False,
                                         options = self.output_options,
-                                        write_p = True) 
+                                        write_p = True)
 
             write_aux = (self.write_aux_always or self.write_aux_init)
             self.solution.write(frame,self.outdir,
@@ -374,7 +384,7 @@ def run(self):
                                             self.file_prefix_p,
                                             write_aux = False, 
                                             options = self.output_options,
-                                            write_p = True) 
+                                            write_p = True)
 
                 self.solution.write(frame,self.outdir,
                                             self.output_format,