"""
Part of comet/pyhed/IO
"""
# COMET - COupled Magnetic resonance Electrical resistivity Tomography
# Copyright (C) 2019 Nico Skibbe
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from __future__ import print_function
import numpy as np
import pygimli as pg
import time
from comet.pyhed.misc import GridtoVector
[docs]def add_vector_to_vtk(vtk, vector, vectorname, dtype_string='double'):
""" Appends a vector fields to an existing vtk file.
Parameters
----------
vtk: string
Path to vtk file, where the field is to be appended.
vector: np.ndarray
Real valued array of shape (3, n) n being either the number of
nodes or the number of cells.
vectorname: string
Name under which the array is to be identified in the vtk file.
dtype_string: string
Format string in the vtk file. Default 'double' is used for float
values.
"""
with open(vtk, 'a+') as file:
# seek end of file
file.seek(0, 2)
# print(file.tell)
file.write('VECTORS ' + vectorname + ' ' + dtype_string + '\n')
# file.write('LOOKUP_TABLE default\n')
for i in range(vector.shape[1]):
v_xyz = str(vector[0, i]) + ' ' +\
str(vector[1, i]) + ' ' +\
str(vector[2, i])
file.write(v_xyz + '\n')
[docs]def savefieldvtk(vtk_name, mesh, field, itype='mesh', components=False,
scalar=False, save=['real', 'imag'], field_name='field',
verbose=True):
""" Basic VTK export routine when it comes to complex vector fields on
unstructured meshes.
Parameters
----------
vtk_name: string
Path to the resulting vtk file.
mesh: string or pg.Mesh or np.ndarray
Pygimli Mesh instance or path to a mesh file. Alternatively a
bare numpy array containing coordinates or meshgrid ranges can be
used.
field: np.ndarray
Complex vector field of shape (3, n) with *n* corresponding either
to mesh cell count or node count.
itype: string [ 'mesh' ]
Defines input type of *mesh*. Possible choices are 'mesh' for
pg.Mesh (instance or file path), 'coords' for direct 3d coordinates
ranges to build a regular meshgrid, or 'grid' if input is a 3d
meshgrid.
components: boolean [ False ]
Separately saves the spatial components of the vector field for
debugging purposes.
scalar: boolean [ False ]
Input is a simple scalar field (e.g. potential).
save: list [ 'real', 'imag' ]
The vector is saved once for each entry in the list. Possible
choices are 'real' to save the real component, 'imag' to save
the imaginary component, 'aps' or 'amp' to save the amplitude, and
'phase' to save the phase component of the field. Only works with
vector fields.
field_name: string
Name under which the array is to be identified in the vtk file.
verbose: boolean [ True ]
Turn on verbose mode.
Returns
-------
True if succesful.
"""
tick = time.time()
from comet.pyhed.IO import checkDirectory
from comet.pyhed import log
checkDirectory(vtk_name, filename=True)
# try auto detect:
if isinstance(mesh, pg.Mesh):
itype == 'mesh'
# start import routine
if itype.lower() == 'coords':
x = mesh[0]
y = mesh[1]
z = mesh[2]
mesh = pg.createGrid(pg.Vector(x), pg.Vector(y), pg.Vector(z))
elif itype.lower() == 'grid':
x = mesh[0][:, 0, 0]
y = mesh[1][0, :, 0]
z = mesh[2][0, 0, :]
mesh = pg.createGrid(pg.Vector(x), pg.Vector(y), pg.Vector(z))
field = GridtoVector(field)
if len(field) == 3 and len(field[0]) != 3:
field_python = np.array(
(field[0], field[1], field[2]),
dtype=np.complex128 if np.iscomplexobj(field) else np.float64)
elif len(field[0]) == 3 and len(field) != 3:
field_python = np.array(
(field[:, 0], field[:, 1], field[:, 2]),
dtype=np.complex128 if np.iscomplexobj(field) else np.float64)
else:
raise Exception('Cannot deal with field of given shape: {}.'.format(
field.shape))
if field_python.shape[1] != mesh.cellCount() and\
field_python.shape[1] != mesh.nodeCount():
raise Exception(
'savefieldvtk: Invalid field shape, expect number of value to '
'fit either the nodecount ({}) or cellcount ({}) of the given '
'mesh. Got {} vectors instead. Abort.'
.format(mesh.nodeCount(), mesh.cellCount(), field_python.shape))
if scalar is True:
if isinstance(field_python[0], complex):
mesh.addData('{}_scalar_real'.format(field_name),
(pg.Vector(field_python.real)))
mesh.addData('{}_scalar_imag'.format(field_name),
(pg.Vector(field_python.imag)))
else:
mesh.addData('{}_scalar_field'.format(field_name),
(pg.Vector(field_python)))
log.info('save scalarfield to vtk ...')
mesh.exportVTK(vtk_name)
log.info('file name: {} ({:.3f} sec)'
.format(vtk_name, time.time() - tick))
return True
elif components is True:
mesh.addData('{}_x_real'.format(field_name),
(pg.Vector(np.real(field_python)[0])))
mesh.addData('{}_y_real'.format(field_name),
(pg.Vector(np.real(field_python)[1])))
mesh.addData('{}_z_real'.format(field_name),
(pg.Vector(np.real(field_python)[2])))
mesh.addData('{}_x_imag'.format(field_name),
(pg.Vector(np.imag(field_python)[0])))
mesh.addData('{}_y_imag'.format(field_name),
(pg.Vector(np.imag(field_python)[1])))
mesh.addData('{}_z_imag'.format(field_name),
(pg.Vector(np.imag(field_python)[2])))
log.info('save vectorfield to vtk ...')
mesh.exportVTK(vtk_name)
if field_python.shape[1] == mesh.cellCount():
log.info('savefieldvtk: Found cell based vector field. '
'Interpolate to nodes.')
field_python = fieldCell2Node(mesh, field_python)
elif field_python.shape[1] == mesh.nodeCount():
pass
else:
raise Exception('add field to vtk wrong shape: {}.'
.format(np.shape(field_python)))
if 'real' in save:
add_vector_to_vtk(vtk_name, np.real(field_python),
'{}_real'.format(field_name))
if 'imag' in save:
add_vector_to_vtk(vtk_name, np.imag(field_python),
'{}_imag'.format(field_name))
if 'abs' in save or 'amp' in save:
ABS = np.absolute(field_python) * np.sign(field_python.real)
add_vector_to_vtk(vtk_name, ABS, '{}_abs'.format(field_name))
if 'phase' in save:
add_vector_to_vtk(vtk_name, np.angle(field_python),
'{}_phase'.format(field_name))
log.info('file name: {} ({:.3f} sec)'.format(vtk_name, time.time() - tick))
return True
[docs]def fieldCell2Node(mesh, field):
shape_in = np.asarray(field).shape
# ensure shape = (n, 3), edge_case for shape (3, 3)
if shape_in[0] == 3 and shape_in[1] != 3:
field = field.T
if field.shape[0] != mesh.cellCount():
raise Exception('Error in fieldCell2Node, expect {} field '
'vectors for given mesh, got {}. Abort.'
.format(mesh.cellCount(), field.shape[0]))
is_complex = np.iscomplexobj(field)
xr = pg.meshtools.cellDataToNodeData(mesh, field[:, 0].real)
yr = pg.meshtools.cellDataToNodeData(mesh, field[:, 1].real)
zr = pg.meshtools.cellDataToNodeData(mesh, field[:, 2].real)
node_real = np.array((xr, yr, zr))
if is_complex:
# magnetization is always real
xi = pg.meshtools.cellDataToNodeData(mesh, field[:, 0].imag)
yi = pg.meshtools.cellDataToNodeData(mesh, field[:, 1].imag)
zi = pg.meshtools.cellDataToNodeData(mesh, field[:, 2].imag)
node_imag = np.array((xi, yi, zi))
ret = node_real + 1j * node_imag
else:
ret = node_real
return ret
# THE END