"""Functions for doing FEM analysis of gdsfactory/qnngds geometries"""
from collections import OrderedDict
from itertools import combinations
import numpy as np
from qnngds.typing import LayerSpec
import qnngds as qg
import matplotlib.pyplot as plt
try:
import shapely
from shapely import LineString
from shapely.geometry import Polygon
from skfem import (
asm,
Basis,
BilinearForm,
ElementDG,
ElementTriP0,
ElementTriP3,
FacetBasis,
Functional,
condense,
solve,
)
from skfem.helpers import dot, grad, inner
from skfem.io import from_meshio
from femwell.mesh import mesh_from_OrderedDict
except ImportError:
raise ImportError("qnngds.analysis.fem requires femwell to be installed")
[docs]class Result:
"""Class for storing results of FEM simulation"""
[docs] def __init__(self, mesh, element, sigma, A, u):
"""Constructor for Result
Args:
mesh (MeshTri1): skfem mesh (first-order triangular mesh)
element (ElementTriP3): piecewise cubic element
sigma (np.ndarray): conductivity for each dof
A (scipy.sparse.csr_matrix): discretized linear system
u (dict[tuple[str,str], np.ndarray]): solution vector for each port pair
"""
self.mesh = mesh
self.element = element
self.sigma = sigma
self.A = A
self.u = u
@BilinearForm
def _laplace_equation(u, v, w):
"""Helper for evaluating laplace equation
Args:
u: Trial basis
w: domain
v: Test basis
"""
return -inner(grad(u), w["sigma"] * grad(v))
@Functional
def _port_flux(w):
"""Calculate flux through port"""
return dot(w.n, grad(w["u"]))
@Functional
def _port_width(w):
"""Calculate port width"""
return 1
[docs]def get_squares(result: Result) -> dict[str, float]:
"""Calculate number of squares from result generated by solve_laplace"""
squares = {}
for port_pair in result.u.keys():
facet_basis = FacetBasis(
result.mesh, result.element, facets=result.mesh.boundaries[port_pair[0]]
)
squares[port_pair] = 1 / asm(_port_flux, facet_basis, u=result.u[port_pair])
return squares
[docs]def visualize_current(
result: Result, port_pair: tuple[str, str], normalize: bool = True, ax=None
):
"""Visualize current density of result generated by solve_laplace
Args:
result (Result): result obtained from running solve_laplace on a mesh
port_pair (tuple[str, str]): pair of ports to visualize
normalize (bool): default True, if False plot the un-normalized current in A/um
ax (plt.axis | None): existing plot axis to add plot to if not None.
Returns:
(plt.axis)
"""
if ax is None:
fig, ax = plt.subplots()
ax.set_aspect(1)
if (port_pair[1], port_pair[0]) in result.u:
port_pair = (port_pair[1], port_pair[0])
elif port_pair not in result.u:
raise KeyError(f"invalid {port_pair=}, please choose one of {result.u.keys()=}")
basis = Basis(result.mesh, result.element)
basis_current = Basis(result.mesh, ElementDG(result.element))
basis0 = basis.with_element(ElementTriP0())
current = basis_current.project(
np.linalg.norm(
-basis0.interpolate(result.sigma)
* grad(basis.interpolate(result.u[port_pair])),
axis=0,
)
)
if normalize:
# calculate current through port
facet_basis = FacetBasis(
result.mesh, result.element, facets=result.mesh.boundaries[port_pair[0]]
)
port_current = asm(_port_flux, facet_basis, u=result.u[port_pair])
# divide by port width
port_width = asm(_port_width, facet_basis)
current /= port_current / port_width
basis_current.plot(
current,
ax=ax,
shading="gouraud",
colorbar=True,
cmap="inferno",
vmin=0,
vmax=np.max(current),
)
doflocs = result.mesh.restrict(result.mesh.subdomains["wire"]).doflocs
ax.set_xlim((1.1 * np.min(doflocs[0]), 1.1 * np.max(doflocs[0])))
ax.set_ylim((1.1 * np.min(doflocs[1]), 1.1 * np.max(doflocs[1])))
ax.set_title("current density")
return ax
[docs]def visualize_mesh(mesh, ax=None):
"""Visualize mesh of wire/component
Args:
mesh (MeshTri1): scikit fem mesh
ax (plt.axis | None): existing plot axis to add plot to if not None.
Returns:
(plt.axis)
"""
if ax is None:
fig, ax = plt.subplots()
ax.set_aspect(1)
mesh.restrict(mesh.subdomains["wire"]).draw(ax=ax)
ax.set_title("mesh")
return ax
[docs]def visualize_boundaries(mesh, ax=None):
"""Visualize boundaries of wire/component
Args:
mesh (MeshTri1): scikit fem mesh
ax (plt.axis | None): existing plot axis to add plot to if not None.
Returns:
(plt.axis)
"""
if ax is None:
fig, ax = plt.subplots()
ax.set_aspect(1)
for subdomain in mesh.subdomains.keys() - {"gmsh:bounding_entities"} - {"air"}:
mesh.restrict(subdomain).draw(ax=ax, boundaries=True, boundaries_only=True)
ax.set_title("boundaries")
return ax
[docs]def make_mesh(device: qg.Device, layer: LayerSpec, tolerance: float = 0.01):
"""Generate a mesh from a qnngds.Device
Args:
device (qnngds.Device): input device to make a mesh of
layer (LayerSpec): layer of device to use.
tolerance (float): simplify tolerance for geometry to reduce mesh complexity
Returns:
(MeshTri1) mesh of the device
"""
pp = device.get_polygons(by_spec=qg.get_layer(layer).tuple)[0]
pp_simplified = np.asarray(LineString(pp).simplify(tolerance=tolerance).coords)
component = Polygon(pp_simplified)
# create boundaries for each port
port_dict = {}
min_port_width = np.inf
for port_name in device.ports:
port = device.ports[port_name]
theta = port.orientation * np.pi / 180
p1 = port.midpoint + port.width / 2 * np.array([np.sin(theta), -np.cos(theta)])
p2 = port.midpoint + port.width / 2 * np.array([-np.sin(theta), np.cos(theta)])
ls = LineString([p1, p2])
port_dict[port.name] = ls
min_port_width = min(min_port_width, port.width)
if min_port_width is np.inf:
raise RuntimeError("port width is too large, aborting")
# set up our simulation geometry
geometry = OrderedDict(
**port_dict,
wire=component,
air=component.buffer(np.sqrt(device.xsize * device.ysize)),
)
for key in geometry.keys():
geometry[key] = shapely.set_precision(geometry[key], 1e-4)
resolutions = dict(wire={"resolution": min_port_width / 2, "distance": 1})
mesh = from_meshio(
mesh_from_OrderedDict(
geometry,
resolutions,
default_resolution_max=min_port_width * 2,
filename="mesh.msh",
)
)
return mesh
[docs]def solve_laplace(mesh) -> Result:
"""Solve laplace equation for a meshed device
Args:
mesh (MeshTri1): mesh of device
Returns:
(Result)
"""
element = ElementTriP3()
basis = Basis(mesh, element)
basis0 = basis.with_element(ElementTriP0())
sigma = basis0.zeros()
# use buffer around component to define Neumann conditions
sigma[basis0.get_dofs(elements="air")] = 1e-9
sigma[basis0.get_dofs(elements="wire")] = 1
A = _laplace_equation.assemble(basis, sigma=basis0.interpolate(sigma))
# solve for one port at a time
ports = mesh.boundaries.keys() - {"wire___air"}
u = {}
for port_pair in combinations(ports, 2):
voltages = basis.zeros()
voltages[basis.get_dofs(port_pair[0])] = 1
u[port_pair] = solve(
*condense(A, basis.zeros(), D=basis.get_dofs(port_pair), x=voltages)
)
return Result(mesh, element, sigma, A, u)