diff --git a/src/festim/hydrogen_transport_problem.py b/src/festim/hydrogen_transport_problem.py index 133b9fd78..07f68c730 100644 --- a/src/festim/hydrogen_transport_problem.py +++ b/src/festim/hydrogen_transport_problem.py @@ -106,7 +106,6 @@ class HydrogenTransportProblem(problem.ProblemBase): model formulation (ufl.form.Form): the formulation of the model solver (dolfinx.nls.newton.NewtonSolver): the solver of the model - multispecies (bool): True if the model has more than one species. temperature_fenics (fem.Constant or fem.Function): the temperature of the model as a fenics object (fem.Constant or fem.Function). @@ -256,10 +255,6 @@ def temperature_time_dependent(self): else: return False - @property - def multispecies(self): - return len(self.species) > 1 - @property def species(self) -> list[_species.Species]: return self._species @@ -569,15 +564,13 @@ def define_D_global(self, species): D.interpolate(D_expr) return D, D_expr - def define_function_spaces(self, element_degree=1): - """Creates the function space of the model, creates a mixed element if - model is multispecies. Creates the main solution and previous solution - function u and u_n. Create global DG function spaces of degree 0 and 1 - for the global diffusion coefficient. + def define_function_spaces(self, element_degree: int = 1): + """Creates the function space of the modelw with a mixed element. Creates the + main solution and previous solution function u and u_n. Create global DG + function spaces of degree 0 and 1 for the global diffusion coefficient. Args: - element_degree (int, optional): Degree order for finite element. - Defaults to 1. + element_degree: Degree order for finite element. Defaults to 1. """ element_CG = basix.ufl.element( @@ -593,19 +586,16 @@ def define_function_spaces(self, element_degree=1): basix.LagrangeVariant.equispaced, ) - if not self.multispecies: - element = element_CG - else: - elements = [] - for spe in self.species: - if isinstance(spe, _species.Species): - if spe.mobile: - elements.append(element_CG) - elif self._element_for_traps == "DG": - elements.append(element_DG) - else: - elements.append(element_CG) - element = basix.ufl.mixed_element(elements) + elements = [] + for spe in self.species: + if isinstance(spe, _species.Species): + if spe.mobile: + elements.append(element_CG) + elif self._element_for_traps == "DG": + elements.append(element_DG) + else: + elements.append(element_CG) + element = basix.ufl.mixed_element(elements) self.function_space = fem.functionspace(self.mesh.mesh, element) @@ -629,31 +619,21 @@ def define_function_spaces(self, element_degree=1): self.u_n = fem.Function(self.function_space) def assign_functions_to_species(self): - """Creates the solution, prev solution, test function and - post-processing solution for each species, if model is multispecies, - created a collapsed function space for each species""" - - if not self.multispecies: - sub_solutions = [self.u] - sub_prev_solution = [self.u_n] - sub_test_functions = [ufl.TestFunction(self.function_space)] - self.species[0].sub_function_space = self.function_space - self.species[0].post_processing_solution = self.u - self.species[0].sub_function = self.u - else: - sub_solutions = list(ufl.split(self.u)) - sub_prev_solution = list(ufl.split(self.u_n)) - sub_test_functions = list(ufl.TestFunctions(self.function_space)) - - for idx, spe in enumerate(self.species): - spe.sub_function_space = self.function_space.sub(idx) - spe.sub_function = self.u.sub( - idx - ) # TODO add this to discontinuous class - spe.post_processing_solution = self.u.sub(idx).collapse() - spe.collapsed_function_space, spe.map_sub_to_main_solution = ( - self.function_space.sub(idx).collapse() - ) + """Creates the solution, prev solution, test function and post-processing + solution for each species, as well as a collapsed function space for each + species""" + + sub_solutions = list(ufl.split(self.u)) + sub_prev_solution = list(ufl.split(self.u_n)) + sub_test_functions = list(ufl.TestFunctions(self.function_space)) + + for idx, spe in enumerate(self.species): + spe.sub_function_space = self.function_space.sub(idx) + spe.sub_function = self.u.sub(idx) # TODO add this to discontinuous class + spe.post_processing_solution = self.u.sub(idx).collapse() + spe.collapsed_function_space, spe.map_sub_to_main_solution = ( + self.function_space.sub(idx).collapse() + ) for idx, spe in enumerate(self.species): spe.solution = sub_solutions[idx] @@ -695,19 +675,14 @@ def create_dirichletbc_form(self, bc): the boundary condition for modifying linear systems. """ # create value_fenics - if not self.multispecies: - function_space_value = bc.species.sub_function_space - else: - function_space_value = bc.species.collapsed_function_space - bc.create_value( temperature=self.temperature_fenics, - function_space=function_space_value, + function_space=bc.species.collapsed_function_space, t=self.t, ) # get dofs - if self.multispecies and isinstance(bc.value_fenics, (fem.Function)): + if isinstance(bc.value_fenics, (fem.Function)): function_space_dofs = ( bc.species.sub_function_space, bc.species.collapsed_function_space, @@ -721,16 +696,10 @@ def create_dirichletbc_form(self, bc): ) # create form - if not self.multispecies and isinstance(bc.value_fenics, (fem.Function)): - # no need to pass the functionspace since value_fenics is already a Function - function_space_form = None - else: - function_space_form = bc.species.sub_function_space - form = fem.dirichletbc( value=bc.value_fenics, dofs=bc_dofs, - V=function_space_form, + V=bc.species.sub_function_space, ) return form @@ -779,10 +748,7 @@ def create_initial_conditions(self): function_space_value = None if callable(condition.value): # if bc.value is a callable then need to provide a functionspace - if not self.multispecies: - function_space_value = condition.species.sub_function_space - else: - function_space_value = condition.species.collapsed_function_space + function_space_value = condition.species.collapsed_function_space condition.create_expr_fenics( mesh=self.mesh.mesh, @@ -792,11 +758,8 @@ def create_initial_conditions(self): # assign to previous solution of species entities = self.volume_meshtags.find(condition.volume.id) - if not self.multispecies: - function_to_interpolate_on = condition.species.prev_solution - else: - idx = self.species.index(condition.species) - function_to_interpolate_on = self.u_n.sub(idx) + idx = self.species.index(condition.species) + function_to_interpolate_on = self.u_n.sub(idx) function_to_interpolate_on.interpolate( condition.expr_fenics, cells0=entities @@ -961,12 +924,10 @@ def update_time_dependent_values(self): def update_post_processing_solutions(self): """Updates the post-processing solutions of each species""" - # update post-processing for mixed function space - if self.multispecies: - for spe in self.species: - spe.post_processing_solution.x.array[:] = self.u.x.array[ - spe.map_sub_to_main_solution - ] + for spe in self.species: + spe.post_processing_solution.x.array[:] = self.u.x.array[ + spe.map_sub_to_main_solution + ] def post_processing(self): """Post processes the model""" @@ -1054,26 +1015,16 @@ def post_processing(self): # computing dofs at each time step is costly so storing it in the export if export._dofs is None: - # if multispecies then we have a mixed function space - # (mixed element) - if self.multispecies: - index = self.species.index(export.field) - V0, export._dofs = self.u.function_space.sub(index).collapse() - coords = V0.tabulate_dof_coordinates()[:, 0] - export._sort_coords = np.argsort(coords) - x = coords[export._sort_coords] - # if not multispecies then we have a single function space - else: - coords = self.u.function_space.tabulate_dof_coordinates()[:, 0] - export._sort_coords = np.argsort(coords) - x = coords[export._sort_coords] + index = self.species.index(export.field) + V0, export._dofs = self.u.function_space.sub(index).collapse() + coords = V0.tabulate_dof_coordinates()[:, 0] + export._sort_coords = np.argsort(coords) + x = coords[export._sort_coords] export.x = x - if self.multispecies: - c = self.u.x.array[export._dofs][export._sort_coords] - else: - c = self.u.x.array[export._dofs] + c = self.u.x.array[export._dofs][export._sort_coords] + export.data.append(c) export.t.append(float(self.t)) @@ -1666,8 +1617,8 @@ def create_solver(self): self.solver.convergence_criterion = self.settings.convergence_criterion self.solver.atol = ( self.settings.atol - if not callable(self.settings.rtol) - else self.settings.rtol(float(self.t)) + if not callable(self.settings.atol) + else self.settings.atol(float(self.t)) ) self.solver.rtol = ( self.settings.rtol @@ -2070,6 +2021,7 @@ def update_post_processing_solutions(self): """Updates the post-processing solutions after each time step""" # need to compute c = theta * K_S # this expression is stored in species.dg_expr + for spe in self.species: if not spe.mobile: continue @@ -2085,12 +2037,6 @@ def create_dirichletbc_form(self, bc: festim.FixedConcentrationBC): dolfinx.fem.bcs.DirichletBC: A representation of the boundary condition for modifying linear systems. """ - # create value_fenics - if not self.multispecies: - function_space_value = bc.species.sub_function_space - else: - function_space_value = bc.species.collapsed_function_space - # create K_S function Q0 = fem.functionspace(self.mesh.mesh, ("DG", 0)) K_S0 = fem.Function(Q0) @@ -2102,15 +2048,16 @@ def create_dirichletbc_form(self, bc: festim.FixedConcentrationBC): K_S = K_S0 * ufl.exp(-E_KS / (festim.k_B * self.temperature_fenics)) + # create value_fenics bc.create_value( temperature=self.temperature_fenics, - function_space=function_space_value, + function_space=bc.species.collapsed_function_space, t=self.t, K_S=K_S, ) # get dofs - if self.multispecies and isinstance(bc.value_fenics, (fem.Function)): + if isinstance(bc.value_fenics, fem.Function): function_space_dofs = ( bc.species.sub_function_space, bc.species.collapsed_function_space, @@ -2124,16 +2071,10 @@ def create_dirichletbc_form(self, bc: festim.FixedConcentrationBC): ) # create form - if not self.multispecies and isinstance(bc.value_fenics, (fem.Function)): - # no need to pass the functionspace since value_fenics is already a Function - function_space_form = None - else: - function_space_form = bc.species.sub_function_space - form = fem.dirichletbc( value=bc.value_fenics, dofs=bc_dofs, - V=function_space_form, + V=bc.species.sub_function_space, ) return form diff --git a/test/test_dirichlet_bc.py b/test/test_dirichlet_bc.py index ff399a7ef..0dd366644 100644 --- a/test/test_dirichlet_bc.py +++ b/test/test_dirichlet_bc.py @@ -73,10 +73,11 @@ def value(x, t): my_model.define_function_spaces() my_model.define_meshtags_and_measures() + my_model.assign_functions_to_species() T = fem.Constant(my_model.mesh.mesh, 550.0) t = fem.Constant(my_model.mesh.mesh, 0.0) - bc.create_value(my_model.function_space, T, t) + bc.create_value(species.collapsed_function_space, T, t) # check that the value_fenics attribute is set correctly assert isinstance(bc.value_fenics, fem.Function) @@ -116,10 +117,11 @@ def value(x, t, T): my_model.define_function_spaces() my_model.define_meshtags_and_measures() + my_model.assign_functions_to_species() T = fem.Constant(my_model.mesh.mesh, 550.0) t = fem.Constant(my_model.mesh.mesh, 0.0) - bc.create_value(my_model.function_space, T, t) + bc.create_value(species.collapsed_function_space, T, t) # check that the value_fenics attribute is set correctly assert isinstance(bc.value_fenics, fem.Function) @@ -205,12 +207,13 @@ def value(x): my_model.define_function_spaces() my_model.define_meshtags_and_measures() + my_model.assign_functions_to_species() T = fem.Constant(my_model.mesh.mesh, 550.0) t = fem.Constant(my_model.mesh.mesh, 0.0) # TEST - bc.create_value(my_model.function_space, T, t) + bc.create_value(species.collapsed_function_space, T, t) # check that the value_fenics attribute is set correctly assert isinstance(bc.value_fenics, fem.Function) diff --git a/test/test_h_transport_problem.py b/test/test_h_transport_problem.py index 6a92ef348..393ccb511 100644 --- a/test/test_h_transport_problem.py +++ b/test/test_h_transport_problem.py @@ -460,6 +460,7 @@ def test_post_processing_update_D_global(): ) my_model.define_function_spaces() + my_model.assign_functions_to_species() my_model.define_meshtags_and_measures() my_model.t = fem.Constant(my_model.mesh.mesh, 1.0) my_model.define_temperature() @@ -510,6 +511,7 @@ def test_post_processing_update_D_global_2(): ) my_model.define_function_spaces() + my_model.assign_functions_to_species() my_model.define_meshtags_and_measures() my_model.t = fem.Constant(my_model.mesh.mesh, 1.0) my_model.define_temperature() @@ -547,6 +549,7 @@ def test_post_processing_update_D_global_volume_1(): ) my_model.define_function_spaces() + my_model.assign_functions_to_species() my_model.define_meshtags_and_measures() my_model.t = fem.Constant(my_model.mesh.mesh, 1.0) my_model.define_temperature() @@ -584,6 +587,7 @@ def test_post_processing_update_D_global_volume_2(): ) my_model.define_function_spaces() + my_model.assign_functions_to_species() my_model.define_meshtags_and_measures() my_model.t = fem.Constant(my_model.mesh.mesh, 1.0) my_model.define_temperature() @@ -867,9 +871,9 @@ def test_create_initial_conditions_expr_fenics(input_value, expected_value): # RUN my_model.initialise() - + prev_solution = my_model.u_n.sub(0) assert np.isclose( - my_model.species[0].prev_solution.x.petsc_vec.array[-1], + prev_solution.x.petsc_vec.array[-1], expected_value, ) @@ -1344,3 +1348,28 @@ def test_exports_cyl_sph_coord_system_raise_not_implemented(coord_sys): match=f"Derived quantity exports are not implemented for {coord_sys} meshes", ): my_model.initialise_exports() + + +def test_traps_with_CG_elements(): + """Test that when creating species from traps, if the element for traps is set to + CG, the created species has a CG element as well""" + + # BUILD + my_mesh = F.Mesh1D(np.linspace(0, 1, num=11)) + my_mat = F.Material(D_0=1.5, E_D=0.1, name="my_mat") + my_vol = F.VolumeSubdomain1D(id=1, borders=[0, 1], material=my_mat) + trap = F.Species(name="trap1", mobile=False) + + my_model = F.HydrogenTransportProblem( + mesh=my_mesh, + subdomains=[my_vol], + species=[trap], + ) + + my_model._element_for_traps = "CG" + + my_model.create_species_from_traps() + my_model.define_function_spaces(element_degree=1) + + # TEST + assert my_model.function_space.sub(0).element.basix_element.discontinuous is False diff --git a/test/test_initial_condition.py b/test/test_initial_condition.py index b5878abb1..116bf0359 100644 --- a/test/test_initial_condition.py +++ b/test/test_initial_condition.py @@ -186,7 +186,7 @@ def f(x): my_problem.initialise() # test that the initial condition is correct - u_prev = H.prev_solution + u_prev = my_problem.u_n.sub(0) np.testing.assert_allclose(u_ref.x.array, u_prev.x.array, atol=1e-14) diff --git a/test/test_species.py b/test/test_species.py index 75ad932c8..289d4c7e0 100644 --- a/test/test_species.py +++ b/test/test_species.py @@ -6,6 +6,7 @@ import ufl from dolfinx.fem import Function, functionspace from dolfinx.mesh import create_unit_cube +from ufl.indexed import Indexed import festim as F @@ -22,7 +23,7 @@ def test_assign_functions_to_species(): mesh=mesh, species=[ F.Species(name="H"), - # F.Species(name="Trap"), + F.Species(name="Trap"), ], ) model.define_function_spaces() @@ -32,9 +33,9 @@ def test_assign_functions_to_species(): assert spe.solution is not None assert spe.prev_solution is not None assert spe.test_function is not None - assert isinstance(spe.solution, dolfinx.fem.Function) - assert isinstance(spe.prev_solution, dolfinx.fem.Function) - assert isinstance(spe.test_function, ufl.Argument) + assert isinstance(spe.solution, ufl.indexed.Indexed) + assert isinstance(spe.prev_solution, ufl.indexed.Indexed) + assert isinstance(spe.test_function, ufl.indexed.Indexed) def test_species_repr_and_str():