multiple grid geometry outputs

gprMax/geometry_outputs.py

@@ -388,3 +388,296 @@ class GeometryObjects(object):
                                 dispersionstr += '{:g} {:g} '.format(material.tau[pole], material.alpha[pole])
                         dispersionstr += material.ID
                         fmaterials.write(dispersionstr + '\n')
+
+class GeometryViewFineMultiGrid:
+    """Geometry view for all grids in the simulation."""
+
+    """"Slicing is not supported by this class :( - only the full extent of the grids
+    are output. The subgrids are output without the non-working regions If you
+    require domainslicing GeometryView seperately for each grid you require and
+    view them at once in Paraview."""
+
+    if sys.byteorder == 'little':
+        byteorder = 'LittleEndian'
+    else:
+        byteorder = 'BigEndian'
+
+    def __init__(self, xs, ys, zs, xf, yf, zf, dx, dy, dz, filename, fileext, G):
+        """
+        Args:
+            xs, xf, ys, yf, zs, zf (int): Extent of the volume in cells.
+            dx, dy, dz (int): Spatial discretisation in cells.
+            filename (str): Filename to save to.
+            fileext (str): File extension of VTK file - either '.vti' for a per cell
+                    geometry view, or '.vtp' for a per cell edge geometry view.
+        """
+        self.G = G
+        self.nx = G.nx
+        self.ny = G.ny
+        self.nz = G.nz
+        self.basefilename = filename
+        self.fileext = '.vtp'
+        self.sg_views = []
+
+        self.additional_lines = 0
+        self.additional_points = 0
+
+        for sg in G.subgrids:
+            # create an object to contain data relevant to the geometry processing
+            sg_gv = SubgridGeometryView(sg)
+            self.sg_views.append(sg_gv)
+            # total additional lines required for subgrid
+            self.additional_lines += sg_gv.n_total_lines
+            # total additional points required for subgrid
+
+            self.additional_points += sg_gv.n_total_points
+
+        self.vtk_numpoints = self.additional_points + (self.nx + 1) * (self.ny + 1) * (self.nz + 1)
+        self.vtk_numpoint_components = 3
+        self.vtk_numlines = self.additional_lines + 2 * self.nx * self.ny + 2 * self.ny * self.nz + 2 * self.nx * self.nz + 3 * self.nx * self.ny * self.nz + self.nx + self.ny + self.nz
+        self.vtk_numline_components = 2
+        self.vtk_connectivity_offset = round_value(int((self.vtk_numpoints * self.vtk_numpoint_components * np.dtype(np.float32).itemsize) + np.dtype(np.uint32).itemsize))
+        self.vtk_offsets_offset = round_value(int(self.vtk_connectivity_offset + (self.vtk_numlines * self.vtk_numline_components * np.dtype(np.uint32).itemsize) + np.dtype(np.uint32).itemsize))
+        self.vtk_materials_offset = round_value(int(self.vtk_offsets_offset + (self.vtk_numlines * np.dtype(np.uint32).itemsize) + np.dtype(np.uint32).itemsize))
+        self.datawritesize = np.dtype(np.float32).itemsize * self.vtk_numpoints * self.vtk_numpoint_components + np.dtype(np.uint32).itemsize * self.vtk_numlines * self.vtk_numline_components + np.dtype(np.uint32).itemsize * self.vtk_numlines + np.dtype(np.uint32).itemsize * self.vtk_numlines
+
+    def set_filename(self, appendmodelnumber):
+        """
+        Construct filename from user-supplied name and model run number.
+
+        Args:
+            appendmodelnumber (str): Text to append to filename.
+        """
+
+        self.filename = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(config.general['inputfilepath'])), self.basefilename + appendmodelnumber))
+        self.filename += self.fileext
+
+    def write_vtk(self, *args):
+        """
+        Writes the geometry information to a VTK file. Either ImageData (.vti) for a
+        per-cell geometry view, or PolygonalData (.vtp) for a per-cell-edge geometry view.
+            N.B. No Python 3 support for VTK at time of writing (03/2015)
+        Args:
+            G (class): Grid class instance - holds essential parameters describing the model.
+        """
+
+        G = self.G
+
+        with open(self.filename, 'wb') as f:
+
+            # refine parameters for subgrid
+
+            f.write('<?xml version="1.0"?>\n'.encode('utf-8'))
+            f.write('<VTKFile type="PolyData" version="1.0" byte_order="{}">\n'.format(GeometryViewFineMultiGrid.byteorder).encode('utf-8'))
+            f.write('<PolyData>\n<Piece NumberOfPoints="{}" NumberOfVerts="0" NumberOfLines="{}" NumberOfStrips="0" NumberOfPolys="0">\n'.format(self.vtk_numpoints, self.vtk_numlines).encode('utf-8'))
+
+            f.write('<Points>\n<DataArray type="Float32" NumberOfComponents="3" format="appended" offset="0" />\n</Points>\n'.encode('utf-8'))
+            f.write('<Lines>\n<DataArray type="UInt32" Name="connectivity" format="appended" offset="{}" />\n'.format(self.vtk_connectivity_offset).encode('utf-8'))
+            f.write('<DataArray type="UInt32" Name="offsets" format="appended" offset="{}" />\n</Lines>\n'.format(self.vtk_offsets_offset).encode('utf-8'))
+            f.write('<CellData Scalars="Material">\n'.encode('utf-8'))
+            f.write('<DataArray type="UInt32" Name="Material" format="appended" offset="{}" />\n'.format(self.vtk_materials_offset).encode('utf-8'))
+            f.write('</CellData>\n'.encode('utf-8'))
+
+            f.write('</Piece>\n</PolyData>\n<AppendedData encoding="raw">\n_'.encode('utf-8'))
+
+            # Write points
+            print('writing points main grid')
+            datasize = np.dtype(np.float32).itemsize * self.vtk_numpoints * self.vtk_numpoint_components
+            f.write(pack('I', datasize))
+            for i in range(0, G.nx + 1):
+                for j in range(0, G.ny + 1):
+                    for k in range(0, self.G.nz + 1):
+                        f.write(pack('fff', i * G.dx, j * G.dy, k * G.dz))
+
+            for sg_v in self.sg_views:
+                print('writing points subgrid')
+                sg_v.write_points(f, G)
+
+            n_x_lines = self.nx * (self.ny + 1) * (self.nz + 1)
+            x_lines = np.zeros((n_x_lines, 2), dtype=np.uint32)
+            x_materials = np.zeros((n_x_lines), dtype=np.uint32)
+
+            n_y_lines = self.ny * (self.nx + 1) * (self.nz + 1)
+            y_lines = np.zeros((n_y_lines, 2), dtype=np.uint32)
+            y_materials = np.zeros((n_y_lines), dtype=np.uint32)
+
+            n_z_lines = self.nz * (self.nx + 1) * (self.ny + 1)
+            z_lines = np.zeros((n_z_lines, 2), dtype=np.uint32)
+            z_materials = np.zeros((n_z_lines), dtype=np.uint32)
+
+            print('calculate connectivity main grid')
+            label = 0
+            counter_x = 0
+            counter_y = 0
+            counter_z = 0
+            for i in range(self.nx + 1):
+                for j in range(self.ny + 1):
+                    for k in range(self.nz + 1):
+
+                        if i < self.nx:
+                            # x connectivity
+                            label_x = label + (self.ny + 1) * (self.nz + 1)
+                            x_lines[counter_x][0] = label
+                            x_lines[counter_x][1] = label_x
+                            # material for the line
+                            x_materials[counter_x] = G.ID[0, i, j, k]
+                            counter_x += 1
+                        if j < self.ny:
+                            label_y = label + self.nz + 1
+                            y_lines[counter_y][0] = label
+                            y_lines[counter_y][1] = label_y
+                            y_materials[counter_y] = G.ID[1, i, j, k]
+                            counter_y += 1
+                        if k < self.nz:
+                            label_z = label + 1
+                            z_lines[counter_z][0] = label
+                            z_lines[counter_z][1] = label_z
+                            z_materials[counter_z] = G.ID[2, i, j, k]
+                            counter_z += 1
+
+                        label = label + 1
+
+            print('calculate connectivity subgrids')
+            for sg_v in self.sg_views:
+                sg_v.populate_connectivity_and_materials(label)
+                # use the last subgrids label for the next view
+                label = sg_v.label
+
+            datasize = np.dtype(np.uint32).itemsize * self.vtk_numlines * self.vtk_numline_components
+            f.write(pack('I', datasize))
+
+            f.write(x_lines.tostring())
+            for sg_v in self.sg_views:
+                f.write(sg_v.x_s_lines.tostring())
+            f.write(y_lines.tostring())
+            for sg_v in self.sg_views:
+                f.write(sg_v.y_s_lines.tostring())
+            f.write(z_lines.tostring())
+            for sg_v in self.sg_views:
+                f.write(sg_v.z_s_lines.tostring())
+
+            # Write cell type (line) offsets
+            vtk_cell_pts = 2
+            datasize = np.dtype(np.uint32).itemsize * self.vtk_numlines
+            f.write(pack('I', datasize))
+            for vtk_offsets in range(vtk_cell_pts, (self.vtk_numline_components * self.vtk_numlines) + vtk_cell_pts, vtk_cell_pts):
+                f.write(pack('I', vtk_offsets))
+
+            datasize = np.dtype(np.uint32).itemsize * self.vtk_numlines
+            f.write(pack('I', datasize))
+
+            f.write(x_materials.tostring())
+            for sg_v in self.sg_views:
+                f.write(sg_v.x_s_materials.tostring())
+            f.write(y_materials.tostring())
+            for sg_v in self.sg_views:
+                f.write(sg_v.y_s_materials.tostring())
+            f.write(z_materials.tostring())
+            for sg_v in self.sg_views:
+                f.write(sg_v.z_s_materials.tostring())
+
+            f.write('\n</AppendedData>\n</VTKFile>'.encode('utf-8'))
+
+            #self.write_gprmax_info(f, G, materialsonly=True)
+
+    def write_gprmax_info(self, f, G, materialsonly=False):
+        """
+        Writes gprMax specific information relating material, source,
+        and receiver names to numeric identifiers.
+        Args:
+            f (filehandle): VTK file.
+            G (class): Grid class instance - holds essential parameters describing the model.
+            materialsonly (boolean): Only write information on materials
+        """
+
+        f.write('\n\n<gprMax>\n'.encode('utf-8'))
+        for material in G.materials:
+            f.write('<Material name="{}">{}</Material>\n'.format(material.ID, material.numID).encode('utf-8'))
+        if not materialsonly:
+            f.write('<PML name="PML boundary region">1</PML>\n'.encode('utf-8'))
+            for index, src in enumerate(G.hertziandipoles + G.magneticdipoles + G.voltagesources + G.transmissionlines):
+                f.write('<Sources name="{}">{}</Sources>\n'.format(src.ID, index + 2).encode('utf-8'))
+            for index, rx in enumerate(G.rxs):
+                f.write('<Receivers name="{}">{}</Receivers>\n'.format(rx.ID, index + 1).encode('utf-8'))
+        f.write('</gprMax>\n'.encode('utf-8'))
+        return None
+
+
+class SubgridGeometryView:
+    def __init__(self, sg):
+
+        self.sg = sg
+        # n component lines in each direction required for subgrid in the working region
+        n_sx_lines = sg.nwx * (sg.nwy + 1) * (sg.nwz + 1)
+        n_sy_lines = sg.nwy * (sg.nwx + 1) * (sg.nwz + 1)
+        n_sz_lines = sg.nwz * (sg.nwx + 1) * (sg.nwy + 1)
+
+        n_total_lines = n_sx_lines + n_sy_lines + n_sz_lines
+        self.n_total_lines = n_total_lines.astype(np.int32)
+
+        # n points in the the working region
+        n_total_points = (sg.nwx + 1) * (sg.nwy + 1) * (sg.nwz + 1)
+        self.n_total_points = n_total_points.astype(np.int32)
+
+        # connectivity array. 2 labels form an x component connection
+        self.x_s_lines = np.zeros((n_sx_lines, 2), dtype=np.uint32)
+        # material array. Each index contains a material index
+        self.x_s_materials = np.zeros((n_sx_lines), dtype=np.uint32)
+
+        self.y_s_lines = np.zeros((n_sy_lines, 2), dtype=np.uint32)
+        self.y_s_materials = np.zeros((n_sy_lines), dtype=np.uint32)
+
+        self.z_s_lines = np.zeros((n_sz_lines, 2), dtype=np.uint32)
+        self.z_s_materials = np.zeros((n_sz_lines), dtype=np.uint32)
+
+        self.label = 0
+
+    def write_points(self, f, G):
+        sg = self.sg
+        for i in range(sg.i0, sg.i0 + sg.nwx + 1):
+            for j in range(sg.j0, sg.j0 + sg.nwy + 1):
+                for k in range(sg.k0, sg.k0 + sg.nwz + 1):
+                    p_x = (sg.i0 * G.dx) + ((i - sg.i0) * sg.dx)
+                    p_y = (sg.j0 * G.dy) + ((j - sg.j0) * sg.dy)
+                    p_z = (sg.k0 * G.dz) + ((k - sg.k0) * sg.dz)
+                    f.write(pack('fff', p_x, p_y, p_z))
+
+    def populate_connectivity_and_materials(self, label):
+        """Label is the starting label. 0 if no other grids are present but +1 the last label used
+        for a multigrid view"""
+        sg = self.sg
+        self.label = label
+
+        # counters to to index numpy arrays
+        counter_x = 0
+        counter_y = 0
+        counter_z = 0
+
+        for i in range(sg.nwx + 1):
+            for j in range(sg.nwy + 1):
+                for k in range(sg.nwz + 1):
+                    i_s = i + sg.n_boundary_cells_x
+                    j_s = j + sg.n_boundary_cells_y
+                    k_s = k + sg.n_boundary_cells_z
+                    if i < sg.nwx:
+                        # x connectivity
+                        label_x = self.label + (sg.nwy + 1) * (sg.nwz + 1)
+                        self.x_s_lines[counter_x][0] = self.label
+                        self.x_s_lines[counter_x][1] = label_x
+                        # material for the line
+                        self.x_s_materials[counter_x] = sg.ID[0, i_s, j_s, k_s]
+                        counter_x += 1
+                    if j < sg.nwy:
+                        label_y = self.label + sg.nwz + 1
+                        self.y_s_lines[counter_y][0] = self.label
+                        self.y_s_lines[counter_y][1] = label_y
+                        self.y_s_materials[counter_y] = sg.ID[1, i_s, j_s, k_s]
+                        counter_y += 1
+                    if k < sg.nwz:
+                        label_z = self.label + 1
+                        self.z_s_lines[counter_z][0] = self.label
+                        self.z_s_lines[counter_z][1] = label_z
+                        self.z_s_materials[counter_z] = sg.ID[2, i_s, j_s, k_s]
+                        counter_z += 1
+
+                    self.label = self.label + 1