Skip to content

Commit

Permalink
add datamode and array type option to add array methods
Browse files Browse the repository at this point in the history
  • Loading branch information
joergbuchwald committed Sep 20, 2023
1 parent 1b619d6 commit 47cbe48
Showing 1 changed file with 24 additions and 23 deletions.
47 changes: 24 additions & 23 deletions VTUinterface/vtuIO.py
Original file line number Diff line number Diff line change
Expand Up @@ -473,7 +473,7 @@ def get_set_data(self, fieldname, pointsetarray=None, data_type="point", interpo
resp_array = np.array(resp_list)
return resp_array

def func_to_field(self, function, fieldname, ofilename=None, cell=False, writefile=True):
def func_to_field(self, function, fieldname, ofilename=None, cell=False, array_type=None, writefile=True, datamode=None):
"""
Add a field to the vtu file (which will be saved directly as "ofilename"
by providing a three argument function(x,y,z)
Expand All @@ -484,6 +484,7 @@ def func_to_field(self, function, fieldname, ofilename=None, cell=False, writefi
fieldname : `str`
ofilename : `str`
cell : `bool`
datamode : `str`
"""
if ofilename is None:
ofilename = self.filename
Expand All @@ -502,20 +503,17 @@ def func_to_field(self, function, fieldname, ofilename=None, cell=False, writefi
fieldarray[i] = function(points[i,0], points[i,1], 0.0)
else:
fieldarray[i] = function(points[i,0], points[i,1], points[i,2])
field_vtk = numpy_to_vtk(fieldarray)
field_vtk = numpy_to_vtk(fieldarray, array_type=None)
if cell is True:
r = self.cdata.AddArray(field_vtk)
self.cdata.GetArray(r).SetName(fieldname)
else:
r = self.pdata.AddArray(field_vtk)
self.pdata.GetArray(r).SetName(fieldname)
if writefile is True:
writer = vtk.vtkXMLUnstructuredGridWriter()
writer.SetFileName(ofilename)
writer.SetInputData(self.output)
writer.Write()
self.write(ofilename, datamode=datamode)

def func_to_m_dim_field(self, functionarray, fieldname, ofilename=None, cell=False, writefile=True):
def func_to_m_dim_field(self, functionarray, fieldname, ofilename=None, cell=False, array_type=None, writefile=True, datamode=None):
"""
Add a multidimensional field to the vtu file (which will be saved directly as "ofilename"
by providing am array of three argument functions.
Expand Down Expand Up @@ -555,18 +553,15 @@ def func_to_m_dim_field(self, functionarray, fieldname, ofilename=None, cell=Fal
points[i,0],
points[i,1],
points[i,2])
field_vtk = numpy_to_vtk(fieldarray)
field_vtk = numpy_to_vtk(fieldarray, array_type=array_type)
if cell is True:
r = self.cdata.AddArray(field_vtk)
self.cdata.GetArray(r).SetName(fieldname)
else:
r = self.pdata.AddArray(field_vtk)
self.pdata.GetArray(r).SetName(fieldname)
if writefile is True:
writer = vtk.vtkXMLUnstructuredGridWriter()
writer.SetFileName(ofilename)
writer.SetInputData(self.output)
writer.Write()
self.write(ofilename, datamode=datamode)

def point_data_to_cell_data(self, fieldname, ofilename=None, writefile=True):
"""
Expand All @@ -588,7 +583,7 @@ def point_data_to_cell_data(self, fieldname, ofilename=None, writefile=True):
cells_orig = self.output.GetCellData()
cells_orig.AddArray(array)
if writefile is True:
self.write(ofilename)
self.write(ofilename, datamode=datamode)

def delete_point_field(self, fieldnames=None, ofilename=None, writefile=True):
"""
Expand All @@ -613,7 +608,7 @@ def delete_point_field(self, fieldnames=None, ofilename=None, writefile=True):
else:
raise TypeError("Fieldnames has the wrong type. Please provide a list or string.")
if writefile is True:
self.write(ofilename)
self.write(ofilename, datamode=datamode)

def delete_cell_field(self, fieldnames=None, ofilename=None, writefile=True):
"""
Expand All @@ -638,7 +633,7 @@ def delete_cell_field(self, fieldnames=None, ofilename=None, writefile=True):
else:
raise TypeError("Fieldnames has the wrong type. Please provide a list or string.")
if writefile is True:
self.write(ofilename)
self.write(ofilename, datamode=datamode)

def delete_integration_point_field(self, fieldnames=None, ofilename=None, writefile=True):
"""
Expand All @@ -663,10 +658,10 @@ def delete_integration_point_field(self, fieldnames=None, ofilename=None, writef
else:
raise TypeError("Fieldnames has the wrong type. Please provide a list or string.")
if writefile is True:
self.write(ofilename)
self.write(ofilename, datamode=datamode)


def add_point_field(self, field, fieldname, ofilename, writefile=True):
def add_point_field(self, field, fieldname, ofilename, writefile=True, array_type=None, datamode=None):
"""
Write a field (numpy array of correct size)
to field "fieldname" as file "ofilename".
Expand All @@ -677,14 +672,16 @@ def add_point_field(self, field, fieldname, ofilename, writefile=True):
fieldname : `str`
ofilename : `str`
writefile : `bool`
array_type : `vtk array type`
datamode : `str`
"""
field_vtk = numpy_to_vtk(field)
field_vtk = numpy_to_vtk(field, array_type=array_type)
r = self.pdata.AddArray(field_vtk)
self.pdata.GetArray(r).SetName(fieldname)
if writefile is True:
self.write(ofilename)
self.write(ofilename, datamode=datamode)

def add_cell_field(self, field, fieldname, ofilename, writefile=True):
def add_cell_field(self, field, fieldname, ofilename, writefile=True, array_type=None, datamode=None):
"""
Write a field (numpy array of correct size)
to field "fieldname" as file "ofilename".
Expand All @@ -695,14 +692,16 @@ def add_cell_field(self, field, fieldname, ofilename, writefile=True):
fieldname : `str`
ofilename : `str`
writefile : `bool`
array_type : `vtk array type`
datamode : `str`
"""
field_vtk = numpy_to_vtk(field)
r = self.cdata.AddArray(field_vtk)
self.cdata.GetArray(r).SetName(fieldname)
if writefile is True:
self.write(ofilename)
self.write(ofilename, datamode=datamode)

def add_integration_point_field(self, field, fieldname, ofilename, writefile=True):
def add_integration_point_field(self, field, fieldname, ofilename, writefile=True, array_type=None, datamode=None):
"""
Write a field (numpy array of correct size)
to field "fieldname" as file "ofilename".
Expand All @@ -713,12 +712,14 @@ def add_integration_point_field(self, field, fieldname, ofilename, writefile=Tru
fieldname : `str`
ofilename : `str`
writefile : `bool`
array_type : `vtk array type`
datamode : `str`
"""
field_vtk = numpy_to_vtk(field)
r = self.ipdata.AddArray(field_vtk)
self.ipdata.GetArray(r).SetName(fieldname)
if writefile is True:
self.write(ofilename)
self.write(ofilename, datamode=datamode)

def write(self, filename, datamode=None):
"""
Expand Down

0 comments on commit 47cbe48

Please sign in to comment.