I’m new to Numba and I’m trying to accelerate the speed of the following function:
@njit
def generate_mesh(
f_min,
f_max,
port,
pml_x,
vertices,
stl="model.stl",
factor=30,
factor_space=15,
fraction=500,
res_fraction=[6,6,6],
cell_ratio=2,
n=[3, 3, 3],
):
#np.set_printoptions(threshold=np.inf)
remesh = False
if remesh:
os.system("gmsh -2 remesh.geo -o tmp.stl -format stl")
# Generate mesh
unit = 1000 # 1000 = mm
c = 299792458
lambda_min = unit * c / f_max
max_cell = lambda_min / factor
max_cell_space = lambda_min / factor_space
min_cell = lambda_min / fraction
# Load tetrahedral mesh
#mesh = trimesh.load("data/tmp/" + stl)
#print(type(mesh))
#print(mesh)
# Fetch node coordinates
#vertices = trimesh.load("data/tmp/" + stl).vertices
# Clear duplicate axis entries
x = np.unique(vertices[:, 0])
y = np.unique(vertices[:, 1])
z = np.unique(vertices[:, 2])
# Refine port
x = np.append(x, [port[0][0], port[1][0]])
y = np.append(y, [port[0][1], port[1][1]])
z = np.append(z, [port[0][2], port[1][2]])
# Complex models
nodes = len(vertices)
#if nodes > 100000:
#mesh = np.array([x,y,z])
### if maximum distance > k * lambda
### PORT NEEDS REFINEMENT EITHER WAY
if nodes < 100000:
# Apply edge refinement
res = [max_cell/res_fraction[0], max_cell/res_fraction[1], max_cell/res_fraction[2]]
mesh = np.array([], dtype="object")
for i, axis in enumerate([x, y, z]):
fine = np.array([])
edge_refinements = np.array([-np.inf])
for line in axis:
edge_refinement = np.array([])
for m in range(-n[i], n[i] + 1):
edge_refinement = np.concatenate((np.array([line + m * res[i]]), edge_refinement))
edge_refinement = np.sort(edge_refinement)
for edge in edge_refinement:
if edge > np.max(edge_refinements)+min_cell or edge in axis:
fine = np.append(fine, edge)
#####edge_refinements = np.append(edge_refinements, edge)
fine = np.array(np.sort(fine))
if i == 0:
fine_0 = fine
elif i == 1:
fine_1 = fine
elif i == 2:
fine_2 = fine
#print(mesh)
#input()
#print(mesh)
mesh = np.array([fine_0, fine_1, fine_2], dtype="object")
###mesh = [x,y,z]
# Sort original vertices
x = np.sort(x)
y = np.sort(y)
z = np.sort(z)
# Add lambda/4 + PML padding cells for bounding box
f_center = (f_min + f_max)/2
lambda_center = unit * c / f_center
padding = lambda_center / 8 + pml_x * max_cell_space
#mesh[0] = np.insert(mesh[0], 0, x[0] - padding)
mesh[0] = np.append([x[0] - padding], mesh[0])
mesh[0] = np.append(mesh[0], x[-1] + padding)
#mesh[1] = np.insert(mesh[1], 0, y[0] - padding)
mesh[1] = np.append([y[0] - padding], mesh[1])
mesh[1] = np.append(mesh[1], y[-1] + padding)
#mesh[2] = np.insert(mesh[2], 0, z[0] - padding)
mesh[2] = np.append([z[0] - padding], mesh[2])
mesh[2] = np.append(mesh[2], z[-1] + padding)
# Global refinement
for i in range(len(mesh)):
ax1 = np.array(mesh[i], dtype=np.float64)
# Add line to every vertex
axis = np.unique(np.concatenate((np.unique(vertices[:, i]), ax1)))
to_delete = np.argwhere(np.ediff1d(axis) <= min_cell) + 1
q = 1
axis_tmp = axis
to_delete_tmp = np.array([])
while (np.argwhere(np.ediff1d(axis_tmp) <= min_cell) + 1).size > 1:
to_delete_tmp = np.delete(to_delete, np.arange(0, to_delete.size, q))
axis_tmp = np.delete(axis, to_delete_tmp)
q += 1
if to_delete_tmp.size != 0:
axis = np.delete(axis, to_delete_tmp)
# Interior + exterior refinement
j_tot = (np.argwhere(np.abs(np.ediff1d(axis)) >= max_cell) + 1).size
for count, j in enumerate(
(np.argwhere(np.abs(np.ediff1d(axis)) >= max_cell) + 1)
):
if count != 0 and count != j_tot - 1: # Interior
n = int((np.abs((axis[j - 1] - axis[j]) / max_cell) + 1)[0])
else: # Exterior
n = int((np.abs((axis[j - 1] - axis[j]) / max_cell_space) + 1)[0])
delta = np.abs(axis[j - 1] - axis[j]) / n
for k in range(n):
axis = np.append(axis, axis[j - 1] + k * delta)
# Assert no cell exceeds minimum cell size (e.g. 35 um trace thickness)
#axis = np.delete(axis, np.argwhere(np.ediff1d(axis) <= min_cell) + 1)
axis = np.unique(axis)
axis = np.sort(axis)
###### add 35 micrometer traces
######axis = np.unique(np.concatenate((np.unique(vertices[:,i]), axis)))
mesh[i] = axis
plot = False
if plot:
plt.title("Mesh Sequence")
plt.plot(mesh[0], ".", label="X-axis")
plt.plot(mesh[1], ".", label="Y-axis")
plt.plot(mesh[2], ".", label="Z-axis")
plt.xlabel("Mesh line index")
plt.ylabel("Position")
plt.legend()
plt.grid()
plt.show()
# Delete temporary STL model
#try:
# os.remove("data/tmp/" + stl)
#except OSError:
# pass
return mesh
However, I get the following error:
Traceback (most recent call last):
File "solver_api.py", line 1077, in <module>
mesh = generate_mesh(f_min, f_max, port, pml_x, vertices, stl="model.stl", factor=factor, factor_space=factor_space, fraction=fraction, res_fraction=[6,6,6], cell_ratio=2, n=[3, 3, 3])
File "/home/abc/.local/lib/python3.8/site-packages/numba/core/dispatcher.py", line 468, in _compile_for_args
error_rewrite(e, 'typing')
File "/home/abc/.local/lib/python3.8/site-packages/numba/core/dispatcher.py", line 409, in error_rewrite
raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Type of variable 'ax1' cannot be determined, operation: call $842load_attr.3($848binary_subscr.6, func=$842load_attr.3, args=[Var($848binary_subscr.6, solver_api.py:249)], kws=[('dtype', Var($852load_attr.8, solver_api.py:249))], vararg=None, target=None), location: solver_api.py (249)
File "solver_api.py", line 249:
def generate_mesh(
<source elided>
for i in range(len(mesh)):
ax1 = np.array(mesh[i], dtype=np.float64)
^
How exactly do I resolve this? I couldn’t figure out how to help Numba “determine its variable type”.
