fix helmholtz problem

Author: GiovanniCanali

pina/_src/equation/equation_factory.py

@@ -382,15 +382,15 @@ class Helmholtz(Equation):  # pylint: disable=R0903
 
             \Delta u + k u - f = 0
 
-    Here, :math:`k` is a parameter of the equation, while :math:`f` is the
+    Here, :math:`k` is the squared wavenumber, while :math:`f` is the
     forcing term.
     """
 
     def __init__(self, k, forcing_term):
         """
         Initialization of the :class:`Helmholtz` class.
 
-        :param k: The parameter of the equation.
+        :param k: The squared wavenumber.
         :type k: float | int
         :param Callable forcing_term: The forcing field function, taking as
             input the points on which evaluation is required.

pina/_src/problem/zoo/helmholtz.py

@@ -37,42 +37,51 @@ class HelmholtzProblem(SpatialProblem):
         "boundary": Condition(domain="boundary", equation=FixedValue(0.0)),
     }
 
-    def __init__(self, alpha=3.0):
+    def __init__(self, k=1.0, alpha_x=1, alpha_y=4):
         """
         Initialization of the :class:`HelmholtzProblem` class.
 
-        :param alpha: Parameter of the forcing term. Default is 3.0.
-        :type alpha: float | int
+        :param k: The squared wavenumber. Default is 1.0.
+        :type k: float | int
+        :param int alpha_x: The frequency in the x-direction. Default is 1.
+        :param int alpha_y: The frequency in the y-direction. Default is 4.
         """
         super().__init__()
-        check_consistency(alpha, (int, float))
-        self.alpha = alpha
+        check_consistency(k, (int, float))
+        check_consistency(alpha_x, int)
+        check_consistency(alpha_y, int)
+        self.k = k
+        self.alpha_x = alpha_x
+        self.alpha_y = alpha_y
 
         def forcing_term(input_):
             """
             Implementation of the forcing term.
             """
+            x, y, pi = input_["x"], input_["y"], torch.pi
+            factor = (self.alpha_x**2 + self.alpha_y**2) * pi**2
             return (
-                (1 - 2 * (self.alpha * torch.pi) ** 2)
-                * torch.sin(self.alpha * torch.pi * input_.extract("x"))
-                * torch.sin(self.alpha * torch.pi * input_.extract("y"))
+                (self.k - factor)
+                * torch.sin(self.alpha_x * pi * x)
+                * torch.sin(self.alpha_y * pi * y)
             )
 
         self.conditions["D"] = Condition(
             domain="D",
-            equation=Helmholtz(self.alpha, forcing_term),
+            equation=Helmholtz(self.k, forcing_term),
         )
 
     def solution(self, pts):
         """
         Implementation of the analytical solution of the Helmholtz problem.
 
         :param LabelTensor pts: Points where the solution is evaluated.
-        :return: The analytical solution of the Poisson problem.
+        :return: The analytical solution of the Helmholtz problem.
         :rtype: LabelTensor
         """
-        sol = torch.sin(self.alpha * torch.pi * pts.extract("x")) * torch.sin(
-            self.alpha * torch.pi * pts.extract("y")
+        x, y, pi = pts["x"], pts["y"], torch.pi
+        sol = torch.sin(self.alpha_x * pi * x) * torch.sin(
+            self.alpha_y * pi * y
         )
         sol.labels = self.output_variables
         return sol

tests/test_problem_zoo/test_helmholtz.py

@@ -3,15 +3,17 @@
 from pina.problem import SpatialProblem
 
 
-@pytest.mark.parametrize("alpha", [1.5, 3])
-def test_constructor(alpha):
+@pytest.mark.parametrize("k", [1.5, 3])
+@pytest.mark.parametrize("alpha_x", [1, 3])
+@pytest.mark.parametrize("alpha_y", [1, 3])
+def test_constructor(k, alpha_x, alpha_y):
 
-    problem = HelmholtzProblem(alpha=alpha)
+    problem = HelmholtzProblem(k=k, alpha_x=alpha_x, alpha_y=alpha_y)
     problem.discretise_domain(n=10, mode="random", domains="all")
     assert problem.are_all_domains_discretised
     assert isinstance(problem, SpatialProblem)
     assert hasattr(problem, "conditions")
     assert isinstance(problem.conditions, dict)
 
     with pytest.raises(ValueError):
-        HelmholtzProblem(alpha="invalid")
+        HelmholtzProblem(k=1, alpha_x=1.5, alpha_y=1)