From 15bd00717b9488441298d3036b45c8f09ebb34fd Mon Sep 17 00:00:00 2001 From: Matthew Scroggs Date: Thu, 24 Oct 2024 17:01:02 +0100 Subject: [PATCH] mypy --- symfem/elements/guzman_neilan.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/symfem/elements/guzman_neilan.py b/symfem/elements/guzman_neilan.py index 7c0113e..ea2081a 100644 --- a/symfem/elements/guzman_neilan.py +++ b/symfem/elements/guzman_neilan.py @@ -72,15 +72,15 @@ def get_sub_cells(reference): ] -def bubbles(reference: Reference) -> typing.List[AnyFunction]: +def bubbles(reference: Reference) -> typing.List[VectorFunction]: """Generate divergence-free bubbles.""" br_bubbles = [] - p = Lagrange(reference, 1, variant="equispaced") + p1 = Lagrange(reference, 1, variant="equispaced") for i in range(reference.sub_entity_count(reference.tdim - 1)): sub_e = reference.sub_entity(reference.tdim - 1, i) - bubble = 1 + bubble = ScalarFunction(1) for j in reference.sub_entities(reference.tdim - 1)[i]: - bubble *= p.get_basis_function(j) + bubble *= p1.get_basis_function(j) br_bubbles.append(VectorFunction(tuple(bubble * j for j in sub_e.normal()))) sub_cells = get_sub_cells(reference) @@ -147,11 +147,10 @@ def bubbles(reference: Reference) -> typing.List[AnyFunction]: assert s_mat @ solution == sympy.Matrix(aim) coeffs = list(solution) - bubble: AnyFunction = f - for i, j in zip(coeffs, sub_basis): - bubble -= i * j + for coeff, basis_f in zip(coeffs, sub_basis): + f -= coeff * basis_f + bubbles.append(f) - bubbles.append(bubble) return bubbles