Initial commit from Polytechnique Montreal

1 files changed, 463 insertions, 0 deletions

diff --git a/PyGan/data/rDragView_proc/vtu_tools.py b/PyGan/data/rDragView_proc/vtu_tools.py
new file mode 100644
index 0000000..392c89a
--- /dev/null
+++ b/PyGan/data/rDragView_proc/vtu_tools.py
@@ -0,0 +1,463 @@
+#!/usr/bin/env python
+
+# Purpose:
+#  Functions related to creating vtu files.
+# 
+# Copyright:
+#  Copyright (C) 2025 Polytechnique Montreal
+#  This library is free software; you can redistribute it and/or modify it under the terms 
+#  of the GNU Lesser General Public License as published by the Free Software Foundation; 
+#  either version 2.1 of the License, or (at your option) any later version
+# 
+# Author(s): Atyab A. Calloo
+
+import numpy as np
+import vtk
+from math import sqrt
+
+
+def create_vtu(Geom, Mcro, Trck, Flux, test_name, verbose):
+   """
+   Uses Geom, Macr, Trck, Flux LCM objects to create the material and flux vtu
+   objects that can then be viewed in Paraview. 
+
+   Args:
+      Geom            : geometry LCM object
+      Macr            : macrolib LCM object
+      Trck            : tracking LCM object
+      Flux            : flux LCM object
+      test_name       : file name of the c2m procedure to be computed
+      verbose         : controls write output level
+   """
+
+   ### RECOVER STATEVECTORS FROM LCM OBJECTS
+   SV_geom = Geom["STATE-VECTOR"]
+   SV_mcro = Mcro["STATE-VECTOR"]
+   SV_trck = Trck["STATE-VECTOR"]
+   SV_flux = Flux["STATE-VECTOR"]
+   GV_mcro = Mcro["GROUP"]
+
+   ### RECOVER RELEVANT MATERIAL, TRACKING AND FLUX DATA
+   n_mat = SV_mcro[1]
+   n_reg = SV_trck[0]
+   n_unk = SV_trck[1]
+   n_dim = SV_trck[8]
+   spltl = SV_trck[25]
+   n_grp = SV_flux[0]
+   geom_type = SV_trck[5] # 5 - Car2D; 7 - Car3D; 8 - Hex2D; 9 - Hex3D
+   ### RECOVER OTHER DATA FROM LCM OBJECTS
+   mat_array  = Trck["MATCOD"]
+
+   ### INITIALIZE FLUX AND NUSIGF ARRAYS
+   flx_array = np.zeros((n_grp,n_unk), dtype='f')
+   nsf_mat_array = np.zeros((n_grp,n_mat), dtype='f')
+   nsf_reg_array = np.zeros((n_grp,n_reg), dtype='f')
+   ### PICK UP FLUX AND NUSIGF ARRAYS 
+   for i_g in range(n_grp):
+      flx_array[i_g,:] = Flux["FLUX"][i_g]
+      nsf_mat_array[i_g,:] = GV_mcro[i_g]["NUSIGF"]
+      ### CREATE NUSIGF ARRAY BY REGION
+      for i_reg in range(n_reg):
+         ind_mat = mat_array[i_reg]
+         if ind_mat == 0 : 
+            nsf_reg_array[i_g,i_reg] = 0
+         elif ind_mat > 0 : 
+            ind_mat=ind_mat-1
+            nsf_reg_array[i_g,i_reg] = nsf_mat_array[i_g,ind_mat]
+
+   ### REBUILD MESHES
+   lx, meshx, ly, meshy, lz, meshz = rebuild_meshes(Geom)
+
+   ### FIND HEXAGON CENTRES IF NEEDED
+   if geom_type==8 or geom_type==9:
+      side = Geom["SIDE"][0]
+      n_hex=int(n_reg/(lz*3*spltl**2))
+      n_ring = int((sqrt(  float((4*n_hex-1)/3)  )+1.)/2.)
+      hex_centres = compute_hex_centres(n_hex,n_ring,side)
+
+   ### CALL TO MAKE_VTU GEOM-DEPENDENT
+   nscat = SV_trck[6]
+   ielem = SV_trck[7]
+   if (geom_type == 5) or (geom_type == 7):
+      if verbose > 0: 
+         print('\n>>>>>>> GENERATING VTUs FOR ',test_name,'.')
+         print('-'*74 + '\n')
+      make_vtu_cartesian_geom(lx,meshx,ly,meshy,lz,meshz,mat_array,"material",test_name )
+      for i_g in range(n_grp):
+         make_vtu_cartesian_geom(lx,meshx,ly,meshy,lz,meshz,flx_array[i_g,:],"flux",test_name ,nscat=nscat,ielem=ielem,i_g=i_g)
+   elif (geom_type == 8) or (geom_type == 9):
+      if verbose > 0: 
+         print('\n>>>>>>> GENERATING VTUs FOR ',test_name,'.')
+         print('-'*74 + '\n')
+      make_vtu_hexagonal_geom(n_hex,side,hex_centres,mat_array,"material",test_name ,spltl=spltl,meshz=meshz,lz=lz)
+      for i_g in range(n_grp):
+         make_vtu_hexagonal_geom(n_hex,side,hex_centres,flx_array[i_g,:],"flux",test_name ,meshz=meshz,lz=lz,nscat=nscat,ielem=ielem,spltl=spltl,i_g=i_g)
+
+   print("\n>>>>>>> VTU CREATION COMPLETED!")
+   print('-'*74 + '\n')
+
+
+########################################################################
+def rebuild_meshes(Geom):
+   """
+   Uses Geom LCM objects to return number of mesh elements and mesh arrays along each 
+   axis, taking split into account.
+
+   Args:
+      Geom            : geometry LCM object
+   Returns:
+      lx              : number of mesh elements along x
+      meshx           : mesh array along x
+      ly              : number of mesh elements along y
+      meshy           : mesh array along y
+      lz              : number of mesh elements along z
+      meshz           : mesh array along z
+   """
+
+   SV_geom = Geom["STATE-VECTOR"]
+   geom_type = SV_geom[0] # 5 - Car2D; 7 - Car3D; 8 - Hex2D; 9 - Hex3D
+
+   mesh_dict = {}
+
+   ### FOR CARTESIAN GEOMETRIES
+   if geom_type==5 or geom_type==7:
+      lx    = SV_geom[2]
+      meshx = Geom["MESHX"]
+      spltx = Geom["SPLITX"]
+      lx, meshx = get_mesh_oneaxis(lx,meshx,spltx)
+      ly    = SV_geom[3]
+      meshy = Geom["MESHY"]
+      splty = Geom["SPLITY"]
+      ly, meshy = get_mesh_oneaxis(ly,meshy,splty)
+   else:
+      lx   =1
+      meshx=1
+      ly   =1
+      meshy=1
+
+   ### FOR 3D GEOMETRIES
+   if geom_type==7 or geom_type==9:
+      lz    = SV_geom[4]
+      meshz = Geom["MESHZ"]
+      spltz = Geom["SPLITZ"]
+      lz, meshz = get_mesh_oneaxis(lz,meshz,spltz)
+   else:
+      lz   =1
+      meshz=1
+      spltz=1
+
+   return lx, meshx, ly, meshy, lz, meshz
+
+
+########################################################################
+def get_mesh_oneaxis(lx,mesh,splt):
+   """
+   Compute number of mesh elements and mesh arrays taking split into account.
+
+   Args:
+      lx              : number of mesh elements along axis w/o splits
+      meshx           : mesh array along axis w/o splits
+      splt            : split value along axis
+   Returns:
+      lx              : number of mesh elements along axis w/ splits
+      meshx           : mesh array along axis w/ splits
+   """
+
+   tmp_mesh=np.empty(1, dtype='f') 
+   val = 0.
+   tmp_mesh[0] = val
+   for i in range(len(splt)):
+      step=(mesh[i+1]-mesh[i])/splt[i]
+      for j in range(splt[i]):
+         val += step
+         tmp_mesh = np.append(tmp_mesh,val)
+
+   lx=len(tmp_mesh)-1
+
+   return lx, tmp_mesh
+
+
+########################################################################
+def compute_hex_centres(n_hex,n_ring,side):
+   """
+   Compute centres of hexagons within hexagonal domain. Centre hexagon is zeroth
+   ring.
+   Args:
+      n_hex           : number of hexagons in domain
+      n_ring          : number of rings in domain; centre hexagon is zeroth ring
+      side            : length of one side of hexagon
+   Returns:
+      hex_centres     : 2D array with x and y coordinates of hexagon centres
+   """
+
+   hex_centres = np.zeros((2,n_hex), dtype='f')
+   ind_hex =0
+   hex_centres[0, 0] = 0.0
+   hex_centres[1, 0] = 0.0
+
+   ### LOOP OVER EACH RING OF HEXAGONS
+   for ind_ring in range(1, n_ring):
+      numhex_ring = ind_ring * 6
+      xcoord = ind_ring * 1.5 * side
+      ycoord = ind_ring * (sqrt(3) / 2) * side
+
+      hex_centres[0, ind_hex + 1] = xcoord
+      hex_centres[1, ind_hex + 1] = ycoord
+      ind_hex += 1
+
+      ind_dir = 1
+      ind_step = 0
+      numhex_perdir = ind_ring
+
+      ### LOOP OVER HEXAGONS IN CURRENT RING
+      ### For each of the six major axis directions, move in the direction
+      ### of that axis for numhex_perdir hexagons (which happens to be the
+      ### same as the ring index).
+      for i_hex in range(1, numhex_ring):
+         ind_step += 1
+         if ind_dir == 1:
+            xcoord -= 1.5 * side
+            ycoord += (sqrt(3) / 2) * side
+         elif ind_dir == 2:
+            xcoord -= 1.5 * side
+            ycoord -= (sqrt(3) / 2) * side
+         elif ind_dir == 3:
+            ycoord -= sqrt(3) * side
+         elif ind_dir == 4:
+            xcoord += 1.5 * side
+            ycoord -= (sqrt(3) / 2) * side
+         elif ind_dir == 5:
+            xcoord += 1.5 * side
+            ycoord += (sqrt(3) / 2) * side
+         elif ind_dir == 6:
+            ycoord += sqrt(3) * side
+
+         hex_centres[0, ind_hex + 1] = xcoord
+         hex_centres[1, ind_hex + 1] = ycoord
+         ind_hex += 1
+
+         if ind_step == numhex_perdir:
+            ind_dir += 1
+            ind_step = 0
+
+   return hex_centres
+
+
+########################################################################
+def writeXML(ugrid, file_name):
+   """
+   Save unstructured grid vtk object (ugrid) to file.
+   Args:
+      ugrid           : vtkUnstructuredGrid object
+      file_name       : file name under which to save ugrid
+   Returns:
+      None
+   """
+   writer = vtk.vtkXMLUnstructuredGridWriter()
+   writer.SetFileName(file_name)
+   writer.SetInputData(ugrid)
+   writer.Write()
+
+
+########################################################################
+def make_vtu_cartesian_geom(
+   lx,meshx,ly,meshy,lz,meshz,
+   data_np,
+   dataset_name,test_name,
+   **kwargs
+   ):
+   """
+   Build vtk unstructured grid object for 2D or 3D Cartesian mesh and populate with
+   given data, material number or flux.
+   Args:
+      lx              : number of mesh elements along x
+      meshx           : mesh array along x
+      ly              : number of mesh elements along y
+      meshy           : mesh array along y
+      lz              : number of mesh elements along z
+      meshz           : mesh array along z
+      data_np         : given data, material number or flux
+      dataset_name    : data name, usually 'material' or 'flux'
+      test_name       : file name of the computed c2m procedure in case of a prior
+                        dragon run or given name by user or default name
+   Returns:
+      None
+   """
+
+   ### ASSIGN OPTIONAL ARGS
+   nscat = kwargs.get("nscat", 1)
+   ielem = kwargs.get("ielem", 1)
+   i_g = kwargs.get("i_g", 0)
+   if dataset_name=="rate": 
+      nsf_reg=kwargs.get("nusigf", 1)
+
+   ### 2D CONSIDERATIONS
+   if lz==1: meshz = [0,0]
+   ndim = 2 if lz==1 else 3
+
+   ### INITIALISE VTK OBJECTS
+   points = vtk.vtkPoints()
+   element = vtk.vtkVoxel()
+   cells = vtk.vtkCellArray()
+   ugrid = vtk.vtkUnstructuredGrid()
+   ### INITIALISE SCALAR CELL DATA
+   DRAGONdataset = vtk.vtkFloatArray()
+   DRAGONdataset.SetName(dataset_name)
+   ugrid.GetCellData().SetScalars(DRAGONdataset)
+
+   ### COMPUTE CONSTANTS
+   flux_offset = 1 if dataset_name=="material" else nscat*(ielem)**ndim
+
+   i_cell=-1
+   ### LOOP OVER Z-AXIS MESH
+   for kk in range(lz):
+
+      ### LOOP OVER Y-AXIS MESH
+      for jj in range(ly):
+         ycoord = [meshy[jj], meshy[jj], meshy[jj+1], meshy[jj+1]]
+
+         ### LOOP OVER X-AXIS MESH
+         for ii in range(lx):
+            xcoord = [meshx[ii], meshx[ii+1], meshx[ii+1], meshx[ii]]
+
+            ### INSERT NEW POINTS AND CELLS (DO NOT MERGE THE TWO LOOPS BELOW!!!)
+            for i_crd in range(0,4):
+               points.InsertNextPoint( xcoord[i_crd],ycoord[i_crd], meshz[kk])
+            for i_crd in range(0,4):
+               points.InsertNextPoint( xcoord[i_crd],ycoord[i_crd], meshz[kk+1])
+            for i in range(0, 8):
+               i_cell += 1
+               element.GetPointIds().SetId(i, i_cell)
+            cells.InsertNextCell(element)
+
+            ugrid.SetPoints(points)
+            ugrid.SetCells(vtk.VTK_HEXAHEDRON, cells)
+
+            ### SET SCALAR CELL DATA
+            ind1 = (kk*lx*ly) + (jj*lx) + ii
+            ind2 = (kk*lx*ly)*flux_offset + (jj*lx)*flux_offset + ii*flux_offset
+            if dataset_name=="rate":
+               DRAGONdataset.InsertTuple1(ind1, nsf_reg[ind1]*data_np[ind2])
+            else:
+               DRAGONdataset.InsertTuple1(ind1, data_np[ind2])
+
+   print(test_name, dataset_name,' PROCESSED. SAVING ...')
+   filename=test_name+'_'+dataset_name+str(i_g)+".vtu"
+   ### SAVE VTU FILE
+   writeXML(ugrid, filename)
+   return ugrid
+
+
+########################################################################
+def make_vtu_hexagonal_geom(
+   n_hex,side,hex_centres,
+   data_np,
+   dataset_name,test_name,
+   **kwargs
+   ):
+   """
+   Build vtk unstructured grid object for 2D or 3D hexagon mesh and populate with
+   given data, material number or flux.
+   Args:
+      n_hex           : number of hexagons in domain
+      side            : length of one side of hexagon
+      hex_centres     : 2D array with x and y coordinates of hexagon centres
+      data_np         : given data, material number or flux
+      dataset_name    : data name, usually 'material' or 'flux'
+      test_name       : file name of the computed c2m procedure in case of a prior
+                        dragon run or given name by user or default name
+   Returns:
+      None
+   """
+   ### ASSIGN OPTIONAL ARGS
+   spltl = kwargs.get("spltl", 1)
+   lz = kwargs.get("lz", 1)
+   meshz = kwargs.get("meshz", 1)
+   nscat = kwargs.get("nscat", 1)
+   ielem = kwargs.get("ielem", 1)
+   i_g = kwargs.get("i_g", 0)
+
+   ### 2D CONSIDERATIONS
+   if lz==1: meshz = [0,0]
+   ndim = 2 if lz==1 else 3
+
+   ### INITIALISE VTK OBJECTS
+   points = vtk.vtkPoints()
+   lozenge = vtk.vtkVoxel()
+   cells = vtk.vtkCellArray()
+   ugrid = vtk.vtkUnstructuredGrid()
+   ### INITIALISE SCALAR CELL DATA
+   DRAGONdataset = vtk.vtkFloatArray()
+   DRAGONdataset.SetName(dataset_name)
+   ugrid.GetCellData().SetScalars(DRAGONdataset)
+
+   ### COMPUTE CONSTANTS
+   x_step=(1.0)*(side/spltl)
+   y_step=(sqrt(3)/2)*(side/spltl)
+   n_loz = 3*spltl**2
+   flux_offset = 1 if dataset_name=="material" else nscat*(ielem)**ndim
+
+   i_cell = -1
+   ### LOOP OVER Z-AXIS LAYERS
+   for i_z in range(0,lz):
+      ### LOOP OVER HEXAGONS
+      for i_hex in range(n_hex):
+         ### FETCH CENTRE OF EACH HEX
+         cx, cy = hex_centres[0, i_hex], hex_centres[1, i_hex]
+
+         ### LOOP OVER EACH HEX SECTION (EACH OF THE THREE MAIN LOZENGES)
+         for hex_section in range(3):
+            ### SET INITIAL VALUES OF XCOORD AND YCOORD
+            if hex_section == 0:
+               x_coords = [0.0, x_step*0.5, x_step, x_step*0.5]
+               y_coords = [0.0, -y_step, 0.0, y_step]
+            elif hex_section == 1:
+               x_coords = [-side, -side+x_step, -side+(x_step*1.5), -side+(x_step * 0.5)]
+               y_coords = [0.0, 0.0, y_step, y_step]
+            else:
+               x_coords = [-side/2, (-side/2)+x_step, (-side/2)+(x_step*0.5), (-side/2)-(x_step*0.5)]
+               y_coords = [-side*sqrt(3)/2, -side*sqrt(3)/2, (-side*sqrt(3)/2)+y_step, (-side*sqrt(3)/2)+y_step]
+
+            ### LOOP OVER Y-MESH INSIDE LOZENGE, INCREMENT XCOORD, YCOORD
+            for i_mshy in range(spltl):
+               if i_mshy != 0:
+                  if hex_section == 0:
+                     x_coords = x_coords - ((x_step/2)*(spltl-2))
+                     y_coords = y_coords + y_step*(spltl)
+                  elif hex_section == 1:
+                     x_coords = x_coords - ((x_step)*(spltl-1.5))
+                     y_coords = y_coords + y_step
+                  else:
+                     x_coords = x_coords - ((x_step)*(spltl-0.5))
+                     y_coords = y_coords + y_step
+
+               ### LOOP OVER X-MESH INSIDE LOZENGE, INCREMENT XCOORD, YCOORD
+               for i_mshx in range(spltl):
+                  if i_mshx != 0:
+                     x_coords = x_coords + (x_step/2) if hex_section == 0 else x_coords + (x_step)
+                     y_coords = y_coords - (y_step) if hex_section == 0 else y_coords
+
+                  ### INSERT NEW POINTS AND CELLS (DO NOT MERGE THE TWO LOOPS BELOW!!!)
+                  for i_crd in range(0,4):
+                     points.InsertNextPoint(x_coords[i_crd]+cx, y_coords[i_crd]+cy, meshz[i_z])
+                  for i_crd in range(0,4):
+                     points.InsertNextPoint(x_coords[i_crd]+cx, y_coords[i_crd]+cy, meshz[i_z+1])
+                  for i in range(0, 8):
+                     i_cell += 1
+                     lozenge.GetPointIds().SetId(i, i_cell)
+                  cells.InsertNextCell(lozenge)
+
+         ugrid.SetPoints(points)
+         ugrid.SetCells(vtk.VTK_HEXAHEDRON, cells)
+
+         ### SET SCALAR CELL DATA
+         for i_loz in range(0,n_loz):
+            ind1 = (i_z*n_hex*n_loz) + (i_hex*n_loz) + i_loz
+            ind2 = (i_z*n_hex*n_loz*flux_offset) + (i_hex*n_loz*flux_offset) + (i_loz*flux_offset)
+            DRAGONdataset.InsertTuple1(ind1, data_np[ind2])
+
+   print(test_name, dataset_name,' PROCESSED. SAVING ...')
+   filename=test_name+'_'+dataset_name+str(i_g)+".vtu"
+   ### SAVE VTU FILE
+   writeXML(ugrid, filename)
+   return ugrid