Add file write options

src/ambit_fe/fluid/fluid_main.py

@@ -1261,15 +1261,15 @@ def solve_initial_prestress(self):
 
         if self.pb.prestress_initial_only:
             # it may be convenient to write the prestress displacement field to a file for later read-in
-            self.pb.io.writefunction(self.pb.uf_pre, self.pb.io.output_path_pre+'/results_'+self.pb.pbase.simname+'_fluiddisplacement_pre.txt')
+            self.pb.io.writefunction(self.pb.uf_pre, self.pb.io.output_path_pre+'/results_'+self.pb.pbase.simname+'_fluiddisplacement_pre')
             if bool(self.pb.io.duplicate_mesh_domains):
                 m=0
                 for j in self.pb.io.duplicate_mesh_domains:
                      # TODO: Might not work for duplicate mesh, since we do not have the input node indices (do we...?)
-                    self.pb.io.writefunction(self.pb.p_[m], self.pb.io.output_path_pre+'/results_'+self.pb.pbase.simname+'_pressure'+str(m+1)+'_pre.txt')
+                    self.pb.io.writefunction(self.pb.p_[m], self.pb.io.output_path_pre+'/results_'+self.pb.pbase.simname+'_pressure'+str(m+1)+'_pre')
                     m+=1
             else:
-                self.pb.io.writefunction(self.pb.p_[0], self.pb.io.output_path_pre+'/results_'+self.pb.pbase.simname+'_pressure_pre.txt')
+                self.pb.io.writefunction(self.pb.p_[0], self.pb.io.output_path_pre+'/results_'+self.pb.pbase.simname+'_pressure_pre')
             utilities.print_status("Prestress only done. To resume, set file path(s) in 'prestress_from_file' and read in uf_pre.", self.pb.comm)
             os._exit(0)
 

src/ambit_fe/ioroutines.py

@@ -331,7 +331,7 @@ def readfunction(self, f, datafile, normalize=False, filetype='id_val'):
     # (e.g. linear pressure field defined on a quadratic input mesh), but not the other way
     # around (since we would not have all input node ids for a higher-order mesh defined using
     # a lower order input mesh)
-    def writefunction(self, f, filenm):
+    def writefunction(self, f, filenm, filetype='id_val'):
 
         # non-ghosted index map and input global node indices
         im_no_ghosts = f.function_space.dofmap.index_map.local_to_global(np.arange(f.function_space.dofmap.index_map.size_local, dtype=np.int32)).tolist()
@@ -366,12 +366,30 @@ def writefunction(self, f, filenm):
         igi_flat_sorted = sorted(igi_flat)
 
         # write to file
-        if self.comm.rank==0:
-            f = open(filenm, 'wt')
-            for i in igi_flat_sorted:
-                ind = np.where(igi_flat_array == i)[0][0]
-                f.write(str(i) + ' ' + ' '.join(map(str, vec_sq[bs*ind:bs*(ind+1)])) + '\n')
-            f.close()
+        if filetype=='id_val':
+            if self.comm.rank==0:
+                f = open(filenm+'.txt', 'wt')
+                for i in igi_flat_sorted:
+                    ind = np.where(igi_flat_array == i)[0][0]
+                    f.write(str(i) + ' ' + ' '.join(map(str, vec_sq[bs*ind:bs*(ind+1)])) + '\n')
+                f.close()
+        elif filetype=='val':
+            if self.comm.rank==0:
+                f = open(filenm+'.txt', 'wt')
+                for i in igi_flat_sorted:
+                    ind = np.where(igi_flat_array == i)[0][0]
+                    f.write(' '.join(map(str, vec_sq[bs*ind:bs*(ind+1)])) + '\n')
+                f.close()
+        elif filetype=='cheart': # CHeart .D file
+            if self.comm.rank==0:
+                f = open(filenm+'.D', 'wt')
+                f.write(str(len(igi_flat_sorted)) + ' ' + str(bs) + '\n')
+                for i in igi_flat_sorted:
+                    ind = np.where(igi_flat_array == i)[0][0]
+                    f.write(' '.join(map(str, vec_sq[bs*ind:bs*(ind+1)])) + '\n')
+                f.close()
+        else:
+            raise ValueError("Unknown filetype!")
 
         self.comm.Barrier()
 
@@ -763,7 +781,7 @@ def writecheckpoint(self, pb, N):
                 key.vector.view(viewer)
                 viewer.destroy()
             elif self.restart_io_type=='rawtxt': # only working for nodal fields!
-                self.writefunction(key, self.output_path+'/checkpoint_'+pb.pbase.simname+'_'+pb.problem_physics+'_'+vecs_to_write[key]+'_'+str(N)+'.txt')
+                self.writefunction(key, self.output_path+'/checkpoint_'+pb.pbase.simname+'_'+pb.problem_physics+'_'+vecs_to_write[key]+'_'+str(N))
             else:
                 raise ValueError("Unknown restart_io_type!")
 
@@ -952,7 +970,7 @@ def writecheckpoint(self, pb, N):
                 key.vector.view(viewer)
                 viewer.destroy()
             elif self.restart_io_type=='rawtxt': # only working for nodal fields!
-                self.writefunction(key, self.output_path+'/checkpoint_'+pb.pbase.simname+'_'+pb.problem_physics+'_'+vecs_to_write[key]+'_'+str(N)+'.txt')
+                self.writefunction(key, self.output_path+'/checkpoint_'+pb.pbase.simname+'_'+pb.problem_physics+'_'+vecs_to_write[key]+'_'+str(N))
             else:
                 raise ValueError("Unknown restart_io_type!")
 
@@ -1064,7 +1082,7 @@ def writecheckpoint(self, pb, N):
                 key.vector.view(viewer)
                 viewer.destroy()
             elif self.restart_io_type=='rawtxt': # only working for nodal fields!
-                self.writefunction(key, self.output_path+'/checkpoint_'+pb.pbase.simname+'_'+pb.problem_physics+'_'+vecs_to_write[key]+'_'+str(N)+'.txt')
+                self.writefunction(key, self.output_path+'/checkpoint_'+pb.pbase.simname+'_'+pb.problem_physics+'_'+vecs_to_write[key]+'_'+str(N))
             else:
                 raise ValueError("Unknown restart_io_type!")
 
@@ -1143,7 +1161,7 @@ def write_restart(self, pb, N, force=False):
                     key.vector.view(viewer)
                     viewer.destroy()
                 elif self.restart_io_type=='rawtxt': # only working for nodal fields!
-                    self.writefunction(key, self.output_path+'/checkpoint_'+pb.pbase.simname+'_'+pb.problem_physics+'_'+vecs_to_write[key]+'_'+str(N)+'.txt')
+                    self.writefunction(key, self.output_path+'/checkpoint_'+pb.pbase.simname+'_'+pb.problem_physics+'_'+vecs_to_write[key]+'_'+str(N))
                 else:
                     raise ValueError("Unknown restart_io_type!")
 

src/ambit_fe/mor/mor_main.py

@@ -288,7 +288,7 @@ def write_modes(self):
                 podfunc.vector[self.ss:self.se] = self.Phi[self.ss:self.se, self.numredbasisvec_true*h+i]
                 outfile.write_function(podfunc)
                 # also as txt (id_val file) for efficient read-in later on...
-                self.pb.io.writefunction(podfunc, self.pb.io.output_path+'/results_'+self.pb.pbase.simname+'_PODmode_P'+str(h+1)+'_'+str(i+1)+'.txt')
+                self.pb.io.writefunction(podfunc, self.pb.io.output_path+'/results_'+self.pb.pbase.simname+'_PODmode_P'+str(h+1)+'_'+str(i+1), filetype=self.filetype)
 
 
     # read modes from files

src/ambit_fe/solid/solid_main.py

@@ -1232,9 +1232,9 @@ def solve_initial_prestress(self):
 
         if self.pb.prestress_initial_only:
             # it may be convenient to write the prestress displacement field to a file for later read-in
-            self.pb.io.writefunction(self.pb.u_pre, self.pb.io.output_path_pre+'/results_'+self.pb.pbase.simname+'_displacement_pre.txt')
+            self.pb.io.writefunction(self.pb.u_pre, self.pb.io.output_path_pre+'/results_'+self.pb.pbase.simname+'_displacement_pre')
             if self.pb.incompressible_2field:
-                self.pb.io.writefunction(self.pb.p, self.pb.io.output_path_pre+'/results_'+self.pb.pbase.simname+'_pressure_pre.txt')
+                self.pb.io.writefunction(self.pb.p, self.pb.io.output_path_pre+'/results_'+self.pb.pbase.simname+'_pressure_pre')
             utilities.print_status("Prestress only done. To resume, set file path(s) in 'prestress_from_file' and read in u_pre.", self.pb.pbase.comm)
             os._exit(0)