emerge 0.5.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

emerge/_emerge/periodic.py

@@ -1,5 +1,5 @@
-from .cs import Axis, _parse_axis, GCS
-from .selection import FaceSelection, SELECTOR_OBJ
+from .cs import Axis, _parse_axis, GCS, _parse_vector
+from .selection import FaceSelection, SELECTOR_OBJ, Selection
 from .geo import GeoPrism, XYPolygon, Alignment, XYPlate
 from .bc import BoundaryCondition
 from typing import Generator
@@ -21,7 +21,7 @@ def _rotnorm(v: np.ndarray) -> np.ndarray:
     ax = ax/np.linalg.norm(ax)
     return ax
 
-def _pair_selection(f1: FaceSelection, f2: FaceSelection, translation: tuple[float, float, float]):
+def _pair_selection(f1: Selection, f2: Selection, translation: tuple[float, float, float]):
     if len(f1.tags) == 1:
         return [f1,], [f2,]
     c1s = [np.array(c) for c in f1.centers]
@@ -32,8 +32,8 @@ def _pair_selection(f1: FaceSelection, f2: FaceSelection, translation: tuple[flo
     for t1, c1 in zip(f1.tags, c1s):
         for t2, c2 in zip(f2.tags, c2s):
             if np.linalg.norm((c1 + ds)-c2) < 1e-6:
-                f1s.append(FaceSelection([t1,]))
-                f2s.append(FaceSelection([t2,]))
+                f1s.append(Selection([t1,]))
+                f2s.append(Selection([t2,]))
     return f1s, f2s
 
 
@@ -42,15 +42,16 @@ def _pair_selection(f1: FaceSelection, f2: FaceSelection, translation: tuple[flo
 #                 BASE PERIODIC CELL CLASS                #
 ############################################################
 
-
+# TODO: This must be moved to mw physics if possible
 class PeriodicCell:
 
     def __init__(self, 
-                 origins: list[tuple[float, float, float]],
-                 vectors: list[tuple[float, float, float] | Axis]):
-        self.origins: list[tuple[float, float, float]] = origins
+                 origins: list[tuple[float, float, float] | list[float] | np.ndarray | Axis],
+                 vectors: list[tuple[float, float, float] | list[float] | np.ndarray | Axis]):
+
+        self.origins: list[tuple[float, float, float]] = [_parse_vector(origin) for origin in origins] # type: ignore
         self.vectors: list[Axis] = [_parse_axis(vec) for vec in vectors]
-        self.excluded_faces: FaceSelection = None
+        self.excluded_faces: Selection | None = None
         self._bcs: list[Periodic] = []
         self._ports: list[BoundaryCondition] = []
 
@@ -71,7 +72,7 @@ class PeriodicCell:
         """
         raise NotImplementedError('This method is not implemented for this subclass.')
     
-    def cell_data(self) -> Generator[tuple[FaceSelection,FaceSelection,tuple[float, float, float]], None, None]:
+    def cell_data(self) -> Generator[tuple[Selection, Selection, np.ndarray], None, None]:
         """An iterator that yields the two faces of the hex cell plus a cell periodicity vector
 
         Yields:
@@ -84,15 +85,17 @@ class PeriodicCell:
         """
         bcs = []
         for f1, f2, a in self.cell_data():
+            f1_new = f1
+            f2_new = f2
             if self.excluded_faces is not None:
-                f1 = f1 - self.excluded_faces
-                f2 = f2 - self.excluded_faces
-            bcs.append(Periodic(f1, f2, a))
+                f1_new = f1 - self.excluded_faces # type: ignore
+                f2_new = f2 - self.excluded_faces # type: ignore
+            bcs.append(Periodic(f1_new, f2_new, tuple(a)))
         self._bcs = bcs
         return bcs
     
     @property
-    def bcs(self) -> list[Periodic]:
+    def bcs(self) -> list[BoundaryCondition]:
         """Returns a list of Periodic boundary conditions for the given PeriodicCell
 
         Args:
@@ -103,7 +106,7 @@ class PeriodicCell:
         """
         if not self._bcs:
             raise ValueError('Periodic Boundary conditions not generated')
-        return self._bcs + self._ports
+        return self._bcs + self._ports # type: ignore
     
     def set_scanangle(self, theta: float, phi: float, degree: bool = True) -> None:
         """Sets the scanangle for the periodic condition. (0,0) is defined along the Z-axis
@@ -126,8 +129,8 @@ class PeriodicCell:
             bc.uy = uy
             bc.uz = uz
         for port in self._ports:
-            port.scan_theta = theta
-            port.scan_phi = phi
+            port.scan_theta = theta # type: ignore
+            port.scan_phi = phi # type: ignore
 
     def port_face(self, z: float):
         raise NotImplementedError('')
@@ -200,9 +203,9 @@ class RectCell(PeriodicCell):
 class HexCell(PeriodicCell):
 
     def __init__(self,
-                 p1: tuple[float, float, float],
-                 p2: tuple[float, float, float],
-                 p3: tuple[float, float, float]):
+                 point1: tuple[float, float, float],
+                 point2: tuple[float, float, float],
+                 point3: tuple[float, float, float]):
         """Generates a Hexagonal periodic tiling by providing 4 coordinates. The layout of the tiling is as following
         Assuming a hexagon with a single vertext at the top and bottom and two vertices on the left and right faces ⬢
 
@@ -211,7 +214,7 @@ class HexCell(PeriodicCell):
             p2 (tuple[float, float, float]): left face bottom vertex
             p3 (tuple[float, float, float]): bottom vertex
         """
-        p1, p2, p3 = [np.array(p) for p in [p1, p2, p3]]
+        p1, p2, p3 = np.array(point1), np.array(point2), np.array(point3)
         p4 = -p1
         self.p1: np.ndarray = p1
         self.p2: np.ndarray = p2
@@ -240,7 +243,7 @@ class HexCell(PeriodicCell):
         poly = XYPolygon(xs, ys).geo(GCS.displace(0,0,zs[0]))
         return poly
     
-    def cell_data(self) -> Generator[FaceSelection, FaceSelection, np.ndarray]:
+    def cell_data(self) -> Generator[tuple[Selection, Selection, np.ndarray], None, None]:
         nrm = np.linalg.norm
 
         o = self.o1[:-1]
@@ -250,7 +253,7 @@ class HexCell(PeriodicCell):
         f2s = SELECTOR_OBJ.inplane(*self.f12[0], *self.f12[1]).exclude(lambda x, y, z: (nrm(np.array([x,y])+o)>w) or (abs((np.array([x,y])+o) @ n ) > 1e-6))
         vec = - (self.p1 + self.p2)
 
-        for f1, f2 in zip(*_pair_selection(f1s, f2s, vec)):
+        for f1, f2 in zip(*_pair_selection(f1s, f2s, vec)): # type: ignore
             yield f1, f2, vec
 
         o = self.o2[:-1]
@@ -259,7 +262,7 @@ class HexCell(PeriodicCell):
         f1s = SELECTOR_OBJ.inplane(*self.f21[0], *self.f21[1]).exclude(lambda x, y, z: (nrm(np.array([x,y])-o)>w) or (abs((np.array([x,y])-o) @ n ) > 1e-6))
         f2s = SELECTOR_OBJ.inplane(*self.f22[0], *self.f22[1]).exclude(lambda x, y, z: (nrm(np.array([x,y])+o)>w) or (abs((np.array([x,y])+o) @ n ) > 1e-6))
         vec = - (self.p2 + self.p3)
-        for f1, f2 in zip(*_pair_selection(f1s, f2s, vec)):
+        for f1, f2 in zip(*_pair_selection(f1s, f2s, vec)): # type: ignore
             yield f1, f2, vec
         
         o = self.o3[:-1]
@@ -268,18 +271,18 @@ class HexCell(PeriodicCell):
         f1s = SELECTOR_OBJ.inplane(*self.f31[0], *self.f31[1]).exclude(lambda x, y, z: (nrm(np.array([x,y])-o)>w) or (abs((np.array([x,y])-o) @ n ) > 1e-6))
         f2s = SELECTOR_OBJ.inplane(*self.f32[0], *self.f32[1]).exclude(lambda x, y, z: (nrm(np.array([x,y])+o)>w) or (abs((np.array([x,y])+o) @ n ) > 1e-6))
         vec = - (self.p3 - self.p1)
-        for f1, f2 in zip(*_pair_selection(f1s, f2s, vec)):
+        for f1, f2 in zip(*_pair_selection(f1s, f2s, vec)): # type: ignore
             yield f1, f2, vec
 
     def volume(self, 
                z1: float,
                z2: float) -> GeoPrism:
         xs, ys, zs = zip(self.p1, self.p2, self.p3)
-        xs = np.array(xs)
-        ys = np.array(ys)
-        xs = np.concatenate([xs, -xs])
-        ys = np.concatenate([ys, -ys])
-        poly = XYPolygon(xs, ys)
+        xs2 = np.array(xs) # type: ignore
+        ys2 = np.array(ys) # type: ignore
+        xs3 = np.concatenate([xs2, -xs2]) # type: ignore
+        ys3 = np.concatenate([ys2, -ys2]) # type: ignore
+        poly = XYPolygon(xs3, ys3)
         length = z2-z1
         return poly.extrude(length, cs=GCS.displace(0,0,z1))
     

emerge/_emerge/physics/microwave/adaptive_freq.py

@@ -49,14 +49,17 @@ Modification history
 
 """
 
+# ty: ignore
+
 import numpy as np
 from typing import Literal
 from loguru import logger
+
 def cc(z):
     return z.conjugate()
 
 def model(s, poles, residues, d, h):
-    return np.sum(r/(s-p) for p, r in zip(poles, residues)) + d + s*h
+    return sum([r/(s-p) for p, r in zip(poles, residues)]) + d + s*h
 
 def vectfit_step(f: np.ndarray, s: np.ndarray, poles: np.ndarray) -> np.ndarray:
     """
@@ -166,12 +169,12 @@ def calculate_residues(f: np.ndarray, s: np.ndarray, poles: np.ndarray, rcond=-1
     b = np.concatenate((b.real, b.imag))
     cA = np.linalg.cond(A)
     if cA > 1e13:
-        print('Warning!: Ill Conditioned Matrix. Consider scaling the problem down')
-        print('Cond(A)', cA)
+        logger.warning('Warning!: Ill Conditioned Matrix. Consider scaling the problem down')
+        logger.warning('Cond(A)', cA)
     x, residuals, rnk, s = np.linalg.lstsq(A, b, rcond=rcond)
 
     # Recover complex values
-    x = np.complex128(x)
+    x = x.astype(np.complex128)
     for i, ci in enumerate(cindex):
        if ci == 1:
            r1, r2 = x[i:i+2]
@@ -190,7 +193,7 @@ def vectfit_auto(f: np.ndarray,
                  inc_real: bool = False, 
                  loss_ratio: float = 1e-2, 
                  rcond: int = -1, 
-                 track_poles: bool = False) -> tuple[np.ndarray, np.ndarray, float, float]:
+                 track_poles: bool = False) -> tuple[np.ndarray, np.ndarray, float, float, np.ndarray]:
     w = s.imag
     pole_locs = np.linspace(w[0], w[-1], n_poles+2)[1:-1]
     lr = loss_ratio
@@ -208,7 +211,7 @@ def vectfit_auto(f: np.ndarray,
 
     if track_poles:
         return poles, residues, d, h, np.array(poles_list)
-    return poles, residues, d, h
+    return poles, residues, d, h, np.array([])
 
 def vectfit_auto_rescale(f: np.ndarray, s: np.ndarray, 
                          n_poles: int = 10, 
@@ -216,10 +219,10 @@ def vectfit_auto_rescale(f: np.ndarray, s: np.ndarray,
                          inc_real: bool = False, 
                          loss_ratio: float = 1e-2, 
                          rcond: int = -1, 
-                         track_poles: bool = False) -> tuple[np.ndarray, np.ndarray, float, float]:
+                         track_poles: bool = False) -> tuple[np.ndarray, np.ndarray, float, float, np.ndarray]:
     s_scale = abs(s[-1])
     f_scale = abs(f[-1])
-    poles_s, residues_s, d_s, h_s = vectfit_auto(f / f_scale, 
+    poles_s, residues_s, d_s, h_s, tracked_poles = vectfit_auto(f / f_scale, 
                                                  s / s_scale,
                                                   n_poles=n_poles,
                                                   n_iter = n_iter,
@@ -231,7 +234,7 @@ def vectfit_auto_rescale(f: np.ndarray, s: np.ndarray,
     residues = residues_s * f_scale * s_scale
     d = d_s * f_scale
     h = h_s * f_scale / s_scale
-    return poles, residues, d, h
+    return poles, residues, d, h, tracked_poles
 
 
 class SparamModel:
@@ -251,7 +254,7 @@ class SparamModel:
             fdense = np.linspace(min(self.f), max(self.f), max(201, 10*self.f.shape[0]))
             success = False
             for nps in range(1,maxpoles):
-                poles, residues, d, h = vectfit_auto_rescale(Sparam, s, n_poles=nps, inc_real=inc_real)
+                poles, residues, d, h, _ = vectfit_auto_rescale(Sparam, s, n_poles=nps, inc_real=inc_real)
                 self.poles: np.ndarray = poles
                 self.residues: np.ndarray = residues
                 self.d = d
@@ -268,7 +271,7 @@ class SparamModel:
                 logger.warning('Could not model S-parameters. Try a denser grid')
 
         else:
-            poles, residues, d, h = vectfit_auto_rescale(Sparam, s, n_poles=n_poles, inc_real=inc_real)
+            poles, residues, d, h, _ = vectfit_auto_rescale(Sparam, s, n_poles=n_poles, inc_real=inc_real)
             self.poles: np.ndarray = poles
             self.residues: np.ndarray = residues
             self.d = d

emerge/_emerge/physics/microwave/assembly/assembler.py

@@ -124,7 +124,7 @@ class Assembler:
                         sig: np.ndarray,
                         k0: float,
                         port: PortBC,
-                        bcs: list[BoundaryCondition]) -> tuple[np.ndarray, np.ndarray, np.ndarray, NedelecLegrange2]:
+                        bcs: list[BoundaryCondition]) -> tuple[csr_matrix, csr_matrix, np.ndarray, NedelecLegrange2]:
         """Computes the boundary mode analysis matrices
 
         Args:
@@ -194,12 +194,9 @@ class Assembler:
                 tids = nedlegfield.tri_to_field[:, i2]
                 pec_ids.extend(list(tids))
 
-        port._field = nedlegfield
-        port._pece = pec_edges
-        port._pecv = pec_vertices
         # Process all port boundary Conditions
-        pec_ids = set(pec_ids)
-        solve_ids = [i for i in range(nedlegfield.n_field) if i not in pec_ids]
+        pec_ids_set: set[int] = set(pec_ids)
+        solve_ids = [i for i in range(nedlegfield.n_field) if i not in pec_ids_set]
 
         return E, B, np.array(solve_ids), nedlegfield
 
@@ -234,7 +231,7 @@ class Assembler:
 
         mesh = field.mesh
         er = er - 1j*sig/(W0*EPS0)*np.repeat(np.eye(3)[:, :, np.newaxis], er.shape[2], axis=2)
-        is_frequency_dependent: bool = np.any((sig > 0) & (sig < self.conductivity_limit))
+        is_frequency_dependent: bool = np.any((sig > 0) & (sig < self.conductivity_limit)) # type: ignore
         
 
         if cache_matrices and not is_frequency_dependent and self.cached_matrices is not None:
@@ -254,7 +251,7 @@ class Assembler:
 
         # ISOLATE BOUNDARY CONDITIONS TO ASSEMBLE
         pec_bcs: list[PEC] = [bc for bc in bcs if isinstance(bc,PEC)]
-        robin_bcs: list[RectangularWaveguide] = [bc for bc in bcs if isinstance(bc,RobinBC)]
+        robin_bcs: list[RobinBC] = [bc for bc in bcs if isinstance(bc,RobinBC)]
         port_bcs: list[PortBC] = [bc for bc in bcs if isinstance(bc, PortBC)]
         periodic_bcs: list[Periodic] = [bc for bc in bcs if isinstance(bc, Periodic)]
 
@@ -264,7 +261,7 @@ class Assembler:
         
         # Process all PEC Boundary Conditions
         logger.debug('    Implementing PEC Boundary Conditions.')
-        pec_ids = []
+        pec_ids: list[int] = []
         # Conductivity above al imit, consider it all PEC
         for itet in range(field.n_tets):
             if sig[itet] > self.conductivity_limit:
@@ -310,13 +307,13 @@ class Assembler:
                         basis = plane_basis_from_points(mesh.nodes[:,nodes]) + 1e-16
                         ibasis = np.linalg.pinv(basis)
                     if bc._include_force:
-
-                        Bempty, b_p = assemble_robin_bc_excited(field, Bempty, tri_ids, Ufunc, gamma, ibasis, bc.cs.origin, gauss_points)
                         
-                        port_vectors[bc.port_number] += b_p
+                        Bempty, b_p = assemble_robin_bc_excited(field, Bempty, tri_ids, Ufunc, gamma, ibasis, bc.cs.origin, gauss_points) # type: ignore
+                        
+                        port_vectors[bc.port_number] += b_p # type: ignore
                     
                     else:
-                        Bempty = assemble_robin_bc(field, Bempty, tri_ids, gamma)
+                        Bempty = assemble_robin_bc(field, Bempty, tri_ids, gamma) # type: ignore
             B_p = field.generate_csr(Bempty)
             K = K + B_p
         
@@ -326,20 +323,20 @@ class Assembler:
         
         # Periodic BCs
         Pmats = []
-        remove = set()
+        remove: set[int] = set()
         has_periodic = False
 
-        for bc in periodic_bcs:
+        for pbc in periodic_bcs:
             has_periodic = True
-            tri_ids_1 = mesh.get_triangles(bc.face1.tags)
-            edge_ids_1 = mesh.get_edges(bc.face1.tags)
-            tri_ids_2 = mesh.get_triangles(bc.face2.tags)
-            edge_ids_2 = mesh.get_edges(bc.face2.tags)
-            dv = np.array(bc.dv)
+            tri_ids_1 = mesh.get_triangles(pbc.face1.tags)
+            edge_ids_1 = mesh.get_edges(pbc.face1.tags)
+            tri_ids_2 = mesh.get_triangles(pbc.face2.tags)
+            edge_ids_2 = mesh.get_edges(pbc.face2.tags)
+            dv = np.array(pbc.dv)
             linked_tris = pair_coordinates(mesh.tri_centers, tri_ids_1, tri_ids_2, dv, 1e-9)
             linked_edges = pair_coordinates(mesh.edge_centers, edge_ids_1, edge_ids_2, dv, 1e-9)
-            dv = np.array(bc.dv)
-            phi = bc.phi(K0)
+            dv = np.array(pbc.dv)
+            phi = pbc.phi(K0)
             
             Pmat, rows = gen_periodic_matrix(tri_ids_1,
                                        edge_ids_1,
@@ -356,15 +353,15 @@ class Assembler:
             Pmat = Pmats[0]
             for P2 in Pmats[1:]:
                 Pmat = Pmat @ P2
-            remove = np.array(sorted(list(remove)))
+            remove_array = np.sort(np.unique(list(remove)))
             all_indices = np.arange(NF)
-            keep_indices = np.setdiff1d(all_indices, remove)
+            keep_indices = np.setdiff1d(all_indices, remove_array)
             Pmat = Pmat[:,keep_indices]
         else:
             Pmat = None
         
-        pec_ids = set(pec_ids)
-        solve_ids = np.array([i for i in range(E.shape[0]) if i not in pec_ids])
+        pec_ids_set = set(pec_ids)
+        solve_ids = np.array([i for i in range(E.shape[0]) if i not in pec_ids_set])
         
         if has_periodic:
             mask = np.zeros((NF,))
@@ -426,11 +423,11 @@ class Assembler:
         NF = E.shape[0]
 
         pecs: list[PEC] = [bc for bc in bcs if isinstance(bc,PEC)]
-        robin_bcs: list[RectangularWaveguide] = [bc for bc in bcs if isinstance(bc,RobinBC)]
+        robin_bcs: list[RectangularWaveguide] = [bc for bc in bcs if isinstance(bc,RobinBC)] # type: ignore
         periodic: list[Periodic] = [bc for bc in bcs if isinstance(bc, Periodic)]
 
         # Process all PEC Boundary Conditions
-        pec_ids = []
+        pec_ids: list = []
         
         logger.debug('    Implementing PEC Boundary Conditions.')
         
@@ -459,23 +456,30 @@ class Assembler:
         if len(robin_bcs) > 0:
             logger.debug('    Implementing Robin Boundary Conditions.')
         
-        for bc in robin_bcs:
-            for tag in bc.tags:
-                face_tags = [tag,]#bc.tags
-
-                tri_ids = mesh.get_triangles(face_tags)
-                nodes = mesh.get_nodes(face_tags)
-                edge_ids = list(mesh.tri_to_edge[:,tri_ids].flatten())
-
-                gamma = bc.get_gamma(k0)
-                
-                ibasis = bc.get_inv_basis()
-                if ibasis is None:
-                    basis = plane_basis_from_points(mesh.nodes[:,nodes]) + 1e-16
-                    ibasis = np.linalg.pinv(basis)
-                B_p = assemble_robin_bc(field, tri_ids, gamma)
-                if bc._include_stiff:
-                    B = B + B_p
+        if len(robin_bcs) > 0:
+            logger.debug('    Implementing Robin Boundary Conditions.')
+        
+            gauss_points = gaus_quad_tri(4)
+            Bempty = field.empty_tri_matrix()
+            for bc in robin_bcs:
+
+                for tag in bc.tags:
+                    face_tags = [tag,]
+
+                    tri_ids = mesh.get_triangles(face_tags)
+                    nodes = mesh.get_nodes(face_tags)
+                    edge_ids = list(mesh.tri_to_edge[:,tri_ids].flatten())
+
+                    gamma = bc.get_gamma(k0)
+                    
+                    ibasis = bc.get_inv_basis()
+                    if ibasis is None:
+                        basis = plane_basis_from_points(mesh.nodes[:,nodes]) + 1e-16
+                        ibasis = np.linalg.pinv(basis)
+                    
+                    Bempty = assemble_robin_bc(field, Bempty, tri_ids, gamma) # type: ignore
+            B_p = field.generate_csr(Bempty)
+            B = B + B_p
         
         if len(periodic) > 0:
             logger.debug('    Implementing Periodic Boundary Conditions.')
@@ -485,17 +489,17 @@ class Assembler:
         remove = set()
         has_periodic = False
 
-        for bc in periodic:
+        for bcp in periodic:
             has_periodic = True
-            tri_ids_1 = mesh.get_triangles(bc.face1.tags)
-            edge_ids_1 = mesh.get_edges(bc.face1.tags)
-            tri_ids_2 = mesh.get_triangles(bc.face2.tags)
-            edge_ids_2 = mesh.get_edges(bc.face2.tags)
-            dv = np.array(bc.dv)
+            tri_ids_1 = mesh.get_triangles(bcp.face1.tags)
+            edge_ids_1 = mesh.get_edges(bcp.face1.tags)
+            tri_ids_2 = mesh.get_triangles(bcp.face2.tags)
+            edge_ids_2 = mesh.get_edges(bcp.face2.tags)
+            dv = np.array(bcp.dv)
             linked_tris = pair_coordinates(mesh.tri_centers, tri_ids_1, tri_ids_2, dv, 1e-9)
             linked_edges = pair_coordinates(mesh.edge_centers, edge_ids_1, edge_ids_2, dv, 1e-9)
-            dv = np.array(bc.dv)
-            phi = bc.phi(k0)
+            dv = np.array(bcp.dv)
+            phi = bcp.phi(k0)
             
             Pmat, rows = gen_periodic_matrix(tri_ids_1,
                                        edge_ids_1,
@@ -513,15 +517,15 @@ class Assembler:
             for P2 in Pmats[1:]:
                 Pmat = Pmat @ P2
             Pmat = Pmat.tocsr()
-            remove = np.array(sorted(list(remove)))
+            remove_array = np.sort(np.array(list(remove)))
             all_indices = np.arange(NF)
-            keep_indices = np.setdiff1d(all_indices, remove)
+            keep_indices = np.setdiff1d(all_indices, remove_array)
             Pmat = Pmat[:,keep_indices]
         else:
             Pmat = None
         
-        pec_ids = set(pec_ids)
-        solve_ids = np.array([i for i in range(E.shape[0]) if i not in pec_ids])
+        pec_ids_set = set(pec_ids)
+        solve_ids = np.array([i for i in range(E.shape[0]) if i not in pec_ids_set])
         
         if has_periodic:
             mask = np.zeros((NF,))

emerge/_emerge/physics/microwave/assembly/generalized_eigen.py

@@ -17,21 +17,33 @@
 
 import numpy as np
 from ....elements.nedleg2 import NedelecLegrange2
-from scipy.sparse import coo_matrix
+from scipy.sparse import csr_matrix
 from numba_progress import ProgressBar, ProgressBarType
 from ....mth.optimized import local_mapping, matinv, compute_distances
 from numba import c16, types, f8, i8, njit, prange
 
+
+
+############################################################
+#                      FIELD MAPPING                      #
+############################################################
+
 @njit(i8[:,:](i8, i8[:,:], i8[:,:], i8[:,:]), cache=True, nogil=True)
 def local_tri_to_edgeid(itri: int, tris, edges, tri_to_edge) -> np.ndarray:
     global_edge_map = edges[:, tri_to_edge[:,itri]]
     return local_mapping(tris[:, itri], global_edge_map)
 
+
+
+############################################################
+#                     PYTHON INTERFACE                     #
+############################################################
+
 def generelized_eigenvalue_matrix(field: NedelecLegrange2,
                            er: np.ndarray, 
                            ur: np.ndarray,
                            basis: np.ndarray,
-                           k0: float,) -> tuple[coo_matrix, coo_matrix]:
+                           k0: float,) -> tuple[csr_matrix, csr_matrix]:
     
     tris = field.mesh.tris
     edges = field.mesh.edges
@@ -47,11 +59,16 @@ def generelized_eigenvalue_matrix(field: NedelecLegrange2,
     
     nfield = field.n_field
 
-    E = coo_matrix((dataE, (rows, cols)), shape=(nfield, nfield)).tocsr()
-    B = coo_matrix((dataB, (rows, cols)), shape=(nfield, nfield)).tocsr()
+    E = csr_matrix((dataE, (rows, cols)), shape=(nfield, nfield))
+    B = csr_matrix((dataB, (rows, cols)), shape=(nfield, nfield))
 
     return E, B
 
+
+############################################################
+#                   MATRIX MULTIPLICATION                  #
+############################################################
+
 @njit(c16[:,:](c16[:,:], c16[:,:]), cache=True, nogil=True)
 def matmul(a, b):
     out = np.empty((2,b.shape[1]), dtype=np.complex128)
@@ -59,7 +76,10 @@ def matmul(a, b):
     out[1,:] = a[1,0]*b[0,:] + a[1,1]*b[1,:]
     return out
 
-### GAUSS QUADRATURE IMPLEMENTATION
+
+############################################################
+#              GAUSS QUADRATURE IMPLEMENTATION             #
+############################################################
 
 @njit(c16(c16[:], c16[:], types.Array(types.float64, 1, 'A', readonly=True)), cache=True, nogil=True)
 def _gqi(v1, v2, W):
@@ -185,7 +205,12 @@ def _nf2_curl(coeff, coords):
     ys = coords[1,:]
     return b3*(c1*(a2 + b2*xs + c2*ys) - c2*(a1 + b1*xs + c1*ys)) - c3*(b1*(a2 + b2*xs + c2*ys) - b2*(a1 + b1*xs + c1*ys)) - 2*(b1*c2 - b2*c1)*(a3 + b3*xs + c3*ys) + 0*1j
 
-####
+
+############################################################
+#     TRIANGLE BARYCENTRIC COORDINATE LIN. COEFFICIENTS    #
+############################################################
+
+
 @njit(types.Tuple((f8[:], f8[:], f8[:], f8))(f8[:], f8[:]), cache = True, nogil=True)
 def tri_coefficients(vxs, vys):
 
@@ -211,13 +236,23 @@ def tri_coefficients(vxs, vys):
     Cs = np.array([c1, c2, c3])*sign
     return As, Bs, Cs, A
 
-#DPTS = gaus_quad_tri(4).astype(np.float64)
+
+############################################################
+#                    CONSTANT DEFINITION                   #
+############################################################
+
 
 DPTS = np.array([[0.22338159, 0.22338159, 0.22338159, 0.10995174, 0.10995174, 0.10995174],
                         [0.10810302, 0.44594849, 0.44594849, 0.81684757, 0.09157621, 0.09157621],
                         [0.44594849, 0.44594849, 0.10810302, 0.09157621, 0.09157621, 0.81684757],
                         [0.44594849, 0.10810302, 0.44594849, 0.09157621, 0.81684757, 0.09157621]], dtype=np.float64)
 
+
+############################################################
+#                 NUMBA OPTIMIZED ASSEMBLER                #
+############################################################
+
+
 @njit(types.Tuple((c16[:,:], c16[:,:]))(f8[:,:], i8[:,:], c16[:,:], c16[:,:], f8), cache=True, nogil=True)
 def generalized_matrix_GQ(tri_vertices, local_edge_map, Ms, Mm, k0):
     '''Nedelec-2 Triangle stiffness and mass submatrix'''
@@ -402,7 +437,7 @@ def _matrix_builder(nodes, tris, edges, tri_to_field, ur, er, k0, pgb: ProgressB
 
     tri_to_edge = tri_to_field[:3,:]
     
-    for itri in prange(ntritot):
+    for itri in prange(ntritot): # type: ignore
         p = itri*196
         if np.mod(itri,10)==0:
             pgb.update(10)

emerge/_emerge/physics/microwave/assembly/robinbc.py

@@ -252,7 +252,7 @@ def ned2_tri_stiff_force(lcs_vertices, gamma, lcs_Uinc, DPTs):
 def compute_bc_entries_excited(vertices_local, tris, Bmat, Bvec, surf_triangle_indices, gamma, Ulocal_all, DPTs, tri_to_field):
     N = 64
     Niter = surf_triangle_indices.shape[0]
-    for i in prange(Niter):
+    for i in prange(Niter): # type: ignore
         itri = surf_triangle_indices[i]
 
         vertex_ids = tris[:, itri]
@@ -388,7 +388,7 @@ def ned2_tri_stiff(vertices, edge_lengths, gamma):
 def compute_bc_entries(vertices, tris, Bmat, all_edge_lengths, surf_triangle_indices, gamma):
     N = 64
     Niter = surf_triangle_indices.shape[0]
-    for i in prange(Niter):
+    for i in prange(Niter): # type: ignore
         itri = surf_triangle_indices[i]
 
         vertex_ids = tris[:, itri]
@@ -404,7 +404,7 @@ def assemble_robin_bc_excited(field: Nedelec2,
                               Bmat: np.ndarray,
                 surf_triangle_indices: np.ndarray,
                 Ufunc: Callable,
-                gamma: np.ndarray,
+                gamma: complex,
                 local_basis: np.ndarray,
                 origin: np.ndarray,
                 DPTs: np.ndarray):
@@ -424,8 +424,8 @@ def assemble_robin_bc_excited(field: Nedelec2,
 
 def assemble_robin_bc(field: Nedelec2,
                       Bmat: np.ndarray,
-                surf_triangle_indices: np.ndarray,
-                gamma: np.ndarray):
+                    surf_triangle_indices: np.ndarray,
+                    gamma: np.ndarray):
     vertices = field.mesh.nodes
     all_edge_lengths = field.mesh.edge_lengths[field.mesh.tri_to_edge]
     Bmat = compute_bc_entries(vertices, field.mesh.tris, Bmat, all_edge_lengths, surf_triangle_indices, gamma)