Tutorials

0. Three-dimensional elastic bar of square cross-section

3D (visco-)elastic waveguide example. The cross-section is a 2D square with free boundary conditions on its boundaries. Material: viscoelastic steel. The SAFE eigenproblem is solved with the varying parameter as the wavenumber (eigenvalues are then frequencies) or as the frequency (eigenvalues are then wavenumbers). Viscoelastic loss is included by introducing imaginary parts (negative) to wave celerities. Results are compared with Euler-Bernoulli analytical solutions for the lowest wavenumber or frequency parameter.

# This file is a tutorial for waveguicsx (*), whose inputs are
# the matrices K0, K1, K2, M and the excitation vector F.
# In the tutorial, K0, K1, K2 and M are finite element matrices generated by FEnicSX (**).
#  (*) waveguicsx is a python library for solving complex waveguide problems
#      Copyright (C) 2023-2024  Fabien Treyssede
#      waveguicsx is free software distributed under the GNU General Public License
#      (https://github.com/treyssede/waveguicsx)
# (**) FEniCSx is an open-source computing platform for solving partial differential equations
#      distributed under the GNU Lesser General Public License (https://fenicsproject.org/)

##################################
# 3D (visco-)elastic waveguide example\
# The cross-section is a 2D square with free boundary conditions on its boundaries\
# Material: viscoelastic steel\
# The waveguide FE formulation (SAFE) leads to the following eigenvalue problem:\
# $(\textbf{K}_1-\omega^2\textbf{M}+\text{i}k(\textbf{K}_2+\textbf{K}_2^\text{T})+k^2\textbf{K}_3)\textbf{U}=\textbf{0}$\
# This eigenproblem is solved with the varying parameter as the wavenumber (eigenvalues are then frequencies) or as the frequency (eigenvalues are then wavenumbers)\
# Viscoelastic loss is included by introducing imaginary parts (negative) to wave celerities\
# Results are compared with Euler-Bernoulli analytical solutions for the lowest wavenumber or frequency parameter

import dolfinx
import ufl
from mpi4py import MPI
from petsc4py import PETSc
from slepc4py import SLEPc
import numpy as np
import matplotlib.pyplot as plt
import pyvista

from waveguicsx.waveguide import Waveguide
#For proper use with a jupyter notebook, uncomment the following line:
#pyvista.set_jupyter_backend("static"); pyvista.start_xvfb() #try: "none", "static", "pythreejs", "ipyvtklink"...

##################################
# Input parameters
a = 2.7e-3 #square half-length (m)
N = 10 #number of finite elements along one half-side
rho, cs, cl = 7932, 3260, 5960 #core density (kg/m3), shear and longitudinal wave celerities (m/s)
kappas, kappal = 0.008, 0.003 #core shear and longitudinal bulk wave attenuations (Np/wavelength)
omega = 2*np.pi*np.linspace(500e3/1000, 500e3, num=50) #angular frequency range (rad/s)
wavenumber = omega/cs #wavenumber range (1/m, eigenvalues are frequency)
nev = 20 #number of eigenvalues

##################################
# Re-scaling
L0 = a #characteristic length
T0 = a/cs #characteristic time
M0 = rho*a**3#characteristic mass
a = a/L0
rho, cs, cl = rho/M0*L0**3, cs/L0*T0, cl/L0*T0
omega = omega*T0
wavenumber = wavenumber*L0
cs, cl = cs/(1+1j*kappas/2/np.pi), cl/(1+1j*kappal/2/np.pi) #complex celerities (core)

##################################
# Create mesh and finite elements (six-node triangles with three dofs per node for the three components of displacement)
mesh = dolfinx.mesh.create_rectangle(MPI.COMM_WORLD, [np.array([-a, -a]), np.array([a, a])], 
                               [2*N, 2*N], dolfinx.mesh.CellType.triangle)
element = ufl.VectorElement("CG", "triangle", 2, 3) #Lagrange element, triangle, quadratic "P2", 3D vector
V = dolfinx.fem.FunctionSpace(mesh, element)

##################################
# Create Material properties (isotropic)
def isotropic_law(rho, cs, cl):
    E, nu = rho*cs**2*(3*cl**2-4*cs**2)/(cl**2-cs**2), 0.5*(cl**2-2*cs**2)/(cl**2-cs**2)
    C11 = C22 = C33 = E/(1+nu)/(1-2*nu)*(1-nu)
    C12 = C13 = C23 = E/(1+nu)/(1-2*nu)*nu
    C44 = C55 = C66 = E/(1+nu)/2
    return ((C11,C12,C13,0,0,0), 
            (C12,C22,C23,0,0,0), 
            (C13,C23,C33,0,0,0), 
            (0,0,0,C44,0,0), 
            (0,0,0,0,C55,0), 
            (0,0,0,0,0,C66))
C = isotropic_law(rho, cs, cl)
C = dolfinx.fem.Constant(mesh, PETSc.ScalarType(C))

##################################
# Create free boundary conditions (or uncomment lines below for Dirichlet)
bcs = []
# Dirichlet test case:
#fdim = mesh.topology.dim - 1
#boundary_facets = dolfinx.mesh.locate_entities_boundary(mesh, dim=fdim, marker=lambda x: np.logical_or(np.isclose(x[0], -a),
#                                                                                                       np.isclose(x[0], +a)+
#                                                                                                       np.isclose(x[1], -a)+
#                                                                                                       np.isclose(x[1], +a)))
#boundary_dofs = dolfinx.fem.locate_dofs_topological(V=V, entity_dim=fdim, entities=boundary_facets)
#bcs = [dolfinx.fem.dirichletbc(value=np.zeros(3, dtype=PETSc.ScalarType), dofs=boundary_dofs, V=V)]

##################################
# Define variational problem (SAFE method)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
Lxy = lambda u: ufl.as_vector([u[0].dx(0), u[1].dx(1), 0, u[0].dx(1)+u[1].dx(0), u[2].dx(0), u[2].dx(1)])
Lz  = lambda u: ufl.as_vector([0, 0, u[2], 0, u[0], u[1]])
k0 = ufl.inner(C*Lxy(u), Lxy(v)) * ufl.dx
k0_form = dolfinx.fem.form(k0)
k1 = ufl.inner(C*Lz(u), Lxy(v)) * ufl.dx
k1_form = dolfinx.fem.form(k1)
k2 = ufl.inner(C*Lz(u), Lz(v)) * ufl.dx
k2_form = dolfinx.fem.form(k2)
m = rho*ufl.inner(u, v) * ufl.dx
mass_form = dolfinx.fem.form(m)

##################################
# Build PETSc matrices
M = dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0)
M.assemble()
K0 = dolfinx.fem.petsc.assemble_matrix(k0_form, bcs=bcs)
K0.assemble()
K1 = dolfinx.fem.petsc.assemble_matrix(k1_form, bcs=bcs, diagonal=0.0)
K1.assemble()
K2 = dolfinx.fem.petsc.assemble_matrix(k2_form, bcs=bcs, diagonal=0.0)
K2.assemble()

##################################
# Solve the eigenproblem with SLEPc\
# The parameter is k, the eigenvalue is omega**2
wg = Waveguide(MPI.COMM_WORLD, M, K0, K1, K2)
wg.set_parameters(wavenumber=wavenumber)
wg.solve(nev) #access to components with: wg.eigenvalues[ik][imode], wg.eigenvectors[ik][idof,imode]
wg.plot()
plt.show()

##################################
# Comparison with Euler-Bernoulli analytical solutions
k = wg.wavenumber[0]
print(f'Computed eigenvalues for the first wavenumber (k={k}):\n {np.around(wg.eigenvalues[0],decimals=6)}')
E, mu, I, Ip, S = rho*cs**2*(3*cl**2-4*cs**2)/(cl**2-cs**2), rho*cs**2, (2*a)**4/12, 0.1406*(2*a)**4, (2*a)**2 #0.1406 is for square section
print(f'Euler-Bernoulli beam solution (only accurate for low frequency):\n \
        Bending wave: {np.around(np.sqrt(E*I/rho/S*k**4),decimals=6)}\n \
        Torsional wave: {np.around(np.sqrt(mu*Ip/rho/(2*I)*k**2),decimals=6)}\n \
        Longitudinal wave: {np.around(np.sqrt(E/rho*k**2),decimals=6)}')

##################################
# Solve the eigenproblem with SLEPc\
# The parameter is omega, the eigenvalue is k
wg = Waveguide(MPI.COMM_WORLD, M, K0, K1, K2)
wg.set_parameters(omega=omega)
wg.solve(nev) #access to components with: wg.eigenvalues[ik][imode], wg.eigenvectors[ik][idof,imode]
wg.plot()
plt.show() #blocking

##################################
# Comparison with Euler-Bernoulli analytical solutions
w = wg.omega[0]
print(f'Computed eigenvalues for the first frequency (omega={w}):\n {np.around(wg.eigenvalues[0],decimals=6)}')
print(f'Euler-Bernoulli beam solution (only accurate for low frequency):\n \
        Bending wave: {np.around(np.sqrt(np.sqrt(rho*S/E/I*w**2)),decimals=6)}\n \
        Torsional wave: {np.around(np.sqrt(rho*2*I/mu/Ip*w**2),decimals=6)}\n \
        Longitudinal wave: {np.around(np.sqrt(rho/E*w**2),decimals=6)}')

##################################
# Mesh visualization
grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(mesh, mesh.topology.dim))
plotter = pyvista.Plotter()
plotter.add_mesh(grid, show_edges=True)
plotter.view_xy()
plotter.show()

##################################
# Mode shape visualization
ik, imode = 0, 5 #parameter index, mode index to visualize
vec = wg.eigenvectors[ik].getColumnVector(imode)
u_grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V))
u_grid["u"] = np.array(vec).real.reshape(int(np.array(vec).size/V.element.value_shape), int(V.element.value_shape)) #V.element.value_shape is equal to 3
u_plotter = pyvista.Plotter()
u_plotter.add_mesh(grid, style="wireframe", color="k") #FE mesh
u_plotter.add_mesh(u_grid.warp_by_vector("u", factor=2.0), opacity=0.8, show_scalar_bar=True, show_edges=False) #do not show edges of higher order elements with pyvista
u_plotter.show_axes()
u_plotter.show()


""

1. Three-dimensional elastic bar of square cross-section with parallelization

3D (visco-)elastic waveguide example. The cross-section is a 2D square with free boundary conditions on its 1D boundaries. Material: viscoelastic steel. Viscoelastic loss is included by introducing imaginary parts (negative) to wave celerities. In this example:

  • the parameter loop (here, the frequency loop) is distributed on all processes

  • FE mesh and matrices are built on each local process

# This file is a tutorial for waveguicsx (*), whose inputs are
# the matrices K0, K1, K2, M and the excitation vector F.
# In the tutorial, K0, K1, K2 and M are finite element matrices generated by FEnicSX (**).
#  (*) waveguicsx is a python library for solving complex waveguide problems
#      Copyright (C) 2023-2024  Fabien Treyssede
#      waveguicsx is free software distributed under the GNU General Public License
#      (https://github.com/treyssede/waveguicsx)
# (**) FEniCSx is an open-source computing platform for solving partial differential equations
#      distributed under the GNU Lesser General Public License (https://fenicsproject.org/)

##################################
# 3D (visco-)elastic waveguide example\
# The cross-section is a 2D square with free boundary conditions on its 1D boundaries\
# Material: viscoelastic steel\
# The waveguide FE formulation (SAFE) leads to the following eigenvalue problem:\
# $(\textbf{K}_1-\omega^2\textbf{M}+\text{i}k(\textbf{K}_2+\textbf{K}_2^\text{T})+k^2\textbf{K}_3)\textbf{U}=\textbf{0}$\
# Viscoelastic loss is included by introducing imaginary parts (negative) to wave celerities\
# In this example:
# - the parameter loop (here, the frequency loop) is distributed on all processes
# - FE mesh and matrices are built on each local process
# Reminder for an execution in parallel mode (e.g. 4 processes):
#  mpiexec -n 4 python3 Elastic_Waveguide_SquareBar3D_ParallelizedLoop.py

import dolfinx
import ufl
from mpi4py import MPI
from petsc4py import PETSc
from slepc4py import SLEPc
import numpy as np
import matplotlib.pyplot as plt
#import pyvista

from waveguicsx.waveguide import Waveguide
#For proper use with a notebook, uncomment the following line:
#pyvista.set_jupyter_backend("static"); pyvista.start_xvfb() #try: "none", "static", "pythreejs", "ipyvtklink"...

##################################
# Input parameters
a = 2.7e-3 #square half-length (m)
N = 10 #number of finite elements along one half-side
rho, cs, cl = 7932, 3260, 5960 #core density (kg/m3), shear and longitudinal wave celerities (m/s)
kappas, kappal = 0.008, 0.003 #core shear and longitudinal bulk wave attenuations (Np/wavelength)
omega = 2*np.pi*np.linspace(500e3/1000, 500e3, num=50) #angular frequency range (rad/s)
nev = 20 #number of eigenvalues

##################################
# Re-scaling
L0 = a #characteristic length
T0 = a/cs #characteristic time
M0 = rho*a**3#characteristic mass
a = a/L0
rho, cs, cl = rho/M0*L0**3, cs/L0*T0, cl/L0*T0
omega = omega*T0
cs, cl = cs/(1+1j*kappas/2/np.pi), cl/(1+1j*kappal/2/np.pi) #complex celerities (core)

##################################
# Create mesh and finite elements (six-node triangles with three dofs per node for the three components of displacement)
mesh = dolfinx.mesh.create_rectangle(MPI.COMM_SELF, [np.array([-a, -a]), np.array([a, a])], 
                               [2*N, 2*N], dolfinx.mesh.CellType.triangle) #MPI.COMM_SELF = FE mesh is built on each local process
element = ufl.VectorElement("CG", "triangle", 2, 3) #Lagrange element, triangle, quadratic "P2", 3D vector
V = dolfinx.fem.FunctionSpace(mesh, element)

##################################
# Create Material properties (isotropic)
def isotropic_law(rho, cs, cl):
    E, nu = rho*cs**2*(3*cl**2-4*cs**2)/(cl**2-cs**2), 0.5*(cl**2-2*cs**2)/(cl**2-cs**2)
    C11 = C22 = C33 = E/(1+nu)/(1-2*nu)*(1-nu)
    C12 = C13 = C23 = E/(1+nu)/(1-2*nu)*nu
    C44 = C55 = C66 = E/(1+nu)/2
    return ((C11,C12,C13,0,0,0), 
            (C12,C22,C23,0,0,0), 
            (C13,C23,C33,0,0,0), 
            (0,0,0,C44,0,0), 
            (0,0,0,0,C55,0), 
            (0,0,0,0,0,C66))
C = isotropic_law(rho, cs, cl)
C = dolfinx.fem.Constant(mesh, PETSc.ScalarType(C))

##################################
# Create free boundary conditions
bcs = []

##################################
# Define variational problem (SAFE method)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
Lxy = lambda u: ufl.as_vector([u[0].dx(0), u[1].dx(1), 0, u[0].dx(1)+u[1].dx(0), u[2].dx(0), u[2].dx(1)])
Lz  = lambda u: ufl.as_vector([0, 0, u[2], 0, u[0], u[1]])
k0 = ufl.inner(C*Lxy(u), Lxy(v)) * ufl.dx
k0_form = dolfinx.fem.form(k0)
k1 = ufl.inner(C*Lz(u), Lxy(v)) * ufl.dx
k1_form = dolfinx.fem.form(k1)
k2 = ufl.inner(C*Lz(u), Lz(v)) * ufl.dx
k2_form = dolfinx.fem.form(k2)
m = rho*ufl.inner(u, v) * ufl.dx
mass_form = dolfinx.fem.form(m)

##################################
# Build PETSc matrices
M = dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0)
M.assemble()
K0 = dolfinx.fem.petsc.assemble_matrix(k0_form, bcs=bcs)
K0.assemble()
K1 = dolfinx.fem.petsc.assemble_matrix(k1_form, bcs=bcs, diagonal=0.0)
K1.assemble()
K2 = dolfinx.fem.petsc.assemble_matrix(k2_form, bcs=bcs, diagonal=0.0)
K2.assemble()

##################################
# Solve the eigenproblem with SLEPc\
# The parameter is k, the eigenvalue is omega**2
# The parameter loop is parallelized
# Parallelization
comm = MPI.COMM_WORLD #use all processes for the loop
size = comm.Get_size()  #number of processors
rank = comm.Get_rank()  #returns the rank of the process that called it within comm_world
# Split the parameter range and scatter to all
if rank == 0: #define on rank 0 only
    param_split = np.array_split(omega, size) #split param in blocks of length size roughly
else:
    param_split = None
param_local = comm.scatter(param_split, root=0) #scatter 1 block per process
# Solve
wg = Waveguide(MPI.COMM_SELF, M, K0, K1, K2) #MPI.COMM_SELF = SLEPc will used FE matrices on each local process
wg.set_parameters(omega=param_local)
wg.solve(nev)
# Some post-processing
wg.compute_energy_velocity()
wg.compute_traveling_direction()
# Gather
wg.omega = comm.reduce([wg.omega], op=MPI.SUM, root=0) #reduce works for lists: brackets are necessary (wg.omega is not a list but a numpy array)
wg.eigenvalues = comm.reduce(wg.eigenvalues, op=MPI.SUM, root=0)
wg.energy_velocity = comm.reduce(wg.energy_velocity, op=MPI.SUM, root=0)
wg.traveling_direction = comm.reduce(wg.traveling_direction, op=MPI.SUM, root=0)
#wg.eigenvectors = comm.reduce(wg.eigenvectors, op=MPI.SUM, root=0) #don't do this line: reduce cannot pickle 'petsc4py.PETSc.Vec' objects (keep the mode shapes distributed on each processor rather than gather them)
# Plot results
if rank == 0:
    wg.omega = np.concatenate(wg.omega) #wg.omega is transformed to a numpy array for a proper use of wg.plot()
    wg.plot()
    wg.plot_energy_velocity(direction=+1)
    #plt.savefig("Elastic_Waveguide_Bar3D_ParallelizedLoop.svg")
    plt.show()

2. Three-dimensional elastic bar of square cross-section buried into a PML external medium

3D (visco-)elastic waveguide example. The cross-section is a 2D square buried into a PML external elastic medium. Material: viscoelastic steel into cement grout. Viscoelastic loss is included by introducing imaginary parts (negative) to wave celerities. The PML has a parabolic profile. Results are to be compared with Fig. 8 of paper: Treyssede, Journal of Computational Physics 314 (2016), 341–354.

# This file is a tutorial for waveguicsx (*), whose inputs are
# the matrices K0, K1, K2, M and the excitation vector F.
# In the tutorial, K0, K1, K2 and M are finite element matrices generated by FEnicSX (**).
#  (*) waveguicsx is a python library for solving complex waveguide problems
#      Copyright (C) 2023-2024  Fabien Treyssede
#      waveguicsx is free software distributed under the GNU General Public License
#      (https://github.com/treyssede/waveguicsx)
# (**) FEniCSx is an open-source computing platform for solving partial differential equations
#      distributed under the GNU Lesser General Public License (https://fenicsproject.org/)

##################################
# 3D (visco-)elastic waveguide example\
# The cross-section is a 2D square buried into a PML external elastic medium\
# Material: viscoelastic steel into cement grout\
# The waveguide FE formulation (SAFE) leads to the following eigenvalue problem:\
# $(\textbf{K}_1-\omega^2\textbf{M}+\text{i}k(\textbf{K}_2+\textbf{K}_2^\text{T})+k^2\textbf{K}_3)\textbf{U}=\textbf{0}$\
# Viscoelastic loss is included by introducing imaginary parts (negative) to wave celerities\
# The PML has a parabolic profile
# Results are to be compared with Fig. 8 of paper: Treyssede, Journal of Computational Physics 314 (2016), 341–354

import dolfinx
import ufl
from mpi4py import MPI
from petsc4py import PETSc
from slepc4py import SLEPc
import numpy as np
import matplotlib.pyplot as plt
import pyvista

from waveguicsx.waveguide import Waveguide
#For proper use with a notebook, uncomment the following line:
#pyvista.set_jupyter_backend("static"); pyvista.start_xvfb() #try: "none", "static", "pythreejs", "ipyvtklink"...

##################################
# Input parameters
a = 2.7e-3 #core half-length (m)
b = 1.5*a #half-length including PML (m)
N = 6 #24 #number of finite elements along one half-side
if not (a*N/b).is_integer():
    raise NotImplementedError("The ratio a/b*N must be an integer, please adjust N")
rho_core, cs_core, cl_core = 7932, 3260, 5960 #core density (kg/m3), shear and longitudinal wave celerities (m/s)
kappas_core, kappal_core = 0.008, 0.003 #core shear and longitudinal bulk wave attenuations (Np/wavelength)
rho_ext, cs_ext, cl_ext = 1600, 1700, 2810 #for the external medium
kappas_ext, kappal_ext = 0.100, 0.043 #for the external medium
alpha = 2+4j #average value of the absorbing function inside the PML
omega = 2*np.pi*np.linspace(0, 26e6/(a*1e3), num=200) #angular frequencies (rad/s)
nev = 10 #number of eigenvalues
target = lambda omega: omega/cl_core.real #target set at the longitudinal wavenumber of core to find L(0,n) modes

##################################
# Re-scaling
L0 = a #characteristic length
T0 = a/cs_core #characteristic time
M0 = rho_core*a**3#characteristic mass
a, b = a/L0, b/L0
rho_core, cs_core, cl_core = rho_core/M0*L0**3, cs_core/L0*T0, cl_core/L0*T0
rho_ext, cs_ext, cl_ext = rho_ext/M0*L0**3, cs_ext/L0*T0, cl_ext/L0*T0
omega = omega*T0
cs_core, cl_core = cs_core/(1+1j*kappas_core/2/np.pi), cl_core/(1+1j*kappal_core/2/np.pi) #complex celerities (core)
cs_ext, cl_ext = cs_ext/(1+1j*kappas_ext/2/np.pi), cl_ext/(1+1j*kappal_ext/2/np.pi) #complex celerities (exterior)

##################################
# Create mesh and finite elements (P2 quadrilaterals or order 8 GLL, with three dofs per node for the three components of displacement)
mesh = dolfinx.mesh.create_rectangle(MPI.COMM_WORLD, [np.array([-b, -b]), np.array([b, b])], 
                               [2*N, 2*N], dolfinx.mesh.CellType.quadrilateral)
#element = ufl.VectorElement("CG", "quadrilateral", 2, 3) #Lagrange element, quadrilateral, quadratic "P2", 3D vector
import basix
element = basix.create_element(basix.ElementFamily.P, basix.CellType.quadrilateral, 8,
                               basix.LagrangeVariant.gll_warped)
element = basix.ufl_wrapper.BasixElement(element)
element = basix.ufl_wrapper.VectorElement(element, 3)
V = dolfinx.fem.FunctionSpace(mesh, element)

##################################
# Create Material properties, discontinuous between core and exterior
def isotropic_law(rho, cs, cl):
    E, nu = rho*cs**2*(3*cl**2-4*cs**2)/(cl**2-cs**2), 0.5*(cl**2-2*cs**2)/(cl**2-cs**2)
    C11 = C22 = C33 = E/(1+nu)/(1-2*nu)*(1-nu)
    C12 = C13 = C23 = E/(1+nu)/(1-2*nu)*nu
    C44 = C55 = C66 = E/(1+nu)/2
    C = np.array([[C11,C12,C13,0,0,0], 
                     [C12,C22,C23,0,0,0], 
                     [C13,C23,C33,0,0,0], 
                     [0,0,0,C44,0,0], 
                     [0,0,0,0,C55,0], 
                     [0,0,0,0,0,C66]])
    C = C[np.triu_indices(6)] #the upper-triangle part of C (21 elements only)...
    C = np.append(C,np.zeros(15)).reshape(36,1) #... stored as a column vector completed with zeros up to 36 elements (error otherwise)
    return C
C_core = isotropic_law(rho_core, cs_core, cl_core)
C_ext = isotropic_law(rho_ext, cs_ext, cl_ext)
def eval_C(x):
    values = C_core * np.logical_and(abs(x[0])<=a, abs(x[1])<=a) \
           + C_ext * np.logical_or(abs(x[0])>=a, abs(x[1])>=a)
    return values
Q = dolfinx.fem.FunctionSpace(mesh, ufl.TensorElement("DG", "quadrilateral", 0, (6,6), symmetry=True)) #symmetry enables to store 21 elements instead of 36
C = dolfinx.fem.Function(Q)
C.interpolate(eval_C) #C.vector.getSize()) should be equal to number of cells x 21
# Same approach for density
def eval_rho(x):
    values = rho_core * np.logical_and(abs(x[0])<=a, abs(x[1])<=a) \
           + rho_ext * np.logical_or(abs(x[0])>=a, abs(x[1])>=a)
    return values
Q = dolfinx.fem.FunctionSpace(mesh, ("DG", 0))
rho = dolfinx.fem.Function(Q)
rho.interpolate(eval_rho)
# Vizualization
plotter = pyvista.Plotter(window_size=[600, 400])
grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(mesh, mesh.topology.dim))
grid.cell_data["Cij"] = C.x.array[0::21].real #index from 0 to 20, 0 being for C11...
#grid.cell_data["Cij"] = rho.x.array.real #for checking density
grid.set_active_scalars("Cij")
plotter.add_mesh(grid, show_edges=True, show_scalar_bar=True)
plotter.add_text('Re(Cij)', 'upper_edge', color='black', font_size=8)
plotter.view_xy()
plotter.show()

##################################
# Create Cartesian PML functions, continuous
def eval_gamma(x): #pml profile: continuous and parabolic
    values = [1+3*(alpha-1)*((abs(x[0])-a)/(b-a))**2*(abs(x[0])>=a), #gammax
              1+3*(alpha-1)*((abs(x[1])-a)/(b-a))**2*(abs(x[1])>=a)] #gammay
    return values
Q = dolfinx.fem.FunctionSpace(mesh, ufl.VectorElement("CG", "quadrilateral", 2, 2))
gamma = dolfinx.fem.Function(Q)
gamma.interpolate(eval_gamma)
# Vizualization
plotter = pyvista.Plotter(window_size=[800, 400], shape=(1,2))
gridx = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(Q))
gridx.point_data["gammax"] = gamma.x.array[0::2].imag
gridx.set_active_scalars("gammax")
plotter.subplot(0,0)
plotter.add_mesh(grid, style="wireframe", color="k") #FE mesh
plotter.add_mesh(gridx, opacity=0.8, show_scalar_bar=True, show_edges=False)
plotter.add_text('Im(gammax)', 'upper_edge', color='black', font_size=8)
plotter.view_xy()
gridy = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(Q))
gridy.point_data["gammay"] = gamma.x.array[1::2].imag
gridy.set_active_scalars("gammay")
plotter.subplot(0,1)
plotter.add_mesh(grid, style="wireframe", color="k") #FE mesh
plotter.add_mesh(gridy, opacity=0.8, show_scalar_bar=True, show_edges=False)
plotter.add_text('Im(gammay)', 'upper_edge', color='black', font_size=8)
plotter.view_xy()
plotter.show()

##################################
# Create free boundary conditions (or uncomment lines below for Dirichlet)
# bcs = []
# Dirichlet test case:
fdim = mesh.topology.dim - 1
boundary_facets = dolfinx.mesh.locate_entities_boundary(mesh, dim=fdim, marker=lambda x: np.logical_or(np.isclose(x[0], -b),
                                                                                                       np.isclose(x[0], +b)+
                                                                                                       np.isclose(x[1], -b)+
                                                                                                       np.isclose(x[1], +b)))
boundary_dofs = dolfinx.fem.locate_dofs_topological(V=V, entity_dim=fdim, entities=boundary_facets)
bcs = [dolfinx.fem.dirichletbc(value=np.zeros(3, dtype=PETSc.ScalarType), dofs=boundary_dofs, V=V)]

##################################
# Define variational problem (SAFE method)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
Lx = lambda u: ufl.as_vector([u[0].dx(0), 0, 0, u[1].dx(0), u[2].dx(0), 0])
Ly = lambda u: ufl.as_vector([0, u[1].dx(1), 0, u[0].dx(1), 0, u[2].dx(1)])
Lz  = lambda u: ufl.as_vector([0, 0, u[2], 0, u[0], u[1]])
gammax = gamma[0]
gammay = gamma[1]
k0 = (ufl.inner(C*(gammay/gammax*Lx(u)+Ly(u)), Lx(v)) + ufl.inner(C*(Lx(u)+gammax/gammay*Ly(u)), Ly(v))) * ufl.dx
k0_form = dolfinx.fem.form(k0)
k1 = (gammay*ufl.inner(C*Lz(u), Lx(v))+gammax*ufl.inner(C*Lz(u), Ly(v))) * ufl.dx
k1_form = dolfinx.fem.form(k1)
k2 = ufl.inner(C*Lz(u), Lz(v)) * gammax * gammay * ufl.dx
k2_form = dolfinx.fem.form(k2)
m = rho*ufl.inner(u, v) * gammax * gammay * ufl.dx
mass_form = dolfinx.fem.form(m)

##################################
# Build PETSc matrices
M = dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0)
M.assemble()
K0 = dolfinx.fem.petsc.assemble_matrix(k0_form, bcs=bcs)
K0.assemble()
K1 = dolfinx.fem.petsc.assemble_matrix(k1_form, bcs=bcs, diagonal=0.0)
K1.assemble()
K2 = dolfinx.fem.petsc.assemble_matrix(k2_form, bcs=bcs, diagonal=0.0)
K2.assemble()

##################################
# Solve the eigenproblem with SLEPc\
# The parameter is omega, the eigenvalue is k
wg = Waveguide(MPI.COMM_WORLD, M, K0, K1, K2)
wg.set_parameters(omega=omega)
wg.evp.setWhichEigenpairs(SLEPc.PEP.Which.TARGET_MAGNITUDE) #here, preferred to TARGET_IMAGINARY
wg.solve(nev, target=target)

##################################
# Plot dispersion curves\
# Results are to be compared with Fig. 8 of Treyssede, Journal of Computational Physics 314 (2016), 341–354
wg.plot_scaler["energy_velocity"] = L0/T0/1000 #units in m/ms
wg.plot_scaler["frequency"] = L0/T0/1000 #frequency units in MHz-mm
sc = wg.plot_energy_velocity()
sc.axes.set_xlim([0, 26])
sc.axes.set_ylim([0, 6])
sc.axes.set_xlabel('Frequency-halfwidth (MHz-mm)')
sc.axes.set_ylabel('Energy velocity (m/ms)')
wg.plot_scaler["attenuation"] = 8.686*1000 #units in dB-mm/m
sc = wg.plot_attenuation()
sc.axes.set_xlim([0, 26])
sc.axes.set_ylim([0, 2000])
sc.axes.set_xlabel('Frequency-halfwidth (MHz-mm)')
sc.axes.set_ylabel('Attenuation (dB-mm/m)')
plt.show()
#wg.plot_spectrum(index=0)
#plt.show()

##################################
# Mode shape visualization
ik, imode = 50, 1 #parameter index, mode index to visualize
vec = wg.eigenvectors[ik].getColumnVector(imode)
u_grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V))
u_grid["u"] = np.array(vec).real.reshape(int(np.array(vec).size/V.element.value_shape), int(V.element.value_shape)) #V.element.value_shape is equal to 3
u_plotter = pyvista.Plotter()
u_plotter.add_mesh(grid, style="wireframe", color="k") #FE mesh
u_plotter.add_mesh(u_grid.warp_by_vector("u", factor=2.0), opacity=0.8, show_scalar_bar=True, show_edges=False) #do not show edges of higher order elements with pyvista
u_plotter.show_axes()
u_plotter.show()

3. Three-dimensional elastic bar of square cross-section buried into a PML external medium using gmsh

3D (visco-)elastic waveguide example. The cross-section is a 2D square buried into a PML external elastic medium. Material: viscoelastic steel into cement grout. Viscoelastic loss is included by introducing imaginary parts (negative) to wave celerities. The PML has a parabolic profile. The FE mesh is built from Gmsh. Results are to be compared with Fig. 8 of paper: Treyssede, Journal of Computational Physics 314 (2016), 341–354.

# This file is a tutorial for waveguicsx (*), whose inputs are
# the matrices K0, K1, K2, M and the excitation vector F.
# In the tutorial, K0, K1, K2 and M are finite element matrices generated by FEnicSX (**).
#  (*) waveguicsx is a python library for solving complex waveguide problems
#      Copyright (C) 2023-2024  Fabien Treyssede
#      waveguicsx is free software distributed under the GNU General Public License
#      (https://github.com/treyssede/waveguicsx)
# (**) FEniCSx is an open-source computing platform for solving partial differential equations
#      distributed under the GNU Lesser General Public License (https://fenicsproject.org/)

##################################
# 3D (visco-)elastic waveguide example\
# The cross-section is a 2D square buried into a PML external elastic medium\
# Material: viscoelastic steel into cement grout\
# The waveguide FE formulation (SAFE) leads to the following eigenvalue problem:\
# $(\textbf{K}_1-\omega^2\textbf{M}+\text{i}k(\textbf{K}_2+\textbf{K}_2^\text{T})+k^2\textbf{K}_3)\textbf{U}=\textbf{0}$\
# Viscoelastic loss is included by introducing imaginary parts (negative) to wave celerities\
# The PML has a parabolic profile\
# The FE mesh is built from Gmsh\
# Results are to be compared with Fig. 8 of paper: Treyssede, Journal of Computational Physics 314 (2016), 341–354

import gmsh #full documentation entirely defined in the `gmsh.py' module
import dolfinx
import ufl
from mpi4py import MPI
from petsc4py import PETSc
from slepc4py import SLEPc
import numpy as np
import matplotlib.pyplot as plt
import pyvista

from waveguicsx.waveguide import Waveguide
#For proper use with a notebook, uncomment the following line:
#pyvista.set_jupyter_backend("static"); pyvista.start_xvfb() #try: "none", "static", "pythreejs", "ipyvtklink"...

##################################
# Input parameters
a = 2.7e-3 #core half-length (m)
b, le = 1.5*a, a/4 #half-length including PML (m), finite element characteristic length (m)
rho_core, cs_core, cl_core = 7932, 3260, 5960 #core density (kg/m3), shear and longitudinal wave celerities (m/s)
kappas_core, kappal_core = 0.008, 0.003 #core shear and longitudinal bulk wave attenuations (Np/wavelength)
rho_ext, cs_ext, cl_ext = 1600, 1700, 2810 #for the external medium
kappas_ext, kappal_ext = 0.100, 0.043 #for the external medium
alpha = 2+4j #average value of the absorbing function inside the PML
omega = 2*np.pi*np.linspace(0, 26e6/(a*1e3), num=400) #angular frequencies (rad/s)
nev = 10 #number of eigenvalues
target = lambda omega: omega/cl_core.real #target set at the longitudinal wavenumber of core to find L(0,n) modes
cell_type = 'quadrilateral' #'triangle', 'quadrilateral'
gll = True #False: Lagrange 2nd order element, True: 8th order element with Gauss-Lobatto-Legendre points (spectral-like)

##################################
# Re-scaling
L0 = a #characteristic length
T0 = a/cs_core #characteristic time
M0 = rho_core*a**3#characteristic mass
a, b, le = a/L0, b/L0, le/L0
rho_core, cs_core, cl_core = rho_core/M0*L0**3, cs_core/L0*T0, cl_core/L0*T0
rho_ext, cs_ext, cl_ext = rho_ext/M0*L0**3, cs_ext/L0*T0, cl_ext/L0*T0
omega = omega*T0
cs_core, cl_core = cs_core/(1+1j*kappas_core/2/np.pi), cl_core/(1+1j*kappal_core/2/np.pi) #complex celerities (core)
cs_ext, cl_ext = cs_ext/(1+1j*kappas_ext/2/np.pi), cl_ext/(1+1j*kappal_ext/2/np.pi) #complex celerities (exterior)

##################################
# Create mesh from Gmsh and finite elements (six-node triangles with three dofs per node for the three components of displacement)
gmsh.initialize()
# Core
gmsh.model.geo.addPoint(+a, -a, 0, le, 1)
gmsh.model.geo.addPoint(+a, +a, 0, le, 2)
gmsh.model.geo.addPoint(-a, +a, 0, le, 3)
gmsh.model.geo.addPoint(-a, -a, 0, le, 4)
gmsh.model.geo.addLine(1, 2, 1)
gmsh.model.geo.addLine(2, 3, 2)
gmsh.model.geo.addLine(3, 4, 3)
gmsh.model.geo.addLine(4, 1, 4)
gmsh.model.geo.addCurveLoop([1, 2, 3, 4], 1)
core = gmsh.model.geo.addPlaneSurface([1])
# External medium
gmsh.model.geo.addPoint(+b, -b, 0, le, 5)
gmsh.model.geo.addPoint(+b, +b, 0, le, 6)
gmsh.model.geo.addPoint(-b, +b, 0, le, 7)
gmsh.model.geo.addPoint(-b, -b, 0, le, 8)
gmsh.model.geo.addLine(5, 6, 5)
gmsh.model.geo.addLine(6, 7, 6)
gmsh.model.geo.addLine(7, 8, 7)
gmsh.model.geo.addLine(8, 5, 8)
gmsh.model.geo.addCurveLoop([5, 6, 7, 8], 2)
exterior = gmsh.model.geo.addPlaneSurface([2, 1])
# Physical groups
gmsh.model.geo.synchronize() #the CAD entities must be synchronized with the Gmsh model
gmsh.model.addPhysicalGroup(2, [core], tag=1) #2: for 2D (surface)
gmsh.model.addPhysicalGroup(2, [exterior], tag=2) #2: for 2D (surface)
gmsh.model.addPhysicalGroup(1, [5, 6, 7, 8], tag=21) #1: for 1D (line)
# Generate mesh
gmsh.model.mesh.generate(2) #generate a 2D mesh
if cell_type == 'quadrilateral':
    gmsh.model.mesh.recombine() #recombine into quadrilaterals
#gmsh.model.mesh.setOrder(2) #interpolation order for the geometry, here 2nd order
# From gmsh to fenicsx
mesh, cell_tags, facet_tags = dolfinx.io.gmshio.model_to_mesh(gmsh.model, MPI.COMM_WORLD, 0, gdim=2)
# # Reminder for save & read
# gmsh.write("Elastic_Waveguide_Bar3D_Open.msh") #save to disk
# mesh, cell_tags, facet_tags = dolfinx.io.gmshio.read_from_msh("Elastic_Waveguide_Bar3D_Open.msh", MPI.COMM_WORLD, rank=0, gdim=2)
gmsh.finalize() #called when done using the Gmsh Python API
# Finite element space
if not gll:
    element = ufl.VectorElement("CG", cell_type, 2, 3) #Lagrange element, triangle, quadratic "P2", 3D vector
else: #spectral element-like
    import basix
    element = basix.create_element(basix.ElementFamily.P, basix.CellType.quadrilateral if cell_type=='quadrilateral' else basix.CellType.triangle, 8,
                                   basix.LagrangeVariant.gll_warped)
    element = basix.ufl_wrapper.BasixElement(element)
    element = basix.ufl_wrapper.VectorElement(element, 3)
V = dolfinx.fem.FunctionSpace(mesh, element)
# Visualize FE mesh with pyvista
Vmesh = dolfinx.fem.FunctionSpace(mesh, ufl.FiniteElement("CG", cell_type, 1)) #cell of order 1 is properly handled with pyvista
plotter = pyvista.Plotter()
grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(Vmesh))
grid.cell_data["Marker"] = cell_tags.values
grid.set_active_scalars("Marker")
plotter.add_mesh(grid, show_edges=True)
grid_nodes = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V)) #add higher-order nodes
plotter.add_mesh(grid_nodes, style='points', render_points_as_spheres=True, point_size=2)
plotter.view_xy()
plotter.show()

##################################
# Create Material properties, discontinuous between core and exterior
def isotropic_law(rho, cs, cl):
    E, nu = rho*cs**2*(3*cl**2-4*cs**2)/(cl**2-cs**2), 0.5*(cl**2-2*cs**2)/(cl**2-cs**2)
    C11 = C22 = C33 = E/(1+nu)/(1-2*nu)*(1-nu)
    C12 = C13 = C23 = E/(1+nu)/(1-2*nu)*nu
    C44 = C55 = C66 = E/(1+nu)/2
    C = np.array([[C11,C12,C13,0,0,0], 
                     [C12,C22,C23,0,0,0], 
                     [C13,C23,C33,0,0,0], 
                     [0,0,0,C44,0,0], 
                     [0,0,0,0,C55,0], 
                     [0,0,0,0,0,C66]])
    C = C[np.triu_indices(6)] #the upper-triangle part of C (21 elements only)
    return PETSc.ScalarType(C)
C_core = isotropic_law(rho_core, cs_core, cl_core)
C_ext = isotropic_law(rho_ext, cs_ext, cl_ext)
Q = dolfinx.fem.FunctionSpace(mesh, ufl.TensorElement("DG", cell_type, 0, (6,6), symmetry=True)) #symmetry enables to store 21 elements instead of 36
C = dolfinx.fem.Function(Q)
cells = cell_tags.find(1) #core (tag=1)
C.x.array[[range(21*c,21*c+21) for c in cells]] = np.tile(C_core, (len(cells),1))
cells = cell_tags.find(2) #exterior (tag=2)
C.x.array[[range(21*c,21*c+21) for c in cells]] = np.tile(C_ext, (len(cells),1))
# Same approach for density
Q = dolfinx.fem.FunctionSpace(mesh, ("DG", 0))
rho = dolfinx.fem.Function(Q)
cells = cell_tags.find(1) #core (tag=1)
rho.x.array[cells] = np.tile(rho_core, len(cells))
cells = cell_tags.find(2) #exterior (tag=2)
rho.x.array[cells] = np.tile(rho_ext, len(cells))
# Vizualization
plotter = pyvista.Plotter(window_size=[600, 400])
gridC = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(mesh, mesh.topology.dim)) #or *dolfinx.plot.create_vtk_mesh(Vmesh)
gridC.cell_data["Cij"] = C.x.array[0::21].real #index from 0 to 35, 0 being for C11...
#gridC.cell_data["Cij"] = rho.x.array.real #for checking density
gridC.set_active_scalars("Cij")
plotter.add_mesh(grid, style="wireframe", color="k") #FE mesh (with vertices of order 1 elements only owing to pyvista) 
plotter.add_mesh(gridC, opacity=0.8, show_scalar_bar=True, show_edges=False)
plotter.add_text('Re(Cij)', 'upper_edge', color='black', font_size=8)
plotter.view_xy()
plotter.show()

##################################
# Create Cartesian PML functions, continuous
def eval_gamma(x): #pml profile: continuous and parabolic
    values = [1+3*(alpha-1)*((abs(x[0])-a)/(b-a))**2*(abs(x[0])>=a), #gammax
              1+3*(alpha-1)*((abs(x[1])-a)/(b-a))**2*(abs(x[1])>=a)] #gammay
    return values
Q = dolfinx.fem.FunctionSpace(mesh, ufl.VectorElement("CG", cell_type, 2, 2))
gamma = dolfinx.fem.Function(Q)
gamma.interpolate(eval_gamma)
# Vizualization
plotter = pyvista.Plotter(window_size=[800, 400], shape=(1,2))
gridx = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(Q))
gridx.point_data["gammax"] = gamma.x.array[0::2].imag
gridx.set_active_scalars("gammax")
plotter.subplot(0,0)
plotter.add_mesh(grid, style="wireframe", color="k") #FE mesh
plotter.add_mesh(gridx, opacity=0.8, show_scalar_bar=True, show_edges=False)
plotter.add_text('Im(gammax)', 'upper_edge', color='black', font_size=8)
plotter.view_xy()
gridy = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(Q))
gridy.point_data["gammay"] = gamma.x.array[1::2].imag
gridy.set_active_scalars("gammay")
plotter.subplot(0,1)
plotter.add_mesh(grid, style="wireframe", color="k") #FE mesh
plotter.add_mesh(gridy, opacity=0.8, show_scalar_bar=True, show_edges=False)
plotter.add_text('Im(gammay)', 'upper_edge', color='black', font_size=8)
plotter.view_xy()
plotter.show()

##################################
# Create free boundary conditions (or uncomment lines below for Dirichlet)
# bcs = []
# Dirichlet test case:
boundary_facets = facet_tags.find(21) #outer boundary (tag=21)
boundary_dofs = dolfinx.fem.locate_dofs_topological(V=V, entity_dim=mesh.topology.dim-1, entities=boundary_facets)
bcs = [dolfinx.fem.dirichletbc(value=np.zeros(3, dtype=PETSc.ScalarType), dofs=boundary_dofs, V=V)]

##################################
# Define variational problem (SAFE method)
#ufl.dx = ufl.dx(metadata={"quadrature_rule": "GLL", "quadrature_degree": 2*(8-1)}) if gll else ufl.dx #uncomment to build diagonal mass matrix
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
Lx = lambda u: ufl.as_vector([u[0].dx(0), 0, 0, u[1].dx(0), u[2].dx(0), 0])
Ly = lambda u: ufl.as_vector([0, u[1].dx(1), 0, u[0].dx(1), 0, u[2].dx(1)])
Lz  = lambda u: ufl.as_vector([0, 0, u[2], 0, u[0], u[1]])
gammax = gamma[0]
gammay = gamma[1]
k0 = (ufl.inner(C*(gammay/gammax*Lx(u)+Ly(u)), Lx(v)) + ufl.inner(C*(Lx(u)+gammax/gammay*Ly(u)), Ly(v))) * ufl.dx
k0_form = dolfinx.fem.form(k0)
k1 = (gammay*ufl.inner(C*Lz(u), Lx(v))+gammax*ufl.inner(C*Lz(u), Ly(v))) * ufl.dx
k1_form = dolfinx.fem.form(k1)
k2 = ufl.inner(C*Lz(u), Lz(v)) * gammax * gammay * ufl.dx
k2_form = dolfinx.fem.form(k2)
m = rho * ufl.inner(u, v) * gammax * gammay * ufl.dx
mass_form = dolfinx.fem.form(m)

##################################
# Build PETSc matrices
M = dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0)
M.assemble()
K0 = dolfinx.fem.petsc.assemble_matrix(k0_form, bcs=bcs)
K0.assemble()
K1 = dolfinx.fem.petsc.assemble_matrix(k1_form, bcs=bcs, diagonal=0.0)
K1.assemble()
K2 = dolfinx.fem.petsc.assemble_matrix(k2_form, bcs=bcs, diagonal=0.0)
K2.assemble()

##################################
# Solve the eigenproblem with SLEPc\
# The parameter is omega, the eigenvalue is k
wg = Waveguide(MPI.COMM_WORLD, M, K0, K1, K2)
#(K2, M) = (wg._diag(K2.getDiagonal()), wg._diag(M.getDiagonal())) if gll else (K2, M) ##uncomment to save allocated memory of K2 and M if they are to be diagonal
wg.set_parameters(omega=omega)
wg.evp.setWhichEigenpairs(SLEPc.PEP.Which.TARGET_MAGNITUDE) #here, preferred to TARGET_IMAGINARY
wg.solve(nev, target=target)

##################################
# Plot dispersion curves\
# Results are to be compared with Fig. 8 of Treyssede, Journal of Computational Physics 314 (2016), 341–354
wg.plot_scaler["energy_velocity"] = L0/T0/1000 #units in m/ms
wg.plot_scaler["frequency"] = L0/T0/1000 #frequency units in MHz-mm
sc = wg.plot_energy_velocity()
sc.axes.set_xlim([0, 26])
sc.axes.set_ylim([0, 6])
sc.axes.set_xlabel('Frequency-halfwidth (MHz-mm)')
sc.axes.set_ylabel('Energy velocity (m/ms)')
#plt.savefig("buried_bar_energy_velocity.png")
wg.plot_scaler["attenuation"] = 8.686*1000 #units in dB-mm/m
sc = wg.plot_attenuation()
sc.axes.set_xlim([0, 26])
sc.axes.set_ylim([0, 2000])
sc.axes.set_xlabel('Frequency-halfwidth (MHz-mm)')
sc.axes.set_ylabel('Attenuation (dB-mm/m)')
#plt.savefig("buried_bar_attenuation.png")
plt.show()
#wg.plot_spectrum(index=0)
#plt.show()

##################################
# Mode shape visualization
ik, imode = 345, 0 #parameter index, mode index to visualize
print(f"f={wg.omega[ik]*L0/T0/1000/(2*np.pi):1.1f}MHz-mm, attenuation={8.686*1000*wg.eigenvalues[ik][imode].imag:1.2f}dB-mm/m")
vec = wg.eigenvectors[ik].getColumnVector(imode)
u_grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V))
u_grid["u"] = np.array(vec).real.reshape(int(np.array(vec).size/V.element.value_shape), int(V.element.value_shape)) #V.element.value_shape is equal to 3
u_plotter = pyvista.Plotter()
u_plotter.add_mesh(grid, style="wireframe", color="k") #FE mesh
u_plotter.add_mesh(u_grid.warp_by_vector("u", factor=1), opacity=1, show_scalar_bar=False, show_edges=False, cmap='viridis') #do not show edges of higher order elements with pyvista
#u_plotter.show_axes()
u_plotter.screenshot('buried_bar_mode_shape.png', transparent_background=True)
u_plotter.show()

""

4. Excitation of a three-dimensional elastic bar of circular cross-section

3D elastic waveguide example. The cross-section is a 2D circle with free boundary conditions on its 1D boundaries, material: elastic steel. The eigenproblem is solved with the varying parameter as the frequency (eigenvalues are then wavenumbers). The forced response is computed for a point force at the center node ot the cross-section. Results are to be compared with Figs. 5, 6 and 7 of paper: Treyssede, Wave Motion 87 (2019), 75-91.

# This file is a tutorial for waveguicsx (*), whose inputs are
# the matrices K0, K1, K2, M and the excitation vector F.
# In the tutorial, K0, K1, K2 and M are finite element matrices generated by FEnicSX (**).
#  (*) waveguicsx is a python library for solving complex waveguide problems
#      Copyright (C) 2023-2024  Fabien Treyssede
#      waveguicsx is free software distributed under the GNU General Public License
#      (https://github.com/treyssede/waveguicsx)
# (**) FEniCSx is an open-source computing platform for solving partial differential equations
#      distributed under the GNU Lesser General Public License (https://fenicsproject.org/)

##################################
# 3D elastic waveguide example\
# The cross-section is a 2D circle with free boundary conditions on its 1D boundaries, material: elastic steel\
# The waveguide FE formulation (SAFE) leads to the following eigenvalue problem:\
# $(\textbf{K}_1-\omega^2\textbf{M}+\text{i}k(\textbf{K}_2+\textbf{K}_2^\text{T})+k^2\textbf{K}_3)\textbf{U}=\textbf{0}$\
# This eigenproblem is solved with the varying parameter as the frequency (eigenvalues are then wavenumbers).\
# The forced response is computed for a point force at the center node ot the cross-section.\
# Results are to be compared with Figs. 5, 6 and 7 of paper: Treyssede, Wave Motion 87 (2019), 75-91.

import gmsh #full documentation entirely defined in the `gmsh.py' module
import dolfinx
import ufl
from mpi4py import MPI
from petsc4py import PETSc
from slepc4py import SLEPc
import numpy as np
import matplotlib.pyplot as plt
import pyvista

from waveguicsx.waveguide import Waveguide
#For proper use with a notebook, uncomment the following line:
#pyvista.set_jupyter_backend("static"); pyvista.start_xvfb() #try: "none", "static", "pythreejs", "ipyvtklink"...

##################################
# Input parameters
a = 2.7e-3 #cross-section radius (m)
le = a/8 #finite element characteristic length (m)
rho, cs, cl = 7800, 3296, 5963 #density (kg/m3), shear and longitudinal wave celerities (m/s)
kappas, kappal = 0*0.008, 0*0.003 #shear and longitudinal bulk wave attenuations (Np/wavelength)
omega = np.arange(0.1, 10.1, 0.1)*cs/a #angular frequencies (rad/s)
nev = 60 #number of eigenvalues

##################################
# Re-scaling
L0 = a #characteristic length
T0 = a/cs #characteristic time
M0 = rho*a**3#characteristic mass
a, le = a/L0, le/L0
rho, cs, cl = rho/M0*L0**3, cs/L0*T0, cl/L0*T0
omega = omega*T0
cs, cl = cs/(1+1j*kappas/2/np.pi), cl/(1+1j*kappal/2/np.pi) #complex celerities

##################################
# Create mesh from Gmsh and finite elements (six-node triangles with three dofs per node for the three components of displacement)
gmsh.initialize()
# Core
origin = gmsh.model.geo.addPoint(+0, 0, 0, le, 1)
gmsh.model.geo.addPoint(+a, 0, 0, le, 2)
gmsh.model.geo.addPoint(0, +a, 0, le, 3)
gmsh.model.geo.addPoint(-a, 0, 0, le, 4)
gmsh.model.geo.addPoint(0, -a, 0, le, 5)
gmsh.model.geo.addCircleArc(2, 1, 3, 1)
gmsh.model.geo.addCircleArc(3, 1, 4, 2)
gmsh.model.geo.addCircleArc(4, 1, 5, 3)
gmsh.model.geo.addCircleArc(5, 1, 2, 4)
gmsh.model.geo.addCurveLoop([1, 2, 3, 4], 1)
disk = gmsh.model.geo.addPlaneSurface([1])
# Physical groups
gmsh.model.geo.synchronize() #the CAD entities must be synchronized with the Gmsh model
gmsh.model.addPhysicalGroup(2, [disk], tag=1) #2: for 2D (surface)
#gmsh.model.addPhysicalGroup(1, [1, 2, 3, 4], tag=11) #1: for 1D (line)
# Generate mesh
gmsh.model.mesh.embed(0, [origin], 2, disk) #ensure node points at the origin
gmsh.model.mesh.generate(2) #generate a 2D mesh
gmsh.model.mesh.setOrder(2) #interpolation order for the geometry, here 2nd order
# From gmsh to fenicsx
mesh, cell_tags, facet_tags = dolfinx.io.gmshio.model_to_mesh(gmsh.model, MPI.COMM_WORLD, 0, gdim=2)
# # Reminder for save & read
# gmsh.write("Elastic_Waveguide_Bar3D_Open.msh") #save to disk
# mesh, cell_tags, facet_tags = dolfinx.io.gmshio.read_from_msh("Elastic_Waveguide_Bar3D_Open.msh", MPI.COMM_WORLD, rank=0, gdim=2)
gmsh.finalize() #called when done using the Gmsh Python API
# Finite element space
element = ufl.VectorElement("CG", "triangle", 2, 3) #Lagrange element, triangle, quadratic "P2", 3D vector
V = dolfinx.fem.FunctionSpace(mesh, element)
# Visualize FE mesh with pyvista
Vmesh = dolfinx.fem.FunctionSpace(mesh, ufl.FiniteElement("CG", "triangle", 1)) #order 1 is properly handled with pyvista
plotter = pyvista.Plotter()
grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(Vmesh))
grid.cell_data["Marker"] = cell_tags.values
grid.set_active_scalars("Marker")
plotter.add_mesh(grid, show_edges=True)
grid_nodes = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V)) #add higher-order nodes
plotter.add_mesh(grid_nodes, style='points', render_points_as_spheres=True, point_size=2)
plotter.view_xy()
plotter.show()

##################################
# Create Material properties (isotropic)
def isotropic_law(rho, cs, cl):
    E, nu = rho*cs**2*(3*cl**2-4*cs**2)/(cl**2-cs**2), 0.5*(cl**2-2*cs**2)/(cl**2-cs**2)
    C11 = C22 = C33 = E/(1+nu)/(1-2*nu)*(1-nu)
    C12 = C13 = C23 = E/(1+nu)/(1-2*nu)*nu
    C44 = C55 = C66 = E/(1+nu)/2
    return np.array([[C11,C12,C13,0,0,0], 
                     [C12,C22,C23,0,0,0], 
                     [C13,C23,C33,0,0,0], 
                     [0,0,0,C44,0,0], 
                     [0,0,0,0,C55,0], 
                     [0,0,0,0,0,C66]])
C = isotropic_law(rho, cs, cl)
C = dolfinx.fem.Constant(mesh, PETSc.ScalarType(C))

##################################
# Create free boundary conditions
bcs = []

##################################
# Define variational problem (SAFE method)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
Lxy = lambda u: ufl.as_vector([u[0].dx(0), u[1].dx(1), 0, u[0].dx(1)+u[1].dx(0), u[2].dx(0), u[2].dx(1)])
Lz  = lambda u: ufl.as_vector([0, 0, u[2], 0, u[0], u[1]])
k0 = ufl.inner(C*Lxy(u), Lxy(v)) * ufl.dx
k0_form = dolfinx.fem.form(k0)
k1 = ufl.inner(C*Lz(u), Lxy(v)) * ufl.dx
k1_form = dolfinx.fem.form(k1)
k2 = ufl.inner(C*Lz(u), Lz(v)) * ufl.dx
k2_form = dolfinx.fem.form(k2)
m = rho*ufl.inner(u, v) * ufl.dx
mass_form = dolfinx.fem.form(m)

##################################
# Build PETSc matrices
M = dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0)
M.assemble()
K0 = dolfinx.fem.petsc.assemble_matrix(k0_form, bcs=bcs)
K0.assemble()
K1 = dolfinx.fem.petsc.assemble_matrix(k1_form, bcs=bcs, diagonal=0.0)
K1.assemble()
K2 = dolfinx.fem.petsc.assemble_matrix(k2_form, bcs=bcs, diagonal=0.0)
K2.assemble()

##################################
# Solve the eigenproblem with SLEPc (the parameter is omega, the eigenvalue is k)
wg = Waveguide(MPI.COMM_WORLD, M, K0, K1, K2)
wg.set_parameters(omega=omega, two_sided=False)
wg.evp.setTolerances(tol=1e-10, max_it=20)
wg.solve(nev=nev, target=0) #access to components with: wg.eigenvalues[ik][imode], wg.eigenvectors[ik][idof,imode]

##################################
# Plot dispersion curves\
# Results are to be compared with Fig. 5 of Treyssede, Wave Motion 87 (2019), 75-91
wg.plot_energy_velocity(direction=+1)
plt.show()

##################################
# Excitation force definition (point force)
dof_coords = V.tabulate_dof_coordinates()
x0 = np.array([0, 0, 0]) #desired coordinate of point force
dof = int(np.argmin(np.linalg.norm(dof_coords - x0, axis=1))) #find nearest dof
print(f'Point force coordinates (nearest dof):  {(dof_coords[dof,:])}') #check
F0 = M.createVecRight()
dof = dof*3 + 2 #orientation along z
#dof = dof - 2 #uncomment this line for an orientation along x instead
F0[dof] = 1
##Uncomment lines below for distributed excitation
#body_force = dolfinx.fem.Constant(mesh, PETSc.ScalarType((0, 0, 1))) #body force of unit amplitude along z
#traction = dolfinx.fem.Constant(mesh, PETSc.ScalarType((0, 0, 0))) #zero traction
#ds = ufl.Measure("ds", domain=mesh)
#f = ufl.inner(body_force, v) * ufl.dx + ufl.inner(traction, v) * ds
#f_form = dolfinx.fem.form(f)
#F = dolfinx.fem.petsc.assemble_vector(f_form)
#F.assemble()

##################################
# Computation of excitability and forced response\
# Results are to be compared with Figs. 6 and 7 of Treyssede, Wave Motion 87 (2019), 75-91
wg.compute_response_coefficient(F=F0, dof=dof)
wg.plot_coefficient()
sc = wg.plot_excitability()
sc.axes.set_yscale('log')
sc.axes.set_ylim(1e-3,0.5e+1)
frequency, response, ll_abs, ll_angle = wg.compute_response(dof=dof, z=[5], spectrum=None, plot=True) #spectrum=excitation.spectrum
ll_abs[0].axes.set_yscale('log')
ll_abs[0].axes.set_ylim(1e-2,1e+1)
ll_abs[0].axes.get_lines()[0].set_color("black")
plt.close()

##################################
# Mode shape visualization
ik, imode = 50, 1 #parameter index, mode index to visualize
vec = wg.eigenvectors[ik].getColumnVector(imode)
u_grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V))
u_grid["u"] = np.array(vec).real.reshape(int(np.array(vec).size/V.element.value_shape), int(V.element.value_shape)) #V.element.value_shape is equal to 3
u_plotter = pyvista.Plotter()
u_plotter.add_mesh(grid, style="wireframe", color="k") #FE mesh
u_plotter.add_mesh(u_grid.warp_by_vector("u", factor=0.5), opacity=0.8, show_scalar_bar=True, show_edges=False) #do not show edges of higher order elements with pyvista
u_plotter.show_axes()
u_plotter.show()

5. Excitation of a three-dimensional elastic bar of circular cross-section with parallelization

3D elastic waveguide example. The cross-section is a 2D circle with free boundary conditions on its 1D boundaries, material: elastic steel. The eigenproblem is solved with the varying parameter as the frequency (eigenvalues are then wavenumbers). The forced response is computed for a point force at the center node ot the cross-section. Results are to be compared with Figs. 5, 6 and 7 of paper: Treyssede, Wave Motion 87 (2019), 75-91. In this example:

  • the parameter loop (here, the frequency loop) is distributed on all processes

  • FE mesh and matrices are built on each local process

# This file is a tutorial for waveguicsx (*), whose inputs are
# the matrices K0, K1, K2, M and the excitation vector F.
# In the tutorial, K0, K1, K2 and M are finite element matrices generated by FEnicSX (**).
#  (*) waveguicsx is a python library for solving complex waveguide problems
#      Copyright (C) 2023-2024  Fabien Treyssede
#      waveguicsx is free software distributed under the GNU General Public License
#      (https://github.com/treyssede/waveguicsx)
# (**) FEniCSx is an open-source computing platform for solving partial differential equations
#      distributed under the GNU Lesser General Public License (https://fenicsproject.org/)

##################################
# 3D elastic waveguide example\
# The cross-section is a 2D circle with free boundary conditions on its 1D boundaries, material: elastic steel\
# The waveguide FE formulation (SAFE) leads to the following eigenvalue problem:\
# $(\textbf{K}_1-\omega^2\textbf{M}+\text{i}k(\textbf{K}_2+\textbf{K}_2^\text{T})+k^2\textbf{K}_3)\textbf{U}=\textbf{0}$\
# This eigenproblem is solved with the varying parameter as the frequency (eigenvalues are then wavenumbers).\
# The forced response is computed for a point force at the center node ot the cross-section.\
# Results are to be compared with Figs. 5, 6 and 7 of paper: Treyssede, Wave Motion 87 (2019), 75-91.
# In this example:
# - the parameter loop (here, the frequency loop) is distributed on all processes
# - FE mesh and matrices are built on each local process
# Reminder for an execution in parallel mode (e.g. 8 processes):
#  mpiexec -n 8 python3 Elastic_Waveguide_CircularBar3D_ForcedResponse_Gmsh_ParallelizedLoop.py

import gmsh #full documentation entirely defined in the `gmsh.py' module
import dolfinx
import ufl
from mpi4py import MPI
from petsc4py import PETSc
from slepc4py import SLEPc
import numpy as np
import matplotlib.pyplot as plt
#import pyvista

from waveguicsx.waveguide import Waveguide
#For proper use with a notebook, uncomment the following line:
#pyvista.set_jupyter_backend("static"); pyvista.start_xvfb() #try: "none", "static", "pythreejs", "ipyvtklink"...

##################################
# Input parameters
a = 2.7e-3 #cross-section radius (m)
le = a/8 #finite element characteristic length (m)
rho, cs, cl = 7800, 3296, 5963 #density (kg/m3), shear and longitudinal wave celerities (m/s)
kappas, kappal = 0*0.008, 0*0.003 #shear and longitudinal bulk wave attenuations (Np/wavelength)
omega = np.arange(0.1, 10.1, 0.1)*cs/a #angular frequencies (rad/s)
nev = 300 #number of eigenvalues

##################################
# Re-scaling
L0 = a #characteristic length
T0 = a/cs #characteristic time
M0 = rho*a**3#characteristic mass
a, le = a/L0, le/L0
rho, cs, cl = rho/M0*L0**3, cs/L0*T0, cl/L0*T0
omega = omega*T0
cs, cl = cs/(1+1j*kappas/2/np.pi), cl/(1+1j*kappal/2/np.pi) #complex celerities

##################################
# Create mesh from Gmsh and finite elements (six-node triangles with three dofs per node for the three components of displacement)
gmsh.initialize()
# Core
origin = gmsh.model.geo.addPoint(+0, 0, 0, le, 1)
gmsh.model.geo.addPoint(+a, 0, 0, le, 2)
gmsh.model.geo.addPoint(0, +a, 0, le, 3)
gmsh.model.geo.addPoint(-a, 0, 0, le, 4)
gmsh.model.geo.addPoint(0, -a, 0, le, 5)
gmsh.model.geo.addCircleArc(2, 1, 3, 1)
gmsh.model.geo.addCircleArc(3, 1, 4, 2)
gmsh.model.geo.addCircleArc(4, 1, 5, 3)
gmsh.model.geo.addCircleArc(5, 1, 2, 4)
gmsh.model.geo.addCurveLoop([1, 2, 3, 4], 1)
disk = gmsh.model.geo.addPlaneSurface([1])
# Physical groups
gmsh.model.geo.synchronize() #the CAD entities must be synchronized with the Gmsh model
gmsh.model.addPhysicalGroup(2, [disk], tag=1) #2: for 2D (surface)
#gmsh.model.addPhysicalGroup(1, [1, 2, 3, 4], tag=11) #1: for 1D (line)
# Generate mesh
gmsh.model.mesh.embed(0, [origin], 2, disk) #ensure node points at the origin
gmsh.model.mesh.generate(2) #generate a 2D mesh
gmsh.model.mesh.setOrder(2) #interpolation order for the geometry, here 2nd order
# From gmsh to fenicsx
mesh, cell_tags, facet_tags = dolfinx.io.gmshio.model_to_mesh(gmsh.model, MPI.COMM_SELF, 0, gdim=2) #MPI.COMM_SELF = FE mesh is built on each local process
# # Reminder for save & read
# gmsh.write("Elastic_Waveguide_Bar3D_Open.msh") #save to disk
# mesh, cell_tags, facet_tags = dolfinx.io.gmshio.read_from_msh("Elastic_Waveguide_Bar3D_Open.msh", MPI.COMM_WORLD, rank=0, gdim=2)
gmsh.finalize() #called when done using the Gmsh Python API
# Finite element space
element = ufl.VectorElement("CG", "triangle", 2, 3) #Lagrange element, triangle, quadratic "P2", 3D vector
V = dolfinx.fem.FunctionSpace(mesh, element)

##################################
# Create Material properties (isotropic)
def isotropic_law(rho, cs, cl):
    E, nu = rho*cs**2*(3*cl**2-4*cs**2)/(cl**2-cs**2), 0.5*(cl**2-2*cs**2)/(cl**2-cs**2)
    C11 = C22 = C33 = E/(1+nu)/(1-2*nu)*(1-nu)
    C12 = C13 = C23 = E/(1+nu)/(1-2*nu)*nu
    C44 = C55 = C66 = E/(1+nu)/2
    return np.array([[C11,C12,C13,0,0,0], 
                     [C12,C22,C23,0,0,0], 
                     [C13,C23,C33,0,0,0], 
                     [0,0,0,C44,0,0], 
                     [0,0,0,0,C55,0], 
                     [0,0,0,0,0,C66]])
C = isotropic_law(rho, cs, cl)
C = dolfinx.fem.Constant(mesh, PETSc.ScalarType(C))

##################################
# Create free boundary conditions
bcs = []

##################################
# Define variational problem (SAFE method)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
Lxy = lambda u: ufl.as_vector([u[0].dx(0), u[1].dx(1), 0, u[0].dx(1)+u[1].dx(0), u[2].dx(0), u[2].dx(1)])
Lz  = lambda u: ufl.as_vector([0, 0, u[2], 0, u[0], u[1]])
k0 = ufl.inner(C*Lxy(u), Lxy(v)) * ufl.dx
k0_form = dolfinx.fem.form(k0)
k1 = ufl.inner(C*Lz(u), Lxy(v)) * ufl.dx
k1_form = dolfinx.fem.form(k1)
k2 = ufl.inner(C*Lz(u), Lz(v)) * ufl.dx
k2_form = dolfinx.fem.form(k2)
m = rho*ufl.inner(u, v) * ufl.dx
mass_form = dolfinx.fem.form(m)

##################################
# Build PETSc matrices
M = dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0)
M.assemble()
K0 = dolfinx.fem.petsc.assemble_matrix(k0_form, bcs=bcs)
K0.assemble()
K1 = dolfinx.fem.petsc.assemble_matrix(k1_form, bcs=bcs, diagonal=0.0)
K1.assemble()
K2 = dolfinx.fem.petsc.assemble_matrix(k2_form, bcs=bcs, diagonal=0.0)
K2.assemble()

##################################
# Excitation force definition (point force)
dof_coords = V.tabulate_dof_coordinates()
x0 = np.array([0, 0, 0]) #desired coordinate of point force
dof = int(np.argmin(np.linalg.norm(dof_coords - x0, axis=1))) #find nearest dof
print(f'Point force coordinates (nearest dof):  {(dof_coords[dof,:])}') #check
F0 = M.createVecRight()
dof = dof*3 + 2 #orientation along z
dof = dof - 2 #uncomment this line for an orientation along x instead
F0[dof] = 1
##Uncomment lines below for distributed excitation
#body_force = dolfinx.fem.Constant(mesh, PETSc.ScalarType((0, 0, 1))) #body force of unit amplitude along z
#traction = dolfinx.fem.Constant(mesh, PETSc.ScalarType((0, 0, 0))) #zero traction
#ds = ufl.Measure("ds", domain=mesh)
#f = ufl.inner(body_force, v) * ufl.dx + ufl.inner(traction, v) * ds
#f_form = dolfinx.fem.form(f)
#F = dolfinx.fem.petsc.assemble_vector(f_form)
#F.assemble()

##################################
# Parallelization
comm = MPI.COMM_WORLD #use all processes for the loop
size = comm.Get_size()  #number of processors
rank = comm.Get_rank()  #returns the rank of the process that called it within comm_world
# Split the parameter range and scatter to all
if rank == 0: #define on rank 0 only
    omega_split = np.array_split(omega, size) #split param in blocks of length size roughly
else:
    omega_split = None
omega_local = comm.scatter(omega_split, root=0) #scatter 1 block per process

##################################
# Solve the eigenproblem with SLEPc (the parameter is omega, the eigenvalue is k)
wg = Waveguide(MPI.COMM_SELF, M, K0, K1, K2) #MPI.COMM_SELF = SLEPc will used FE matrices on each local process
wg.set_parameters(omega=omega_local)
wg.evp.setTolerances(tol=1e-10, max_it=20)
wg.solve(nev=nev, target=0) #access to components with: wg.eigenvalues[ik][imode], wg.eigenvectors[ik][idof,imode]

##################################
# Computation of excitability and forced response
wg.compute_response_coefficient(F=F0, dof=dof)
frequency, response = wg.compute_response(dof=dof, z=5, spectrum=None) #spectrum=excitation.spectrum

##################################
# Gather
wg.omega = comm.reduce([wg.omega], op=MPI.SUM, root=0) #reduce works for lists: brackets are necessary (wg.omega is not a list but a numpy array)
wg.eigenvalues = comm.reduce(wg.eigenvalues, op=MPI.SUM, root=0)
#wg.eigenvectors = comm.reduce(wg.eigenvectors, op=MPI.SUM, root=0) #don't do this line: reduce cannot pickle 'petsc4py.PETSc.Vec' objects (keep the mode shapes distributed on each processor rather than gather them)
wg.coefficient = comm.reduce(wg.coefficient, op=MPI.SUM, root=0)
wg.excitability = comm.reduce(wg.excitability, op=MPI.SUM, root=0)
frequency = comm.reduce([frequency], op=MPI.SUM, root=0)
response = comm.reduce([response], op=MPI.SUM, root=0)

##################################
# Plots
if rank == 0:
    wg.omega = np.concatenate(wg.omega) #wg.omega is transformed to a numpy array for a proper use of wg.plot()
    wg.plot()
    wg.plot_coefficient()
    sc = wg.plot_excitability()
    sc.axes.set_yscale('log')
    sc.axes.set_ylim(1e-3,0.5e+1)
    frequency = np.concatenate(frequency)
    response = np.concatenate(response) #use np.concatenate(response, axis=1)  if z contains more than one value
    fig, ax = plt.subplots(1, 1)  
    ax.plot(frequency, np.abs(response.T), linewidth=1, linestyle="-", color="k")
    ax.set_xlabel('frequency')
    ax.set_ylabel('|u|')
    ax.set_yscale('log')
    ax.set_ylim(1e-2,1e+1)
    fig.tight_layout()
    #plt.savefig("figure_name.svg")
    plt.show()

6. Time response of a two-dimensional plate excited near its first ZGV resonance

2D (visco-)elastic waveguide example (Lamb modes in a plate excited near 1st ZGV resonance). The cross-section is a 1D line with free boundary conditions on its boundaries. Material: viscoelastic steel. The eigenproblem is solved with the varying parameter as the frequency (eigenvalues are then wavenumbers). Viscoelastic loss can be included by introducing imaginary parts (negative) to wave celerities. Results are to be compared with Figs. 5b, 7a and 8a of paper: Treyssede and Laguerre, JASA 133 (2013), 3827-3837. Note: the depth direction is x, the axis of propagation is z.

# This file is a tutorial for waveguicsx (*), whose inputs are
# the matrices K0, K1, K2, M and the excitation vector F.
# In the tutorial, K0, K1, K2 and M are finite element matrices generated by FEnicSX (**).
#  (*) waveguicsx is a python library for solving complex waveguide problems
#      Copyright (C) 2023-2024  Fabien Treyssede
#      waveguicsx is free software distributed under the GNU General Public License
#      (https://github.com/treyssede/waveguicsx)
# (**) FEniCSx is an open-source computing platform for solving partial differential equations
#      distributed under the GNU Lesser General Public License (https://fenicsproject.org/)

##################################
# 2D (visco-)elastic waveguide example (Lamb modes in a plate excited near 1st ZGV resonance)\
# The cross-section is a 1D line with free boundary conditions on its boundaries\
# material: viscoelastic steel\
# The waveguide FE formulation (SAFE) leads to the following eigenvalue problem:\
# $(\textbf{K}_1-\omega^2\textbf{M}+\text{i}k(\textbf{K}_2+\textbf{K}_2^\text{T})+k^2\textbf{K}_3)\textbf{U}=\textbf{0}$\
# This eigenproblem is solved with the varying parameter as the frequency (eigenvalues are then wavenumbers)\
# Viscoelastic loss can be included by introducing imaginary parts (negative) to wave celerities\
# Results are to be compared with Figs. 5b, 7a and 8a of paper: Treyssede and Laguerre, JASA 133 (2013), 3827-3837\
# Note: the depth direction is x, the axis of propagation is z

import dolfinx
import ufl
from mpi4py import MPI
from petsc4py import PETSc
from slepc4py import SLEPc
import numpy as np
import matplotlib.pyplot as plt
import pyvista

from waveguicsx.waveguide import Waveguide, Signal
#For proper use with a notebook, uncomment the following line:
#pyvista.set_jupyter_backend("static"); pyvista.start_xvfb() #try: "none", "static", "pythreejs", "ipyvtklink"...

##################################
# Input parameters
h = 0.01 #plate thickness (m)
N = 10 #number of finite elements along one half-side
rho, cs, cl = 7800, 3218, 6020 #core density (kg/m3), shear and longitudinal wave celerities (m/s)
kappas, kappal = 0*0.008, 0*0.003 #core shear and longitudinal bulk wave attenuations (Np/wavelength)
nev = 20 #number of eigenvalues

##################################
# Excitation spectrum
excitation = Signal(alpha=0*np.log(50)/5e-3)
excitation.toneburst(fs=1000e3, T=5e-3, fc=250e3, n=5)
excitation.plot()
excitation.plot_spectrum()
#excitation.ifft(coeff=1); excitation.plot() #uncomment for check only
plt.show()
omega = 2*np.pi*excitation.frequency #angular frequency range (rad/s)

##################################
# Re-scaling
L0 = h #characteristic length
T0 = h/cs #characteristic time
M0 = rho*h**3#characteristic mass
h = h/L0
rho, cs, cl = rho/M0*L0**3, cs/L0*T0, cl/L0*T0
omega = omega*T0
cs, cl = cs/(1+1j*kappas/2/np.pi), cl/(1+1j*kappal/2/np.pi) #complex celerities (core)

##################################
# Create mesh and finite elements (three-node lines with two dofs per node for the two components of displacement)
mesh = dolfinx.mesh.create_interval(MPI.COMM_WORLD, N, np.array([0, h]))
element = ufl.VectorElement("CG", "interval", 2, 2) #Lagrange element, line element, quadratic "P2", 2D vector
V = dolfinx.fem.FunctionSpace(mesh, element)

##################################
# Create Material properties (isotropic)
def isotropic_law(rho, cs, cl):
    E, nu = rho*cs**2*(3*cl**2-4*cs**2)/(cl**2-cs**2), 0.5*(cl**2-2*cs**2)/(cl**2-cs**2)
    C11 = C22 = E/(1+nu)/(1-2*nu)*(1-nu)
    C12 = E/(1+nu)/(1-2*nu)*nu
    C33 = E/(1+nu)/2
    return ((C11,C12,0), 
            (C12,C22,0), 
            (0,0,C33))
C = isotropic_law(rho, cs, cl)
C = dolfinx.fem.Constant(mesh, PETSc.ScalarType(C))

##################################
# Create free boundary conditions
bcs = []

##################################
# Define variational problem (SAFE method)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
Lxy = lambda u: ufl.as_vector([u[0].dx(0), 0, u[1].dx(0)])
Lz = lambda u: ufl.as_vector([0, u[1], u[0]])
k0 = ufl.inner(C*Lxy(u), Lxy(v)) * ufl.dx
k0_form = dolfinx.fem.form(k0)
k1 = ufl.inner(C*Lz(u), Lxy(v)) * ufl.dx
k1_form = dolfinx.fem.form(k1)
k2 = ufl.inner(C*Lz(u), Lz(v)) * ufl.dx
k2_form = dolfinx.fem.form(k2)
m = rho*ufl.inner(u, v) * ufl.dx
mass_form = dolfinx.fem.form(m)

##################################
# Build PETSc matrices
M = dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0)
M.assemble()
K0 = dolfinx.fem.petsc.assemble_matrix(k0_form, bcs=bcs)
K0.assemble()
K1 = dolfinx.fem.petsc.assemble_matrix(k1_form, bcs=bcs, diagonal=0.0)
K1.assemble()
K2 = dolfinx.fem.petsc.assemble_matrix(k2_form, bcs=bcs, diagonal=0.0)
K2.assemble()

##################################
# Solve the eigenproblem with SLEPc\
# The parameter is omega, the eigenvalue is k
wg = Waveguide(MPI.COMM_WORLD, M, K0, K1, K2)
wg.set_parameters(omega=omega)
wg.solve(nev) #access to components with: wg.eigenvalues[ik][imode], wg.eigenvectors[ik][idof,imode]
wg.plot()
plt.show() #blocking

###############################################################################
# Excitation force: point force at x=0 oriented in the x-direction
dof_coords = V.tabulate_dof_coordinates()
x0 = np.array([0, 0, 0]) #desired coordinate of point force
dof = int(np.argmin(np.linalg.norm(dof_coords - x0, axis=1))) #find nearest dof
print(f'Point force coordinates (nearest dof):  {(dof_coords[dof,:])}') #check
F = M.createVecRight()
dof = dof*2 + 0 #x-direction
F[dof] = 1

###############################################################################
# Computation of excitabilities and forced response\
# Results are to be compared with Figs. 5b and 7a of Treyssede and Laguerre, JASA 133 (2013), 3827-3837
wg.compute_response_coefficient(F=F, dof=dof)
wg.set_plot_scaler()
sc = wg.plot_excitability()
sc.axes.set_yscale('log')
sc.axes.set_ylim(1e-3,1e2)
wg.set_plot_scaler(length=L0, time=T0, mass=M0, dim=2) #to visualize dimensional results
frequency, response, ll_abs, ll_angle = wg.compute_response(dof=dof, z=h/2, spectrum=excitation.spectrum, plot=True)
ll_abs[0].axes.get_lines()[0].set_color("black")
ll_abs[0].axes.set_xlim([0, 0.5e6])
ll_abs[0].axes.set_ylim([0, 1e-4])
plt.close()

###############################################################################
# Time response\
# Results are to be compared with Fig. 8a of Treyssede and Laguerre, JASA 133 (2013), 3827-3837
response = Signal(frequency=frequency, spectrum=response, alpha=0*np.log(50)/5e-3*T0)
#response.plot_spectrum()
response.ifft(coeff=1)
fig, ax = plt.subplots(1,1)
fig.set_figheight(4)
fig.set_figwidth(10)
ax = response.plot(ax=ax)
ax.set_xlim([0, 5e-3])
ax.set_xlabel('time (s)')
ax.set_ylim([-2e-3, 2e-3])
ax.set_ylabel('displacement (m)')
#plt.savefig('plate_zgv_transient.png')
plt.show()

##################################
# Mesh visualization
grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(mesh, mesh.topology.dim))
plotter = pyvista.Plotter()
plotter.add_mesh(grid, show_edges=True)
plotter.add_mesh(grid, style='points', render_points_as_spheres=True, point_size=10)
plotter.view_xy()
plotter.show()

##################################
# Mode shape visualization
ik, imode = 100, 5 #parameter index, mode index to visualize
vec = wg.eigenvectors[ik].getColumnVector(imode)*wg.plot_scaler["eigenvectors"] #note: multiplying by wg.plot_scaler["eigenvectors"] enables to normalize the cross-section power flow of eigenmodes to 1 Watt(/m)
u_grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V))
u_grid["u"] = np.array(vec).real.reshape(int(np.array(vec).size/V.element.value_shape), int(V.element.value_shape)) #V.element.value_shape is equal to 2
u_grid["u"] = np.insert(u_grid["u"], 1, 0, axis=1) #insert a zero column to the second component (the y component)
u_plotter = pyvista.Plotter()
u_plotter.add_mesh(grid, style="wireframe", color="k") #FE mesh
u_plotter.add_mesh(u_grid.warp_by_vector("u", factor=1e-2), opacity=0.8, show_scalar_bar=True, show_edges=False) #do not show edges of higher order elements with pyvista
u_plotter.view_zx()
u_plotter.show_axes()
u_plotter.show()

###############################################################################
# Save matrices into file (for basic usage of waveguicsx, see README examples)
viewer = PETSc.Viewer().createBinary('BasicExample.dat', 'w')
K0.view(viewer=viewer)
K1.view(viewer=viewer)
K2.view(viewer=viewer)
M.view(viewer=viewer)
F.view(viewer=viewer)

7. Dispersion curves of a rail

3D elastic waveguide example. The cross-section is a 2D rail profile (60E1, 60.21kg/m) with free boundary conditions on its 1D boundaries, material: elastic steel. This eigenproblem is solved with the varying parameter as the frequency (eigenvalues are then wavenumbers). The FE mesh is built from gmsh with a .geo file.

# This file is a tutorial for waveguicsx (*), whose inputs are
# the matrices K0, K1, K2, M and the excitation vector F.
# In the tutorial, K0, K1, K2 and M are finite element matrices generated by FEnicSX (**).
#  (*) waveguicsx is a python library for solving complex waveguide problems
#      Copyright (C) 2023-2024  Fabien Treyssede
#      waveguicsx is free software distributed under the GNU General Public License
#      (https://github.com/treyssede/waveguicsx)
# (**) FEniCSx is an open-source computing platform for solving partial differential equations
#      distributed under the GNU Lesser General Public License (https://fenicsproject.org/)

##################################
# 3D elastic waveguide example\
# The cross-section is a 2D rail profile (60E1, 60.21kg/m) with free boundary conditions on its 1D boundaries, material: elastic steel\
# The waveguide FE formulation (SAFE) leads to the following eigenvalue problem:\
# $(\textbf{K}_1-\omega^2\textbf{M}+\text{i}k(\textbf{K}_2+\textbf{K}_2^\text{T})+k^2\textbf{K}_3)\textbf{U}=\textbf{0}$\
# This eigenproblem is solved with the varying parameter as the frequency (eigenvalues are then wavenumbers).\
# The FE mesh is built from gmsh with a .geo file.

import dolfinx
import ufl
from mpi4py import MPI
from petsc4py import PETSc
from slepc4py import SLEPc
import numpy as np
import matplotlib.pyplot as plt
import pyvista

from waveguicsx.waveguide import Waveguide
#For proper use with a jupyter notebook, uncomment the following line:
#pyvista.set_jupyter_backend("none"); pyvista.start_xvfb() #try also: "static", "pythreejs", "ipyvtklink"...

##################################
# Input parameters
a = np.sqrt(76.70e-4)/2 #characteristic length (m) = halfwidth of a square of identical cross-section (rail cross-section : 76.70cm2)
le = a/8 #finite element characteristic length (m)
rho, cs, cl = 7850.0, 3207.7, 6001.0 #density (kg/m3), shear and longitudinal wave celerities (m/s)
kappas, kappal = 0*0.008, 0*0.003 #shear and longitudinal bulk wave attenuations (Np/wavelength)
omega, nev = 2*np.pi*np.linspace(100, 5000, 50), 20 #angular frequencies (rad/s) and number of eigenvalues (low-frequency regime)
#wavenumber, nev = 2*np.pi*np.linspace(1e3, 1e5, 100)/cs, 100 #angular frequencies (rad/s) and number of eigenvalues (high-frequency regime)
#reminder:
#E, nu, rho = 2.1e+11, 0.3, 7850
#cs = np.sqrt(E/rho/(2*(1+nu)))
#cl = np.sqrt(E/rho*(1-nu)/((1+nu)*(1-2*nu)))

##################################
# Re-scaling
L0 = a #characteristic length
T0 = a/cs #characteristic time
M0 = rho*a**3#characteristic mass
a, le = a/L0, le/L0
rho, cs, cl = rho/M0*L0**3, cs/L0*T0, cl/L0*T0
omega = omega*T0
#wavenumber = wavenumber*L0
cs, cl = cs/(1+1j*kappas/2/np.pi), cl/(1+1j*kappal/2/np.pi) #complex celerities

##################################
# Build mesh from Gmsh and import to dolfinx (six-node triangles with three dofs per node for the three components of displacement)

!gmsh rail60E1_60.21kg_m.geo -2 -order 2 -setnumber Mesh.MeshSizeFactor {le*L0*1000} -setnumber Mesh.ScalingFactor {1/(1000*L0)} #units in .geo are in mm
mesh, cell_tags, facet_tags = dolfinx.io.gmshio.read_from_msh("rail60E1_60.21kg_m.msh", MPI.COMM_WORLD, rank=0, gdim=2)

##################################
# Visualize FE mesh with pyvista
Vmesh = dolfinx.fem.FunctionSpace(mesh, ufl.FiniteElement("CG", "triangle", 1)) #order 1 is properly handled with pyvista
plotter = pyvista.Plotter()
grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(Vmesh))
grid.cell_data["Marker"] = cell_tags.values
grid.set_active_scalars("Marker")
plotter.add_mesh(grid, show_edges=True, show_scalar_bar=False)
plotter.view_xy()
plotter.show()
# Finite element space
element = ufl.VectorElement("CG", "triangle", 2, 3) #Lagrange element, triangle, quadratic "P2", 3D vector
V = dolfinx.fem.FunctionSpace(mesh, element)

##################################
# Create Material properties (isotropic)
def isotropic_law(rho, cs, cl):
    E, nu = rho*cs**2*(3*cl**2-4*cs**2)/(cl**2-cs**2), 0.5*(cl**2-2*cs**2)/(cl**2-cs**2)
    C11 = C22 = C33 = E/(1+nu)/(1-2*nu)*(1-nu)
    C12 = C13 = C23 = E/(1+nu)/(1-2*nu)*nu
    C44 = C55 = C66 = E/(1+nu)/2
    return np.array([[C11,C12,C13,0,0,0], 
                     [C12,C22,C23,0,0,0], 
                     [C13,C23,C33,0,0,0], 
                     [0,0,0,C44,0,0], 
                     [0,0,0,0,C55,0], 
                     [0,0,0,0,0,C66]])
C = isotropic_law(rho, cs, cl)
C = dolfinx.fem.Constant(mesh, PETSc.ScalarType(C))

##################################
# Create free boundary conditions
bcs = []

##################################
# Define variational problem (SAFE method)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
Lxy = lambda u: ufl.as_vector([u[0].dx(0), u[1].dx(1), 0, u[0].dx(1)+u[1].dx(0), u[2].dx(0), u[2].dx(1)])
Lz  = lambda u: ufl.as_vector([0, 0, u[2], 0, u[0], u[1]])
k0 = ufl.inner(C*Lxy(u), Lxy(v)) * ufl.dx
k0_form = dolfinx.fem.form(k0)
k1 = ufl.inner(C*Lz(u), Lxy(v)) * ufl.dx
k1_form = dolfinx.fem.form(k1)
k2 = ufl.inner(C*Lz(u), Lz(v)) * ufl.dx
k2_form = dolfinx.fem.form(k2)
m = rho*ufl.inner(u, v) * ufl.dx
mass_form = dolfinx.fem.form(m)

##################################
# Build PETSc matrices
M = dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0)
M.assemble()
K0 = dolfinx.fem.petsc.assemble_matrix(k0_form, bcs=bcs)
K0.assemble()
K1 = dolfinx.fem.petsc.assemble_matrix(k1_form, bcs=bcs, diagonal=0.0)
K1.assemble()
K2 = dolfinx.fem.petsc.assemble_matrix(k2_form, bcs=bcs, diagonal=0.0)
K2.assemble()

##################################
# Solve the eigenproblem with SLEPc (the parameter is omega, the eigenvalue is k)
wg = Waveguide(MPI.COMM_WORLD, M, K0, K1, K2)
wg.set_parameters(omega=omega)
#wg.set_parameters(wavenumber=wavenumber)
wg.evp.setTolerances(tol=1e-8, max_it=20)
wg.solve(nev=nev, target=0) #access to components with: wg.eigenvalues[ik][imode], wg.eigenvectors[ik][idof,imode]

##################################
# Plot dispersion curves in the low-frequency regime (for comparison, see e.g. Zhang et al., Mechanical Systems and Signal Processing, 150, 1-22 (2021))
wg.set_plot_scaler(length=L0, time=T0, mass=M0)
sc = wg.plot()
sc.axes.set_xlim([0, 25])
sc.axes.set_ylim([0, 5000])
sc = wg.plot_phase_velocity()
sc.axes.set_xlim([0, 5000])
sc.axes.set_ylim([0, 6000])
sc = wg.plot_group_velocity(direction=+1)
sc.axes.set_xlim([0, 5000])
sc.axes.set_ylim([0, 6000])
plt.show()

###############################################################################
# Plot phase velocity dispersion curves in the high-frequency regime (for comparison, see e.g. Ge et al., Structural Health Monitoring, 21, 1287-1308 (2022))
if False:
    wg.set_plot_scaler(length=L0, time=T0, mass=M0)
    wg.plot_scaler['frequency'] = wg.plot_scaler['frequency']/1000
    sc = wg.plot_phase_velocity()
    sc.axes.set_xlim([0, 1e2])
    sc.axes.set_ylim([0, 12000])
    sc.axes.set_xlabel('frequency (kHz)')
    sc.axes.set_ylabel('phase velocity (m/s)')
    #plt.savefig('rail_phase_velocity.png')
    plt.show()

##################################
# Mode shape visualization
ik, imode = 49, 9 #parameter index, mode index to visualize
print(f"f={wg.omega[ik]/(2*np.pi*T0):1.1f}Hz, l={2*np.pi*L0/wg.eigenvalues[ik][imode].real:1.2f}m")
vec = wg.eigenvectors[ik].getColumnVector(imode)
grid_interp = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(Vmesh))
grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V))
grid["u"] = np.array(vec).real.reshape(int(np.array(vec).size/V.element.value_shape), int(V.element.value_shape)) #V.element.value_shape is equal to 3
grid_interp = grid_interp.interpolate(grid) #interpolation onto a finite element mesh of order 1 to plot both the FE mesh and the results with pyvista
plotter = pyvista.Plotter()
plotter.add_mesh(grid_interp.warp_by_vector("u", factor=0.1), show_scalar_bar=False, show_edges=True)
#plotter.show_axes()
plotter.view_xy()
plotter.screenshot('rail_mode_shape.png', transparent_background=True)
plotter.show()

8. Reflection of Lamb modes by the free edge of a plate

Scattering in 2D elastic waveguide example (reflection of Lamb modes by the free edge of a plate). The cross-section is a 1D line with free boundary conditions on its boundaries. The inhomogeneous part, including the free edge, is a 2D rectangle. Material: elastic steel. The problem is solved using FEM with transparent boundary condition (tbc) in the inlet cross-section. The inlet eigenproblem is solved using SAFE as a function of frequency (eigenvalues are wavenumbers). Results can be compared with Fig. 4 of paper: Karunasena et al., CMAME 125 (1995), 221-233 (see also Gregory and Gladwell, J. of Elast. 13 (1983), 185-206).

# This file is a tutorial for waveguicsx (*), whose inputs are:
# - the matrices K, M, C and the internal excitation vector F
# - the matrices K0, K1, K2 and Ms for each transparent BC.
# In the tutorial, these matrices are finite element matrices generated by FEnicSX (**).
#  (*) waveguicsx is a python library for solving complex waveguide problems
#      Copyright (C) 2023-2024  Fabien Treyssede
#      waveguicsx is free software distributed under the GNU General Public License
#      (https://github.com/treyssede/waveguicsx)
# (**) FEniCSx is an open-source computing platform for solving partial differential equations
#      distributed under the GNU Lesser General Public License (https://fenicsproject.org/)

##################################
# Scattering in 2D elastic waveguide example (reflection of Lamb modes by the free edge of a plate)\
# The cross-section is a 1D line with free boundary conditions on its boundaries\
# The inhomogeneous part, including the free edge, is a 2D rectangle\
# material: elastic steel\
# The problem is solved using FEM with transparent boundary condition (tbc) in the inlet cross-section\
# The inlet eigenproblem is solved using SAFE as a function of frequency (eigenvalues are wavenumbers)\
# Results can be compared with Fig. 4 of paper: Karunasena et al., CMAME 125 (1995), 221-233 (see also\
# Gregory and Gladwell, J. of Elast. 13 (1983), 185-206).

import gmsh #full documentation entirely defined in the `gmsh.py' module
import dolfinx
import ufl
from mpi4py import MPI
from petsc4py import PETSc
from slepc4py import SLEPc
import numpy as np
import matplotlib.pyplot as plt
import pyvista

from waveguicsx.waveguide import Waveguide
from waveguicsx.scattering import Scattering
#For proper use with a jupyter notebook, uncomment the following line:
#pyvista.set_jupyter_backend("static"); pyvista.start_xvfb() #try also: "static", "pythreejs", "ipyvtklink"...

##################################
# Input parameters
h = 1 #plate thickness (m)
L = 1 #plate length (m) for the inhomogeneous part
le = h/5 #finite element characteristic length (m)
E, nu, rho = 2.0e+11, 0.25, 7800 #Young's modulus (Pa), Poisson ratio, density (kg/m3)
kappas, kappal = 0.008*0, 0.003*0 #shear and longitudinal bulk wave attenuations (Np/wavelength)
cs, cl = np.sqrt(E/rho/(2*(1+nu))), np.sqrt(E/rho*(1-nu)/((1+nu)*(1-2*nu))) #shear and longitudinal wave celerities
omega = 2*np.sqrt(3)*np.linspace(1.48, 1.60, num=100)*cs/h #angular frequencies (rad/s)
nev = 30 #number of eigenvalues
free_end = True #True: only one tbc (at one end), False: two tbcs (at both ends, to check that an incident mode is outgoing properly)

##################################
# Re-scaling
L0 = h #characteristic length
T0 = h/cs #characteristic time
M0 = rho*h**3 #characteristic mass
h, L, le = h/L0, L/L0, le/L0
rho, cs, cl = rho/M0*L0**3, cs/L0*T0, cl/L0*T0
omega = omega*T0
cs, cl = cs/(1+1j*kappas/2/np.pi), cl/(1+1j*kappal/2/np.pi) #complex celerities

##################################
# Create FE mesh (2D)
gmsh.initialize()
gmsh.model.geo.addPoint(0, 0, 0, le, 1)
gmsh.model.geo.addPoint(h, 0, 0, le, 2)
gmsh.model.geo.addPoint(h, L, 0, le, 3)
gmsh.model.geo.addPoint(0, L, 0, le, 4)
tbc0 = gmsh.model.geo.addLine(1, 2, 1)
gmsh.model.geo.addLine(2, 3, 2)
tbc1 = gmsh.model.geo.addLine(3, 4, 3)
gmsh.model.geo.addLine(4, 1, 4)
gmsh.model.geo.addCurveLoop([1, 2, 3, 4], 1)
domain = gmsh.model.geo.addPlaneSurface([1])
gmsh.model.geo.synchronize()
gmsh.model.addPhysicalGroup(1, [tbc0], tag=0) #tag=0: 1st transparent boundary condition, 'inlet'
gmsh.model.addPhysicalGroup(1, [tbc1], tag=1) #tag=1: 2nd transparent boundary condition, 'outlet'
gmsh.model.addPhysicalGroup(2, [domain], tag=2) #tag=2: the whole FE domain (2D mesh)
gmsh.model.mesh.generate(2)
gmsh.model.mesh.setOrder(2)
#gmsh.model.mesh.recombine() #recombine into quadrilaterals
mesh, cell_tags, facet_tags = dolfinx.io.gmshio.model_to_mesh(gmsh.model, MPI.COMM_WORLD, 0, gdim=2)
gmsh.finalize()
tbc_tags = [0] if free_end else [0, 1]

##################################
# Finite element space (2D)
element = ufl.VectorElement("CG", "triangle", 2, 2) #Lagrange element, triangle, quadratic "P2", 2D vector
V = dolfinx.fem.FunctionSpace(mesh, element)

##################################
# Visualize FE mesh with pyvista
Vmesh = dolfinx.fem.FunctionSpace(mesh, ufl.FiniteElement("CG", "triangle", 1)) #cell of order 1 is properly handled with pyvista
plotter = pyvista.Plotter()
grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(Vmesh))
grid.cell_data["Marker"] = cell_tags.values
grid.set_active_scalars("Marker")
plotter.add_mesh(grid, show_edges=True)
grid_nodes = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V)) #add higher-order nodes
plotter.add_mesh(grid_nodes, style='points', render_points_as_spheres=True, point_size=2)
plotter.view_xy()
plotter.show()

##################################
# Create Material properties (isotropic)
def isotropic_law(rho, cs, cl):
    E, nu = rho*cs**2*(3*cl**2-4*cs**2)/(cl**2-cs**2), 0.5*(cl**2-2*cs**2)/(cl**2-cs**2)
    C11 = C22 = E/(1+nu)/(1-2*nu)*(1-nu)
    C12 = E/(1+nu)/(1-2*nu)*nu
    C33 = E/(1+nu)/2
    return ((C11,C12,0), 
            (C12,C22,0), 
            (0,0,C33))
C = isotropic_law(rho, cs, cl)
C = dolfinx.fem.Constant(mesh, PETSc.ScalarType(C))

##################################
# Create free boundary conditions
bcs = []

###############################################################################
# Define the 2D variational formulation (FE problem inside the box)
dx = ufl.Measure("dx", domain=mesh)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
L = lambda u: ufl.as_vector([u[0].dx(0), u[1].dx(1), u[1].dx(0)+u[0].dx(1)])
k = ufl.inner(C*L(u), L(v)) * dx
k_form = dolfinx.fem.form(k)
m = rho*ufl.inner(u, v) * dx
mass_form = dolfinx.fem.form(m)
M = dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0)
M.assemble()
K = dolfinx.fem.petsc.assemble_matrix(k_form, bcs=bcs)
K.assemble()
dofs = dolfinx.fem.Function(V)
dofs.vector.setValues(range(M.size[0]), range(M.size[0])) #global dof vector

###############################################################################
# Determination of sign of outward normal of transparent boundaries (for check)
ds = ufl.Measure("ds", domain=mesh, subdomain_data=facet_tags)
n = ufl.FacetNormal(mesh)
ey = dolfinx.fem.Constant(mesh, PETSc.ScalarType((0, 1))) #ey: unit vector along the waveguide axis
tbc_normal = []
for tag in tbc_tags:
    normal = ufl.inner(n, ey) * ds(tag) #integration of the y-component of normal along the boundary
    normal = dolfinx.fem.form(normal)
    normal = np.sign(dolfinx.fem.assemble_scalar(normal)).real
    tbc_normal.append(normal.astype('int8'))
print(tbc_normal)

##################################
# For each transparent boundary condition:
# - extract the FE mesh (1D)
# - define the variational form (SAFE formulation)
# - assemble matrices
# - locate tbc dofs in the global matrix
Ms, K0, K1, K2, tbc_dofs = [], [], [], [], []
for tag in tbc_tags:  
    #Extract the 1D mesh
    safe_mesh = dolfinx.mesh.create_submesh(mesh, mesh.topology.dim-1, facet_tags.indices[facet_tags.values==tag])[0]
    #Define the variational form
    safe_element = ufl.VectorElement("CG", safe_mesh.ufl_cell(), 2, 2)
    safe_V = dolfinx.fem.FunctionSpace(safe_mesh, safe_element)
    dx = ufl.Measure("dx", domain=safe_mesh)
    u = ufl.TrialFunction(safe_V)
    v = ufl.TestFunction(safe_V)
    Lxy = lambda u: ufl.as_vector([u[0].dx(0), 0, u[1].dx(0)])
    Lz = lambda u: ufl.as_vector([0, u[1], u[0]])
    k0 = ufl.inner(C*Lxy(u), Lxy(v)) * dx
    k0_form = dolfinx.fem.form(k0)
    k1 = ufl.inner(C*Lz(u), Lxy(v)) * dx
    k1_form = dolfinx.fem.form(k1)
    k2 = ufl.inner(C*Lz(u), Lz(v)) * dx
    k2_form = dolfinx.fem.form(k2)
    m = rho*ufl.inner(u, v) * dx
    mass_form = dolfinx.fem.form(m)
    Ms.append(dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0))
    Ms[-1].assemble()
    K0.append(dolfinx.fem.petsc.assemble_matrix(k0_form, bcs=bcs))
    K0[-1].assemble()
    K1.append(dolfinx.fem.petsc.assemble_matrix(k1_form, bcs=bcs, diagonal=0.0))
    K1[-1].assemble()
    K2.append(dolfinx.fem.petsc.assemble_matrix(k2_form, bcs=bcs, diagonal=0.0))
    K2[-1].assemble()
    #Locate tbc dofs in the global matrix (trick using interpolate)
    tbc_dofs.append(dolfinx.fem.Function(safe_V))
    tbc_dofs[-1].interpolate(dofs)
    tbc_dofs[-1] = tbc_dofs[-1].vector.array.real.round().astype('int32')    

##################################
# Scattering initialization
tbcs =  [('waveguide0', -tbc_dofs[0])] if free_end else [('waveguide0', -tbc_dofs[0]), ('waveguide1', +tbc_dofs[1])] #-1 because tbc_normal[0]=-1
ws = Scattering(MPI.COMM_WORLD, M, K, 0*M, tbcs)
#Solve waveguide problem of 1st tbc
print(f'\nTransparent boundary condition 0\n')
ws.waveguide0 = Waveguide(MPI.COMM_WORLD, Ms[0], K0[0], K1[0], K2[0])
ws.waveguide0.set_parameters(omega=omega)
ws.waveguide0.solve(nev)
ws.waveguide0.compute_traveling_direction()
ws.waveguide0.compute_poynting_normalization()
#Solve waveguide problem of 2nd tbc
if not free_end:
    print(f'\nTransparent boundary condition 1\n')
    ws.waveguide1 = Waveguide(MPI.COMM_WORLD, Ms[1], K0[1], K1[1], K2[1])
    ws.waveguide1.set_parameters(omega=omega)
    ws.waveguide1.solve(nev)
    ws.waveguide1.compute_traveling_direction()
    ws.waveguide1.compute_poynting_normalization()
print(f'\n')

##################################
# Solving scattering problem
index = np.argmin(np.abs(ws.waveguide0.eigenvalues[0]-4.36)) #4.36 is the S1 wavenumber value at angular frequency 5.13 (roughly)
mode = ws.waveguide0.track_mode(0, index, threshold=0.98, plot=True) #track ingoing S1 mode over the whole frequency range
ws.set_ingoing_mode('waveguide0', mode) #set mode as a single ingoing mode, coeff is 1 (its power is also 1 thanks to poynting normalization)
ws.set_parameters()
print(f'KSP solver type: {ws.ksp.getType()}')
ws.solve()
ws.plot_energy_balance() #checking tbc modal truncature
plt.show()

###############################################################################
# Plot modal coefficients vs. angular frequency\
# Results can be compared with Fig. 4 of paper: Karunasena et al., CMAME 125 (1995), 221-233\
# (see also Gregory and Gladwell, J. of Elast. 13 (1983), 185-206).

ws.waveguide0.compute_complex_power()
sc = ws.waveguide0.plot(y=('complex_power', lambda x:np.abs(np.real(x))), direction=-1)
sc.axes.set_ylim(0, 1.2)
sc.axes.set_ylabel('Reflected power')
#ws.waveguide0.plot_complex_power() #plot both real and imaginary part (signed)
#ws.waveguide0.plot_energy_velocity(c=('coefficient', np.abs)) #plot ve curves colored by |q|
if not free_end:
    ws.waveguide1.compute_complex_power()
    sc = ws.waveguide1.plot(y=('complex_power', lambda x:np.abs(np.real(x))), direction=+1)
    sc.axes.set_ylim(0, 1.2)
    sc.axes.set_ylabel('Transmitted power')

###############################################################################
# Example of plot for a single mode
index = np.argmin(np.abs(ws.waveguide0.eigenvalues[0]+4.36)) #-4.36 is the opposite S1 wavenumber value at angular frequency 5.13 (roughly)
mode2 = ws.waveguide0.track_mode(0, index, threshold=0.98, plot=False) #track reflected S1 mode
sc = ws.waveguide0.plot_complex_power(mode=mode2)
sc[0].axes.legend()
plt.show()

##################################
# Example of plot for displacement field at a given angular frequency
i = 50 #angular frequency index
vec = ws.displacement[i]
grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V))
#grid["ux"] = np.array(vec).real[0::2] #x-component of displacement
grid["uy"] = np.array(vec).real[1::2] #y-component of displacement
plotter = pyvista.Plotter()
plotter.add_mesh(grid)
plotter.show_axes()
plotter.view_xy()
plotter.show()

###############################################################################
# Save matrices into file (for basic usage of waveguicsx, see README examples)
viewer = PETSc.Viewer().createBinary('BasicScatteringExample.dat', 'w')
dofs = PETSc.Vec().createSeq(tbc_dofs[0].size)
dofs.setValues(range(tbc_dofs[0].size), tbc_dofs[0])
for obj in [M, K, Ms[0], K0[0], K1[0], K2[0], dofs]:
    obj.view(viewer=viewer)

9. Reflection and transmission of Pochhammer-Chree modes inside a cylinder

Scattering in 3D elastic waveguide example. Reflection of Pochhammer-Chree modes by the free edge of a cylinder or by notch. The cross-section is a 2D disk with free boundary conditions on its boundaries. The inhomogeneous part, including free edge or notch, is a 3D cylinder. Material: elastic steel. The problem is solved using FEM with transparent boundary condition in the inlet and outlet cross-sections. The tbc eigenproblem is solved using SAFE as a function of frequency (eigenvalues are wavenumbers)Results can be compared with the following papers, for free edge and notch respectively:

  • Gregory and Gladwell, Q. J. Mech. Appl. Math. 42 (1989), 327–337

  • Benmeddour et al., IJSS 48 (2011), 764-774.

# This file is a tutorial for waveguicsx (*), whose inputs are:
# - the matrices K, M, C and the internal excitation vector F
# - the matrices K0, K1, K2 and Ms for each transparent BC.
# In the tutorial, these matrices are finite element matrices generated by FEnicSX (**).
#  (*) waveguicsx is a python library for solving complex waveguide problems
#      Copyright (C) 2023-2024  Fabien Treyssede
#      waveguicsx is free software distributed under the GNU General Public License
#      (https://github.com/treyssede/waveguicsx)
# (**) FEniCSx is an open-source computing platform for solving partial differential equations
#      distributed under the GNU Lesser General Public License (https://fenicsproject.org/)

##################################
# Scattering in 3D elastic waveguide example
# Reflection of Pochhammer-Chree modes by the free edge of a cylinder or by notch\
# The cross-section is a 2D disk with free boundary conditions on its boundaries\
# The inhomogeneous part, including free edge or notch, is a 3D cylinder\
# material: elastic steel\
# The problem is solved using FEM with transparent boundary condition in the inlet and outlet cross-sections\
# The tbc eigenproblem is solved using SAFE as a function of frequency (eigenvalues are wavenumbers)\
# Results can be compared with the following papers, for free edge and notch respectively:
# - Gregory and Gladwell, Q. J. Mech. Appl. Math.  42 (1989), 327–337
# - Benmeddour et al., IJSS 48 (2011), 764-774.

import gmsh #full documentation entirely defined in the `gmsh.py' module
import dolfinx
import ufl
from mpi4py import MPI
from petsc4py import PETSc
from slepc4py import SLEPc
import numpy as np
import matplotlib.pyplot as plt
import pyvista

from waveguicsx.waveguide import Waveguide
from waveguicsx.scattering import Scattering
#For proper use with a jupyter notebook, uncomment the following line:
#pyvista.set_jupyter_backend("static"); pyvista.start_xvfb() #try also: "static", "pythreejs", "ipyvtklink"...

##################################
# Input parameters
a = 1 #cylinder radius (m)
L = 2*a #cylinder length (m) for the inhomogeneous part
le = a/2 #finite element characteristic length (m)
E, nu, rho = 2.0e+11, 0.25, 7800 #Young's modulus (Pa), Poisson ratio, density (kg/m3)
kappas, kappal = 0.008*0, 0.003*0 #shear and longitudinal bulk wave attenuations (Np/wavelength)
cs, cl = np.sqrt(E/rho/(2*(1+nu))), np.sqrt(E/rho*(1-nu)/((1+nu)*(1-2*nu)))
#omega = np.linspace(1.8, 2.4, num=50)*cl/a #angular frequencies (rad/s), for comparison with Gregory and Gladwell
omega = np.linspace(0.1, 3.0, num=30)*cl/a #for comparison with Benmeddour et al.
nev = 40 #number of eigenvalues
free_end = False #True: only one tbc (at one end), False: two tbcs (at both ends, note: check that an incident mode is outgoing properly)
notch_depth, notch_thickness = 1*a, 0.05*a #case without notch: set notch_depth to 0

##################################
# Re-scaling
L0 = a #characteristic length
T0 = a/cs #characteristic time
M0 = rho*a**3 #characteristic mass
a, L, le = a/L0, L/L0, le/L0
rho, cs, cl = rho/M0*L0**3, cs/L0*T0, cl/L0*T0
omega = omega*T0
cs, cl = cs/(1+1j*kappas/2/np.pi), cl/(1+1j*kappal/2/np.pi) #complex celerities (core)

##################################
# Create FE mesh (2D)
gmsh.initialize()
gmsh.model.occ.addCylinder(0, 0, 0, 0, 0, L, a, 0)
if notch_depth!=0: #notch case
    gmsh.model.occ.addBox(a, a, L/2, -2*a, -notch_depth, notch_thickness, 1)
    gmsh.model.occ.cut([(3, 0)], [(3, 1)], 2)
gmsh.model.occ.synchronize()
gmsh.model.mesh.generate(3)
volumes = gmsh.model.getEntities(dim=3)
gmsh.model.addPhysicalGroup(volumes[0][0], [volumes[0][1]], 0) #volume tag: 0
surfaces = gmsh.model.occ.getEntities(dim=2)
for s in surfaces:
    center_of_mass = gmsh.model.occ.getCenterOfMass(s[0], s[1]) #useful to identify surface numbers of tbcs
    if np.isclose(center_of_mass[2], 0):
        gmsh.model.addPhysicalGroup(s[0], [s[1]], 0) #inlet tbc tag: 0
    if np.isclose(center_of_mass[2], L):
        gmsh.model.addPhysicalGroup(s[0], [s[1]], 1) #outlet tbc tag: 1
gmsh.model.mesh.setOrder(2)
mesh, cell_tags, facet_tags = dolfinx.io.gmshio.model_to_mesh(gmsh.model, MPI.COMM_WORLD, 0, gdim=3)
gmsh.finalize()
tbc_tags = [0] if free_end else [0, 1]

##################################
# Finite element space (3D)
element = ufl.VectorElement("CG", "tetrahedron", 2, 3) #Lagrange element, tetrahedron, quadratic, 3D vector
V = dolfinx.fem.FunctionSpace(mesh, element)

##################################
# Visualize FE mesh with pyvista
Vmesh = dolfinx.fem.FunctionSpace(mesh, ufl.FiniteElement("CG", "tetrahedron", 1)) #cell of order 1 is properly handled with pyvista
plotter = pyvista.Plotter()
grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(Vmesh))
grid.cell_data["Marker"] = cell_tags.values
grid.set_active_scalars("Marker")
plotter.add_mesh(grid, show_edges=True)
grid_nodes = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V)) #add higher-order nodes
plotter.add_mesh(grid_nodes, style='points', render_points_as_spheres=True, point_size=2)
plotter.show()

##################################
# Create Material properties (isotropic)
def isotropic_law(rho, cs, cl):
    E, nu = rho*cs**2*(3*cl**2-4*cs**2)/(cl**2-cs**2), 0.5*(cl**2-2*cs**2)/(cl**2-cs**2)
    C11 = C22 = C33 = E/(1+nu)/(1-2*nu)*(1-nu)
    C12 = C13 = C23 = E/(1+nu)/(1-2*nu)*nu
    C44 = C55 = C66 = E/(1+nu)/2
    return np.array([[C11,C12,C13,0,0,0], 
                     [C12,C22,C23,0,0,0], 
                     [C13,C23,C33,0,0,0], 
                     [0,0,0,C44,0,0], 
                     [0,0,0,0,C55,0], 
                     [0,0,0,0,0,C66]])
C = isotropic_law(rho, cs, cl)
C = dolfinx.fem.Constant(mesh, PETSc.ScalarType(C))

##################################
# Create free boundary conditions
bcs = []

###############################################################################
# Define the 2D variational formulation
dx = ufl.Measure("dx", domain=mesh)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
L = lambda u: ufl.as_vector([u[0].dx(0), u[1].dx(1), u[2].dx(2),
                            u[0].dx(1)+u[1].dx(0), u[0].dx(2)+u[2].dx(0), u[1].dx(2)+u[2].dx(1)])
k = ufl.inner(C*L(u), L(v)) * dx
k_form = dolfinx.fem.form(k)
m = rho*ufl.inner(u, v) * dx
mass_form = dolfinx.fem.form(m)
M = dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0)
M.assemble()
K = dolfinx.fem.petsc.assemble_matrix(k_form, bcs=bcs)
K.assemble()
dofs = dolfinx.fem.Function(V)
dofs.vector.setValues(range(M.size[0]), range(M.size[0])) #global dof vector

###############################################################################
# Determination of sign of outward normal of transparent boundaries (for check)
ds = ufl.Measure("ds", domain=mesh, subdomain_data=facet_tags)
n = ufl.FacetNormal(mesh)
ez = dolfinx.fem.Constant(mesh, PETSc.ScalarType((0, 0, 1))) #ez: unit vector along the waveguide axis
tbc_normal = []
for tag in tbc_tags:
    normal = ufl.inner(n, ez) * ds(tag) #integration of the z-component of normal along the boundary
    normal = dolfinx.fem.form(normal)
    normal = np.sign(dolfinx.fem.assemble_scalar(normal)).real
    tbc_normal.append(normal.astype('int8'))
print(tbc_normal)

##################################
# For each transparent boundary condition:
# - extract the FE mesh (2D)
# - define the variational form (SAFE formulation)
# - assemble matrices
# - locate tbc dofs in the global matrix
Ms, K0, K1, K2, tbc_dofs = [], [], [], [], []
for tag in tbc_tags:  
    #Extract the 2D mesh
    safe_mesh = dolfinx.mesh.create_submesh(mesh, mesh.topology.dim-1, facet_tags.indices[facet_tags.values==tag])[0]
    #Define the variational form
    safe_element = ufl.VectorElement("CG", safe_mesh.ufl_cell(), 2, 3)
    safe_V = dolfinx.fem.FunctionSpace(safe_mesh, safe_element)
    dx = ufl.Measure("dx", domain=safe_mesh)
    u = ufl.TrialFunction(safe_V)
    v = ufl.TestFunction(safe_V)
    Lxy = lambda u: ufl.as_vector([u[0].dx(0), u[1].dx(1), 0, u[0].dx(1)+u[1].dx(0), u[2].dx(0), u[2].dx(1)])
    Lz  = lambda u: ufl.as_vector([0, 0, u[2], 0, u[0], u[1]])
    k0 = ufl.inner(C*Lxy(u), Lxy(v)) * dx
    k0_form = dolfinx.fem.form(k0)
    k1 = ufl.inner(C*Lz(u), Lxy(v)) * dx
    k1_form = dolfinx.fem.form(k1)
    k2 = ufl.inner(C*Lz(u), Lz(v)) * dx
    k2_form = dolfinx.fem.form(k2)
    m = rho*ufl.inner(u, v) * dx
    mass_form = dolfinx.fem.form(m)
    Ms.append(dolfinx.fem.petsc.assemble_matrix(mass_form, bcs=bcs, diagonal=0.0))
    Ms[-1].assemble()
    K0.append(dolfinx.fem.petsc.assemble_matrix(k0_form, bcs=bcs))
    K0[-1].assemble()
    K1.append(dolfinx.fem.petsc.assemble_matrix(k1_form, bcs=bcs, diagonal=0.0))
    K1[-1].assemble()
    K2.append(dolfinx.fem.petsc.assemble_matrix(k2_form, bcs=bcs, diagonal=0.0))
    K2[-1].assemble()
    #Locate tbc dofs in the global matrix (trick using interpolate)
    tbc_dofs.append(dolfinx.fem.Function(safe_V))
    tbc_dofs[-1].interpolate(dofs)
    tbc_dofs[-1] = tbc_dofs[-1].vector.array.real.round().astype('int32')    

##################################
# Scattering initialization
tbcs =  [('waveguide0', -tbc_dofs[0])] if free_end else [('waveguide0', -tbc_dofs[0]), ('waveguide1', +tbc_dofs[1])] #-1 because tbc_normal[0]=-1
ws = Scattering(MPI.COMM_WORLD, M, K, 0*M, tbcs)
#Solve waveguide problem of 1st tbc
print(f'\nTransparent boundary condition 0\n')
ws.waveguide0 = Waveguide(MPI.COMM_WORLD, Ms[0], K0[0], K1[0], K2[0])
ws.waveguide0.set_parameters(omega=omega)
ws.waveguide0.solve(nev)
ws.waveguide0.compute_traveling_direction()
ws.waveguide0.compute_poynting_normalization()
#Solve waveguide problem of 2nd tbc
if not free_end:
    print(f'\nTransparent boundary condition 1\n')
    ws.waveguide1 = Waveguide(MPI.COMM_WORLD, Ms[1], K0[1], K1[1], K2[1])
    ws.waveguide1.set_parameters(omega=omega)
    ws.waveguide1.solve(nev)
    ws.waveguide1.compute_traveling_direction()
    ws.waveguide1.compute_poynting_normalization()
print(f'\n')

##################################
# Solving scattering problem
# index = np.argmin(np.abs(ws.waveguide0.eigenvalues[0]-2.59)) #2.59 is the L(0,1) wavenumber value at angular frequency 3.12 (roughly), for comparison with Gregory and Gladwell
index = np.argmin(np.abs(ws.waveguide0.eigenvalues[0]-0.11)) #0.11 is the L(0,1) wavenumber value at angular frequency 0.17 (roughly), for comparison with Benmeddour et al.
mode = ws.waveguide0.track_mode(0, index, threshold=0.98, plot=True) #track mode over the whole frequency range
ws.set_ingoing_mode('waveguide0', mode) #set mode as a single ingoing mode, coeff is 1 (its power is also 1 thansk to poynting normalization)
ws.set_parameters(solver='direct')
ws.solve()
ws.plot_energy_balance() #checking tbc modal truncature
plt.show()
#ws.ksp.view()

###############################################################################
# Plot modal coefficients w.r.t angular frequency\
# Results can be compared with paper: Gregory and Gladwell, Q. J. Mech. Appl. Math.  42 (1989), 327–337
ws.waveguide0.compute_complex_power()
sc = ws.waveguide0.plot(y=('complex_power', lambda x:np.abs(np.real(x))), direction=-1)
sc.axes.set_ylim(0, 1.2)
sc.axes.set_ylabel('Reflected power')
#ws.waveguide0.plot_complex_power() #plot both real and imaginary part (signed)
#ws.waveguide0.plot_energy_velocity(c=('coefficient', np.abs)) #plot ve curves colored by |q|
if not free_end:
    ws.waveguide1.compute_complex_power()
    sc = ws.waveguide1.plot(y=('complex_power', lambda x:np.abs(np.real(x))), direction=+1)
    sc.axes.set_ylim(0, 1.2)
    sc.axes.set_ylabel('Transmitted power')

###############################################################################
# Example of plot for a single mode
# index = np.argmin(np.abs(ws.waveguide0.eigenvalues[0]+2.59)) #-2.59 is the opposite L(0,1) wavenumber value at angular frequency 3.12 (roughly)
index = np.argmin(np.abs(ws.waveguide0.eigenvalues[0]+0.11)) #-0.11 is the opposite L(0,1) wavenumber value at angular frequency 0.17 (roughly)
mode2 = ws.waveguide0.track_mode(0, index, threshold=0.98, plot=False) #track mode (0, 0))
sc = ws.waveguide0.plot_complex_power(mode=mode2)
sc[0].axes.legend()
plt.show()

##################################
# Example of plot for displacement field at a given angular frequency
i = 20 #angular frequency index
vec = ws.displacement[i]
grid = pyvista.UnstructuredGrid(*dolfinx.plot.create_vtk_mesh(V))
#grid["ux"] = np.array(vec).real[0::3] #x-component of displacement
#grid["uy"] = np.array(vec).real[1::3] #y-component of displacement
grid["uz"] = np.array(vec).real[2::3] #z-component of displacement
plotter = pyvista.Plotter()
plotter.add_mesh(grid)
plotter.show_axes()
plotter.show()

""