From d6a14f0a146854a0b6d5a94af1f5bcaa47c6db90 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Thu, 22 May 2025 11:13:35 +0300
Subject: [PATCH 01/19] update pde class to enable multiple inputs for the PDE
 model that is based on FEniCS

---
 cuqipy_fenics/pde.py | 103 +++++++++++++++++++++++--------------------
 1 file changed, 54 insertions(+), 49 deletions(-)

diff --git a/cuqipy_fenics/pde.py b/cuqipy_fenics/pde.py
index 757a9ff..863f597 100644
--- a/cuqipy_fenics/pde.py
+++ b/cuqipy_fenics/pde.py
@@ -2,6 +2,7 @@
 from abc import ABC, abstractmethod
 from cuqi.pde import PDE
 from cuqi.array import CUQIarray
+from cuqi.utilities import get_non_default_args
 import dolfin as dl
 from copy import copy
 import warnings
@@ -146,9 +147,13 @@ def __init__(self, PDE_form, mesh, solution_function_space,
         for key, value in linalg_solve_kwargs.items():
             self._solver.parameters[key] = value
 
-        # Initialize the parameter
-        self.parameter = dl.Function(self.parameter_function_space)
-
+        # Initialize the parameter (one or more)
+        # If only one parameter is passed, it is converted to a list
+        if not isinstance(self.parameter_function_space, (list, tuple)):
+            parameter_function_space_list = [self.parameter_function_space]
+        self.parameter= {}
+        for i, k in enumerate(self._non_default_args):
+            self.parameter[k] = dl.Function(parameter_function_space_list[i])
 
     @property
     def parameter(self):
@@ -169,7 +174,8 @@ def parameter(self, value):
         # Subsequent times setting the parameter (avoid assigning the parameter
         # to new object, set parameter array in place instead)
         elif self._is_parameter_new(value):
-            self._parameter.vector().set_local(value.vector().get_local())
+            for key in self._non_default_args:
+                self._parameter[key].vector().set_local(value[key].vector().get_local())
             # The operator in the solver is no longer valid
             self._flags["is_operator_valid"] = False
 
@@ -179,6 +185,15 @@ def parameter(self, value):
         self._gradient = None
         self.rhs = None
 
+    @property
+    def parameter_args(self):
+        """Get the args form of the parameter"""
+        args = list(self.parameter.values())
+        return args
+
+    @property
+    def _non_default_args(self):
+        return get_non_default_args(self.PDE_form)[:-2] # Exclude the last two arguments (u and p) from the list of non-default args
 
     @property
     def forward_solution(self):
@@ -243,10 +258,14 @@ def observation_operator(self):
     @observation_operator.setter
     def observation_operator(self, value):
         """ Set the observation operator """
-        self._observation_operator = self._create_observation_operator(value)
+        if value == None or callable(value):
+            self._observation_operator = value
+        else:
+            raise NotImplementedError(
+                "observation_operator must be a callable function or None")
 
     @abstractmethod
-    def assemble(self,parameter):
+    def assemble(self, *args, **kwargs):
         """ Assemble the PDE weak form """
         raise NotImplementedError
 
@@ -265,34 +284,34 @@ def gradient_wrt_parameter(self):
         """ Compute gradient of the PDE weak form w.r.t. the parameter"""
         raise NotImplementedError 
 
-    @abstractmethod
-    def _create_observation_operator(self, observation_operator):
-        raise NotImplementedError
-
     def _is_parameter_new(self, input_parameter):
         """ A helper function to check if the `input_parameter` is different 
         from the current parameter (cached in self._parameter). """
 
-        if hasattr(self, '_parameter') \
-            and np.allclose(self._parameter.vector().get_local(),
-                            input_parameter.vector().get_local(),
-                            atol=dl.DOLFIN_EPS, rtol=dl.DOLFIN_EPS):
-            return False
-        else:
+        if not hasattr(self, '_parameter'):
             return True
 
+        is_new = False
+        for key in self._non_default_args:
+            if not np.allclose(self._parameter[key].vector().get_local(),
+                                    input_parameter[key].vector().get_local(),
+                                    atol=dl.DOLFIN_EPS, rtol=dl.DOLFIN_EPS):
+                is_new = True
+        return is_new
 
 class SteadyStateLinearFEniCSPDE(FEniCSPDE):
     """ Class representation of steady state linear PDEs defined in FEniCS. It accepts the same arguments as the base class `cuqipy_fenics.pde.FEniCSPDE`."""
 
-    def assemble(self, parameter=None):
+    def assemble(self, *args, **kwargs):
         self._solution_trial_function = dl.TrialFunction(
             self.solution_function_space)
         self._solution_test_function = dl.TestFunction(
             self.solution_function_space)
 
-        if parameter is not None:
-            self.parameter = parameter
+        kwargs = self._parse_args_add_to_kwargs(
+            *args, map_name="assemble", **kwargs
+        )
+        self.parameter = kwargs
 
         # Either assemble the lhs and rhs forms separately or the full PDE form
         if self.lhs_form is not None:
@@ -328,11 +347,10 @@ def _assemble_full(self):
                 and self._flags["is_operator_valid"] and\
                 self.rhs is not None:
             return
-
-        diff_op = dl.lhs(self.PDE_form(self.parameter,
+        diff_op = dl.lhs(self.PDE_form(*self.parameter_args,
                                        self._solution_trial_function,
                                        self._solution_test_function))
-        self.rhs = dl.rhs(self.PDE_form(self.parameter,
+        self.rhs = dl.rhs(self.PDE_form(*self.parameter_args,
                                         self._solution_trial_function,
                                         self._solution_test_function))
 
@@ -353,7 +371,7 @@ def _assemble_lhs(self):
                 and self._flags["is_operator_valid"]:
             return
 
-        diff_op = dl.assemble(self.lhs_form(self.parameter,
+        diff_op = dl.assemble(self.lhs_form(*self.parameter_args,
                                             self._solution_trial_function,
                                             self._solution_test_function))
 
@@ -367,7 +385,7 @@ def _assemble_rhs(self):
                 and self.rhs is not None:
             return
 
-        self.rhs = dl.assemble(self.rhs_form(self.parameter,
+        self.rhs = dl.assemble(self.rhs_form(*self.parameter_args,
                                              self._solution_test_function))
         for bc in self._dirichlet_bcs: bc.apply(self.rhs)
 
@@ -381,9 +399,9 @@ def observe(self,PDE_solution_fun):
         if self.observation_operator is None: 
             return PDE_solution_fun
         else:
-            return self._apply_obs_op(self.parameter, PDE_solution_fun)
+            return self._apply_obs_op(*self.parameter_args, PDE_solution_fun)
 
-    def gradient_wrt_parameter(self, direction, wrt, **kwargs):
+    def gradient_wrt_parameter(self, direction, *args, **kwargs):
         """ Compute the gradient of the PDE with respect to the parameter
 
         Note: This implementation is largely based on the code:
@@ -391,6 +409,10 @@ def gradient_wrt_parameter(self, direction, wrt, **kwargs):
 
         See also: Gunzburger, M. D. (2002). Perspectives in flow control and optimization. Society for Industrial and Applied Mathematics, for adjoint based derivative derivation. 
         """
+
+        kwargs = self._parse_args_add_to_kwargs(
+            *args, map_name="gradient_wrt_parameter", **kwargs
+        )
         # Raise an error if the adjoint boundary conditions are not provided
         if self._adjoint_dirichlet_bcs is None:
             raise ValueError(
@@ -402,8 +424,8 @@ def gradient_wrt_parameter(self, direction, wrt, **kwargs):
 
         # Compute forward solution
         # TODO: Use stored forward solution if available and wrt == self.parameter
-        self.parameter = wrt
-        self.assemble()
+        self.parameter = kwargs
+        self.assemble(*self.parameter_args)
         self.forward_solution, _ = self.solve()
 
         # Compute adjoint solution
@@ -411,7 +433,7 @@ def gradient_wrt_parameter(self, direction, wrt, **kwargs):
         test_solution = dl.TestFunction(self.solution_function_space)
 
         # note: temp_form is a weak form used for building the adjoint operator
-        temp_form = self.PDE_form(wrt, self.forward_solution, trial_adjoint)
+        temp_form = self.PDE_form(*self.parameter_args, self.forward_solution, trial_adjoint)
         adjoint_form = dl.derivative(
             temp_form, self.forward_solution, test_solution)
 
@@ -431,8 +453,8 @@ def gradient_wrt_parameter(self, direction, wrt, **kwargs):
 
         # Compute gradient
         # note: temp_form is a weak form used for building the gradient
-        temp_form = self.PDE_form(wrt, self.forward_solution, adjoint)
-        gradient_form = dl.derivative(temp_form, wrt, test_parameter)
+        temp_form = self.PDE_form(*self.parameter_args, self.forward_solution, adjoint)
+        gradient_form = dl.derivative(temp_form, *self.parameter_args, test_parameter)
         gradient = dl.Function(self.parameter_function_space)
         dl.assemble(gradient_form, tensor=gradient.vector())
         return gradient
@@ -450,21 +472,4 @@ def _apply_obs_op(self, PDE_parameter_fun, PDE_solution_fun,):
             raise NotImplementedError("obs_op output must be a number, a numpy array or a ufl.algebra.Operator type")
     
 
-    def _create_observation_operator(self, observation_operator):
-        """
-        """
-        if observation_operator == 'potential':
-            observation_operator = lambda m, u: u 
-        elif observation_operator == 'gradu_squared':
-            observation_operator = lambda m, u: dl.inner(dl.grad(u),dl.grad(u))
-        elif observation_operator == 'power_density':
-            observation_operator = lambda m, u: m*dl.inner(dl.grad(u),dl.grad(u))
-        elif observation_operator == 'sigma_u':
-            observation_operator = lambda m, u: m*u
-        elif observation_operator == 'sigma_norm_gradu':
-            observation_operator = lambda m, u: m*dl.sqrt(dl.inner(dl.grad(u),dl.grad(u)))
-        elif observation_operator == None or callable(observation_operator):
-            observation_operator = observation_operator
-        else:
-            raise NotImplementedError
-        return observation_operator
+

From 461c5d62d2c8dbecc055ca757e65046179cf7baf Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Tue, 27 May 2025 21:04:48 +0300
Subject: [PATCH 02/19] update PDE_form and get_non_default_args properties

---
 cuqipy_fenics/pde.py | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/cuqipy_fenics/pde.py b/cuqipy_fenics/pde.py
index 863f597..7c32aa3 100644
--- a/cuqipy_fenics/pde.py
+++ b/cuqipy_fenics/pde.py
@@ -6,6 +6,7 @@
 import dolfin as dl
 from copy import copy
 import warnings
+from functools import partial
 from .utilities import _LazyUFLLoader
 ufl = _LazyUFLLoader()
 
@@ -151,6 +152,8 @@ def __init__(self, PDE_form, mesh, solution_function_space,
         # If only one parameter is passed, it is converted to a list
         if not isinstance(self.parameter_function_space, (list, tuple)):
             parameter_function_space_list = [self.parameter_function_space]
+        else:
+            parameter_function_space_list = self.parameter_function_space
         self.parameter= {}
         for i, k in enumerate(self._non_default_args):
             self.parameter[k] = dl.Function(parameter_function_space_list[i])
@@ -193,7 +196,11 @@ def parameter_args(self):
 
     @property
     def _non_default_args(self):
-        return get_non_default_args(self.PDE_form)[:-2] # Exclude the last two arguments (u and p) from the list of non-default args
+        form = self._form
+        if isinstance(self._form, tuple):
+            # extract non-default args from the lhs first form
+            form = self._form[0]
+        return get_non_default_args(form)[:-2] # Exclude the last two arguments (u and p) from the list of non-default args
 
     @property
     def forward_solution(self):
@@ -204,7 +211,24 @@ def forward_solution(self):
     def PDE_form(self):
         """ Get the PDE form """
         if isinstance(self._form, tuple):
-            return lambda m, u, p: self._form[0](m, u, p) - self._form[1](m, p)
+            # Create a string for the lambda function that represents the PDE form
+            form_str = (
+                "lambda form_lhs, form_rhs, "
+                + ", ".join(self._non_default_args)
+                + ", u, p: form_lhs("
+                + ", ".join(self._non_default_args)
+                + ", u, p) - form_rhs("
+                + ", ".join(self._non_default_args)
+                + ", p)"
+            )
+            # Create a lambda function that represents the PDE form
+            form = eval(form_str)
+
+            # partial evaluation of the form 
+            form_partial = partial(form, form_lhs=self._form[0],
+                                      form_rhs=self._form[1])
+
+            return form_partial
         else:
             return self._form
 

From 0514b9d7af228998484466a7ceb390b2d9cbc638 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Tue, 27 May 2025 21:05:40 +0300
Subject: [PATCH 03/19] add tests for multiple inputs case

---
 tests/test_PDEModel.py | 89 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 89 insertions(+)

diff --git a/tests/test_PDEModel.py b/tests/test_PDEModel.py
index 910419f..50362e4 100644
--- a/tests/test_PDEModel.py
+++ b/tests/test_PDEModel.py
@@ -404,3 +404,92 @@ def test_observation_operator_setter():
     # Check that the solutions are the same
     assert np.allclose(u1_obs, u2_obs) and np.allclose(u2_obs, u3_obs)
     assert len(u1_obs) == 5
+
+class PoissonMultipleInputs:
+    """Define the variational PDE problem for the Poisson equation with multiple
+     unknowns in two ways: as a full form, and as lhs and rhs forms"""
+
+    def __init__(self):
+
+        # Create the mesh and define function spaces for the solution and the
+        # parameter
+        self.mesh = dl.UnitIntervalMesh(10)
+
+        # Define the boundary condition
+        self.bc_value = dl.Constant(0.0)
+
+    @property
+    def form(self):
+        return lambda m, source_term, u, v:\
+            ufl.exp(m)*ufl.inner(ufl.grad(u), ufl.grad(v))*ufl.dx\
+            + source_term*v*ufl.dx
+
+    @property
+    def lhs_form(self):
+        return lambda m, source_term, u, v:\
+            ufl.exp(m)*ufl.inner(ufl.grad(u), ufl.grad(v))*ufl.dx
+
+    @property
+    def rhs_form(self):
+        return lambda m, source_term, v: -source_term*v*ufl.dx
+
+    @property
+    def solution_function_space(self):
+        return dl.FunctionSpace(self.mesh, "Lagrange", 2)
+
+    @property
+    def parameter_function_space(self):
+        return [dl.FunctionSpace(self.mesh, "Lagrange", 1), dl.FunctionSpace(self.mesh, "Lagrange", 1)]
+
+    @property
+    def bcs(self):
+        if not hasattr(self, "_bcs") or self._bcs is None:
+            self._bcs = dl.DirichletBC(self.solution_function_space,
+                                      self.bc_value, "on_boundary")
+        return self._bcs
+
+    @bcs.setter
+    def bcs(self, bcs):
+        self._bcs = bcs
+
+def test_form_multiple_inputs():
+    """Test creating PDEModel with full form, and with lhs and rhs forms with multiple inputs"""
+    # Create the variational problem
+    poisson = PoissonMultipleInputs()
+
+    # Set up model parameters
+    m = dl.Function(poisson.parameter_function_space[0])
+    m.vector()[:] = 1.0
+    source_term = dl.Function(poisson.parameter_function_space[1])
+    source_term.interpolate(dl.Expression('x[0]', degree=1))
+
+    # Create a PDE object with full form
+    PDE_with_full_form = cuqipy_fenics.pde.SteadyStateLinearFEniCSPDE(
+        poisson.form,
+        poisson.mesh,
+        poisson.solution_function_space,
+        poisson.parameter_function_space,
+        poisson.bcs)
+
+    # Solve the PDE
+    PDE_with_full_form.assemble(m, source_term)
+    PDE_with_full_form.assemble(m=m, source_term=source_term)    
+    u1, info = PDE_with_full_form.solve()
+
+    # Create a PDE object with lhs and rhs forms
+    PDE_with_lhs_rhs_forms = cuqipy_fenics.pde.SteadyStateLinearFEniCSPDE(
+        (poisson.lhs_form, poisson.rhs_form),
+        poisson.mesh,
+        poisson.solution_function_space,
+        poisson.parameter_function_space,
+        poisson.bcs)
+    # Solve the PDE
+    PDE_with_lhs_rhs_forms.assemble(m, source_term)
+    u2, info = PDE_with_lhs_rhs_forms.solve()
+
+    # Check that the solutions are the same
+    assert np.allclose(u1.vector().get_local(), u2.vector().get_local())
+
+    # Check passing wrong parameter_function_space raises an error
+    with pytest.raises(ValueError, match= r"assemble input is specified by keywords arguments \['m2', 'source_term'\] that does not match the non_default_args of assemble \['m', 'source_term'\]"):
+        PDE_with_full_form.assemble(m2=m, source_term=source_term)
\ No newline at end of file

From 9b19a0fd1e85121611fba38dc804a09748dfa14a Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Fri, 30 May 2025 10:41:15 +0300
Subject: [PATCH 04/19] Update gradient computation to incorporate multiple
 parameters

---
 cuqipy_fenics/pde.py | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/cuqipy_fenics/pde.py b/cuqipy_fenics/pde.py
index 7c32aa3..a233342 100644
--- a/cuqipy_fenics/pde.py
+++ b/cuqipy_fenics/pde.py
@@ -453,7 +453,6 @@ def gradient_wrt_parameter(self, direction, *args, **kwargs):
         self.forward_solution, _ = self.solve()
 
         # Compute adjoint solution
-        test_parameter = dl.TestFunction(self.parameter_function_space)
         test_solution = dl.TestFunction(self.solution_function_space)
 
         # note: temp_form is a weak form used for building the adjoint operator
@@ -478,10 +477,24 @@ def gradient_wrt_parameter(self, direction, *args, **kwargs):
         # Compute gradient
         # note: temp_form is a weak form used for building the gradient
         temp_form = self.PDE_form(*self.parameter_args, self.forward_solution, adjoint)
-        gradient_form = dl.derivative(temp_form, *self.parameter_args, test_parameter)
-        gradient = dl.Function(self.parameter_function_space)
-        dl.assemble(gradient_form, tensor=gradient.vector())
-        return gradient
+        parameter_function_space = self.parameter_function_space
+        gradient_list = []
+        if not isinstance(self.parameter_function_space, (list, tuple)):
+            parameter_function_space = [self.parameter_function_space]
+
+        for i, k in enumerate(self._non_default_args):
+            test_parameter = dl.TestFunction(parameter_function_space[i])
+            gradient_form = dl.derivative(temp_form, self.parameter_args[i], test_parameter)
+            gradient = dl.Function(parameter_function_space[i])
+            dl.assemble(gradient_form, tensor=gradient.vector())
+            gradient_list.append(gradient)
+
+        # If only one parameter is passed, return a single gradient
+        if len(gradient_list) == 1:
+            gradient_list = gradient_list[0]
+        else:
+            gradient_list = tuple(gradient_list)
+        return gradient_list
 
 
     def _apply_obs_op(self, PDE_parameter_fun, PDE_solution_fun,):

From f6dc6296a95187606dc238fcfc4ba1b3c61f8681 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Thu, 5 Jun 2025 10:54:17 +0300
Subject: [PATCH 05/19] remove if main == name

---
 tests/test_PDEModel.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/test_PDEModel.py b/tests/test_PDEModel.py
index 50362e4..8105e99 100644
--- a/tests/test_PDEModel.py
+++ b/tests/test_PDEModel.py
@@ -1,3 +1,4 @@
+#%%
 import dolfin as dl
 import cuqi
 import cuqipy_fenics
@@ -492,4 +493,4 @@ def test_form_multiple_inputs():
 
     # Check passing wrong parameter_function_space raises an error
     with pytest.raises(ValueError, match= r"assemble input is specified by keywords arguments \['m2', 'source_term'\] that does not match the non_default_args of assemble \['m', 'source_term'\]"):
-        PDE_with_full_form.assemble(m2=m, source_term=source_term)
\ No newline at end of file
+        PDE_with_full_form.assemble(m2=m, source_term=source_term)

From 0ba732f2096a1f3fafcaa46c93464fee4832a0f8 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Thu, 5 Jun 2025 11:42:18 +0300
Subject: [PATCH 06/19] fix poisson class: only one instance of function space
 should be created

---
 tests/test_PDEModel.py | 39 +++++++++++++++++++++------------------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/tests/test_PDEModel.py b/tests/test_PDEModel.py
index 8105e99..0c033b1 100644
--- a/tests/test_PDEModel.py
+++ b/tests/test_PDEModel.py
@@ -1,4 +1,3 @@
-#%%
 import dolfin as dl
 import cuqi
 import cuqipy_fenics
@@ -310,23 +309,32 @@ class Poisson:
     """Define the variational PDE problem for the Poisson equation in two
     ways: as a full form, and as lhs and rhs forms"""
 
-    def __init__(self):
+    def __init__(self, mesh):
 
-        # Create the mesh and define function spaces for the solution and the
-        # parameter
-        self.mesh = dl.UnitIntervalMesh(10)
+        # Set the mesh
+        self._mesh =mesh
 
         # Define the boundary condition
         self.bc_value = dl.Constant(0.0)
 
-        # the source term
-        self.source_term = dl.Expression('x[0]', degree=1)
+        # The source term
+        self._source_term = dl.Expression('x[0]', degree=1)
+
+        # Solution function space
+        self._solution_function_space = dl.FunctionSpace(self.mesh, "Lagrange", 2)
+
+        # Parameter function space
+        self._parameter_function_space = dl.FunctionSpace(self.mesh, "Lagrange", 1)
+
+    @property
+    def mesh(self):
+        return self._mesh
 
     @property
     def form(self):
         return lambda m, u, v:\
             ufl.exp(m)*ufl.inner(ufl.grad(u), ufl.grad(v))*ufl.dx\
-            + self.source_term*v*ufl.dx
+            + self._source_term*v*ufl.dx
 
     @property
     def lhs_form(self):
@@ -339,22 +347,17 @@ def rhs_form(self):
 
     @property
     def solution_function_space(self):
-        return dl.FunctionSpace(self.mesh, "Lagrange", 2)
+        return self._solution_function_space
 
     @property
     def parameter_function_space(self):
-        return dl.FunctionSpace(self.mesh, "Lagrange", 1)
+        return self._parameter_function_space
 
     @property
     def bcs(self):
-        if not hasattr(self, "_bcs") or self._bcs is None:
-            self._bcs = dl.DirichletBC(self.solution_function_space,
-                                      self.bc_value, "on_boundary")
-        return self._bcs
-
-    @bcs.setter
-    def bcs(self, bcs):
-        self._bcs = bcs
+        return dl.DirichletBC(
+            self.solution_function_space, self.bc_value, "on_boundary"
+        )
 
 
 def test_observation_operator_setter():

From b36a0414c3b58faab5a1e5cfd83dd50a5dfaedff Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Sat, 7 Jun 2025 12:30:13 +0300
Subject: [PATCH 07/19] add gradient test for poisson problem

---
 tests/test_PDEModel.py | 57 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/tests/test_PDEModel.py b/tests/test_PDEModel.py
index 0c033b1..e5ec739 100644
--- a/tests/test_PDEModel.py
+++ b/tests/test_PDEModel.py
@@ -5,6 +5,7 @@
 import pytest
 import time
 import sys
+from scipy import optimize
 ufl = cuqipy_fenics.utilities._import_ufl()
 
 def test_model_input():
@@ -409,6 +410,62 @@ def test_observation_operator_setter():
     assert np.allclose(u1_obs, u2_obs) and np.allclose(u2_obs, u3_obs)
     assert len(u1_obs) == 5
 
+def test_gradient_poisson():
+    """Test the gradient of the Poisson PDEModel is computed correctly"""
+    # Set random seed
+    np.random.seed(0)
+
+    # Create the variational problem
+    poisson = Poisson(dl.UnitIntervalMesh(40))
+
+    # Set up model parameters
+    m = dl.Function(poisson.parameter_function_space)
+    m.vector()[:] = 1.0
+
+    # Create a PDE object
+    PDE = cuqipy_fenics.pde.SteadyStateLinearFEniCSPDE(
+        poisson.form,
+        poisson.mesh,
+        poisson.solution_function_space,
+        poisson.parameter_function_space,
+        poisson.bcs,
+        poisson.bcs)
+
+    # Create a PDE model
+    PDE_model = cuqi.model.PDEModel(
+        PDE,
+        domain_geometry=cuqipy_fenics.geometry.FEniCSContinuous(
+            poisson.parameter_function_space
+        ),
+        range_geometry=cuqipy_fenics.geometry.FEniCSContinuous(
+            poisson.solution_function_space
+        ),
+    )
+
+    # Create a prior distribution
+    m_prior = cuqi.distribution.Gaussian(
+        mean=np.zeros(PDE_model.domain_dim), cov=1, geometry=PDE_model.domain_geometry
+    )
+
+    # Create a likelihood
+    y = cuqi.distribution.Gaussian(
+        mean=PDE_model(m_prior), cov=np.ones(PDE_model.range_dim)*.1**2, geometry=PDE_model.range_geometry)
+    y = y(y=PDE_model(m.vector().get_local()))
+
+    # Evaluate the adjoint based gradient at value m2
+    m2 = dl.Function(poisson.parameter_function_space)
+    m2.vector()[:] = np.random.randn(PDE_model.domain_dim)
+    adjoint_grad = y.gradient(m_prior=m2.vector().get_local())
+
+    # Compute the FD gradient
+    step = 1e-7   # finite diff step
+    FD_grad = optimize.approx_fprime(m2.vector().get_local(), y.logd, step)
+
+    # Check that the adjoint gradient and FD gradient are close
+    assert np.allclose(adjoint_grad, FD_grad, rtol=1e-1),\
+        f"Adjoint gradient: {adjoint_grad}, FD gradient: {FD_grad}"
+    assert np.linalg.norm(adjoint_grad-FD_grad)/ np.linalg.norm(FD_grad) < 1e-2
+
 class PoissonMultipleInputs:
     """Define the variational PDE problem for the Poisson equation with multiple
      unknowns in two ways: as a full form, and as lhs and rhs forms"""

From 529fad047194e724ad66835722d7351e96e52b99 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Sat, 7 Jun 2025 12:39:32 +0300
Subject: [PATCH 08/19] Add gradient test for model with multiple inputs

---
 tests/test_PDEModel.py | 69 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 67 insertions(+), 2 deletions(-)

diff --git a/tests/test_PDEModel.py b/tests/test_PDEModel.py
index e5ec739..6fcf8ea 100644
--- a/tests/test_PDEModel.py
+++ b/tests/test_PDEModel.py
@@ -479,6 +479,14 @@ def __init__(self):
         # Define the boundary condition
         self.bc_value = dl.Constant(0.0)
 
+        # Set the solution function space
+        self._solution_function_space =  dl.FunctionSpace(self.mesh, "Lagrange", 2)
+
+        # Set the parameter function space
+        self._parameter_function_space = [
+            dl.FunctionSpace(self.mesh, "Lagrange", 1),
+            dl.FunctionSpace(self.mesh, "Lagrange", 1)]
+
     @property
     def form(self):
         return lambda m, source_term, u, v:\
@@ -496,11 +504,11 @@ def rhs_form(self):
 
     @property
     def solution_function_space(self):
-        return dl.FunctionSpace(self.mesh, "Lagrange", 2)
+        return self._solution_function_space
 
     @property
     def parameter_function_space(self):
-        return [dl.FunctionSpace(self.mesh, "Lagrange", 1), dl.FunctionSpace(self.mesh, "Lagrange", 1)]
+        return self._parameter_function_space
 
     @property
     def bcs(self):
@@ -554,3 +562,60 @@ def test_form_multiple_inputs():
     # Check passing wrong parameter_function_space raises an error
     with pytest.raises(ValueError, match= r"assemble input is specified by keywords arguments \['m2', 'source_term'\] that does not match the non_default_args of assemble \['m', 'source_term'\]"):
         PDE_with_full_form.assemble(m2=m, source_term=source_term)
+
+def test_gradient_poisson_multiple_inputs():
+    """Test the gradient of the Poisson PDEModel with multiple inputs is computed correctly"""
+    # Set random seed
+    np.random.seed(0)
+
+    # Create the variational problem
+    poisson = PoissonMultipleInputs()
+
+    # Set up model parameters
+    m = dl.Function(poisson.parameter_function_space[0])
+    m.vector()[:] = 1.0
+    source_term = dl.Function(poisson.parameter_function_space[1])
+    source_term.interpolate(dl.Expression('x[0]', degree=1))
+
+    # Create a PDE object
+    PDE = cuqipy_fenics.pde.SteadyStateLinearFEniCSPDE(
+        poisson.form,
+        poisson.mesh,
+        poisson.solution_function_space,
+        poisson.parameter_function_space,
+        poisson.bcs)
+
+    # Create a PDE model
+    PDE_model = cuqi.model.PDEModel(
+        PDE,
+        domain_geometry=cuqipy_fenics.geometry.FEniCSContinuous(
+            poisson.parameter_function_space[0]
+        ),
+        range_geometry=cuqipy_fenics.geometry.FEniCSContinuous(
+            poisson.solution_function_space
+        ),
+    )
+
+    # Create a prior distribution
+    m_prior = cuqi.distribution.Gaussian(
+        mean=np.zeros(PDE_model.domain_dim), cov=1, geometry=PDE_model.domain_geometry
+    )
+
+    # Create a likelihood
+    y = cuqi.distribution.Gaussian(
+        mean=PDE_model(m_prior, source_term), cov=np.ones(PDE_model.range_dim)*.1**2, geometry=PDE_model.range_geometry)
+    y = y(y=PDE_model(m.vector().get_local(), source_term))
+
+    # Evaluate the adjoint based gradient at value m2
+    m2 = dl.Function(poisson.parameter_function_space[0])
+    m2.vector()[:] = np.random.randn(PDE_model.domain_dim)
+    adjoint_grad = y.gradient(m_prior=m2.vector().get_local())
+
+    # Compute the FD gradient
+    step = 1e-7   # finite diff step
+    FD_grad = optimize.approx_fprime(m2.vector().get_local(), y.logd, step)
+
+    # Check that the adjoint gradient and FD gradient are close
+    assert np.allclose(adjoint_grad, FD_grad, rtol=1e-1),\
+        f"Adjoint gradient: {adjoint_grad}, FD gradient: {FD_grad}"
+    assert np.linalg.norm(adjoint_grad-FD_grad)/ np.linalg.norm(FD_grad) < 1e-2

From 90e4568b032ccdf3fa8c8c464d8b271c044f8b06 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Sat, 7 Jun 2025 18:39:52 +0300
Subject: [PATCH 09/19] update gradient test

---
 tests/test_PDEModel.py | 49 ++++++++++++++++++++++++++++--------------
 1 file changed, 33 insertions(+), 16 deletions(-)

diff --git a/tests/test_PDEModel.py b/tests/test_PDEModel.py
index 6fcf8ea..8e83fb6 100644
--- a/tests/test_PDEModel.py
+++ b/tests/test_PDEModel.py
@@ -583,39 +583,56 @@ def test_gradient_poisson_multiple_inputs():
         poisson.mesh,
         poisson.solution_function_space,
         poisson.parameter_function_space,
+        poisson.bcs,
         poisson.bcs)
 
+    domain_geom1 = cuqipy_fenics.geometry.FEniCSContinuous(
+                poisson.parameter_function_space[0]
+            )
+    domain_geom2 = cuqipy_fenics.geometry.FEniCSContinuous(
+                poisson.parameter_function_space[1]
+            )
+
     # Create a PDE model
     PDE_model = cuqi.model.PDEModel(
         PDE,
-        domain_geometry=cuqipy_fenics.geometry.FEniCSContinuous(
-            poisson.parameter_function_space[0]
-        ),
+        domain_geometry=(domain_geom1, domain_geom2),
         range_geometry=cuqipy_fenics.geometry.FEniCSContinuous(
-            poisson.solution_function_space
-        ),
+            poisson.solution_function_space)
     )
 
     # Create a prior distribution
     m_prior = cuqi.distribution.Gaussian(
-        mean=np.zeros(PDE_model.domain_dim), cov=1, geometry=PDE_model.domain_geometry
+        mean=np.zeros(domain_geom1.par_dim), cov=1, geometry=domain_geom1
+    )
+    source_term_prior = cuqi.distribution.Gaussian(
+        mean=np.zeros(domain_geom2.par_dim), cov=1, geometry=domain_geom2
     )
 
     # Create a likelihood
     y = cuqi.distribution.Gaussian(
-        mean=PDE_model(m_prior, source_term), cov=np.ones(PDE_model.range_dim)*.1**2, geometry=PDE_model.range_geometry)
-    y = y(y=PDE_model(m.vector().get_local(), source_term))
+        mean=PDE_model(m_prior, source_term_prior), cov=np.ones(PDE_model.range_dim)*.1**2, geometry=PDE_model.range_geometry)
+    y = y(y=PDE_model(m.vector().get_local(), source_term.vector().get_local()))
 
     # Evaluate the adjoint based gradient at value m2
     m2 = dl.Function(poisson.parameter_function_space[0])
-    m2.vector()[:] = np.random.randn(PDE_model.domain_dim)
-    adjoint_grad = y.gradient(m_prior=m2.vector().get_local())
+    m2.vector()[:] = np.random.randn(domain_geom1.par_dim)
+    source_term2 = dl.Function(poisson.parameter_function_space[1])
+    source_term2.interpolate(dl.Expression('sin(2*x[0]*pi)', degree=1))
+
+    # Compute the adjoint gradient
+    adjoint_grad = y.gradient(m_prior=m2.vector().get_local(), source_term_prior=source_term2.vector().get_local())
 
     # Compute the FD gradient
-    step = 1e-7   # finite diff step
-    FD_grad = optimize.approx_fprime(m2.vector().get_local(), y.logd, step)
+    step = 1e-9   # finite diff step
+    FD_grad_m = optimize.approx_fprime(m2.vector().get_local(), lambda m2: y.logd(m2, source_term2.vector().get_local()), step)
+    FD_grad_source_term = optimize.approx_fprime(source_term2.vector().get_local(), lambda source_term2: y.logd(m2.vector().get_local(), source_term2), step)
+
+    # Assert gradients are close
+    assert np.allclose(adjoint_grad['m_prior'], FD_grad_m, rtol=2e-1),\
+        f"Adjoint gradient w.r.t. m: {adjoint_grad['m_prior']}, FD gradient w.r.t. m: {FD_grad_m}"
+    assert np.allclose(adjoint_grad['source_term_prior'], FD_grad_source_term, rtol=1e-1),\
+        f"Adjoint gradient w.r.t. source term: {adjoint_grad['source_term_prior']}, FD gradient w.r.t. source term: {FD_grad_source_term}"
+    assert np.linalg.norm(adjoint_grad['m_prior']-FD_grad_m)/ np.linalg.norm(FD_grad_m) < 1e-2
+    assert np.linalg.norm(adjoint_grad['source_term_prior']-FD_grad_source_term)/ np.linalg.norm(FD_grad_source_term) < 1e-2
 
-    # Check that the adjoint gradient and FD gradient are close
-    assert np.allclose(adjoint_grad, FD_grad, rtol=1e-1),\
-        f"Adjoint gradient: {adjoint_grad}, FD gradient: {FD_grad}"
-    assert np.linalg.norm(adjoint_grad-FD_grad)/ np.linalg.norm(FD_grad) < 1e-2

From 3d80e3d26dd15a9e58b2ff6a791145894285c507 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Sat, 7 Jun 2025 19:02:27 +0300
Subject: [PATCH 10/19] Uodate tests

---
 tests/test_PDEModel.py | 43 +++++++++++++++++++++---------------------
 1 file changed, 22 insertions(+), 21 deletions(-)

diff --git a/tests/test_PDEModel.py b/tests/test_PDEModel.py
index 8e83fb6..107ce97 100644
--- a/tests/test_PDEModel.py
+++ b/tests/test_PDEModel.py
@@ -43,7 +43,7 @@ def test_model_input():
 def test_solver_choice():
     """Test passing different solvers to PDEModel"""
     # Create the variational problem
-    poisson = Poisson()
+    poisson = Poisson(dl.UnitIntervalMesh(10))
 
     # Set up model parameters
     m = dl.Function(poisson.parameter_function_space)
@@ -83,8 +83,7 @@ def test_reuse_assembled():
     """Test that reusing the assembled and factored lhs gives the same solution
      and a better performance"""
     # Create the variational problem
-    poisson = Poisson()
-    poisson.mesh = dl.UnitSquareMesh(50, 50)
+    poisson = Poisson(dl.UnitSquareMesh(50, 50))
 
     # Set up model parameters
     m = dl.Function(poisson.parameter_function_space)
@@ -140,7 +139,7 @@ def test_reuse_assembled():
 def test_form():
     """Test creating PDEModel with full form, and with lhs and rhs forms"""
     # Create the variational problem
-    poisson = Poisson()
+    poisson = Poisson(dl.UnitIntervalMesh(10))
 
     # Set up model parameters
     m = dl.Function(poisson.parameter_function_space)
@@ -185,24 +184,13 @@ def test_with_updated_rhs(copy_reference, case):
         pytest.skip("Test on MAC OS only")
 
     # Set up first poisson problem
-    poisson1 = Poisson()
-    poisson1.mesh = dl.UnitSquareMesh(40, 40)
+    poisson1 = Poisson(dl.UnitSquareMesh(40, 40))
     poisson1.source_term = dl.Constant(1.0)
 
     # Set up second poisson problem
-    poisson2 = Poisson()
-    poisson2.mesh = poisson1.mesh
+    poisson2 = Poisson(poisson1.mesh)
     poisson2.source_term = dl.Expression("sin(2*x[0]*pi)*sin(2*x[1]*pi)", degree=1)
 
-    # Set up boundary function (where the Dirichlet boundary conditions are applied)
-    def u_boundary(x, on_boundary):
-        return on_boundary and\
-            (x[0] < dl.DOLFIN_EPS or x[0] > 1.0 - dl.DOLFIN_EPS)
-
-    poisson1.bcs = dl.DirichletBC(poisson1.solution_function_space,
-                                 dl.Constant(0.0), u_boundary)
-    poisson2.bcs = poisson1.bcs
-
     # Create two PDE objects with different rhs terms
     PDE1 = cuqipy_fenics.pde.SteadyStateLinearFEniCSPDE(
         (poisson1.lhs_form, poisson1.rhs_form),
@@ -335,7 +323,7 @@ def mesh(self):
     def form(self):
         return lambda m, u, v:\
             ufl.exp(m)*ufl.inner(ufl.grad(u), ufl.grad(v))*ufl.dx\
-            + self._source_term*v*ufl.dx
+            + self.source_term*v*ufl.dx
 
     @property
     def lhs_form(self):
@@ -359,13 +347,22 @@ def bcs(self):
         return dl.DirichletBC(
             self.solution_function_space, self.bc_value, "on_boundary"
         )
+    
+    @property
+    def source_term(self):
+        """Return the source term"""
+        return self._source_term
+    
+    @source_term.setter
+    def source_term(self, value):
+        """Set the source term"""
+        self._source_term = value
 
 
 def test_observation_operator_setter():
     """Test that the observation setter works as expected"""
     # Create the variational problem
-    poisson = Poisson()
-    poisson.mesh = dl.UnitSquareMesh(50, 50)
+    poisson = Poisson(dl.UnitSquareMesh(50, 50))
 
     # Set up model parameters
     m = dl.Function(poisson.parameter_function_space)
@@ -474,7 +471,7 @@ def __init__(self):
 
         # Create the mesh and define function spaces for the solution and the
         # parameter
-        self.mesh = dl.UnitIntervalMesh(10)
+        self._mesh = dl.UnitIntervalMesh(10)
 
         # Define the boundary condition
         self.bc_value = dl.Constant(0.0)
@@ -487,6 +484,10 @@ def __init__(self):
             dl.FunctionSpace(self.mesh, "Lagrange", 1),
             dl.FunctionSpace(self.mesh, "Lagrange", 1)]
 
+    @property
+    def mesh(self):
+        return self._mesh
+
     @property
     def form(self):
         return lambda m, source_term, u, v:\

From ae863794543e8322387a236b3cb15c678f0cdf83 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Sat, 7 Jun 2025 20:26:51 +0300
Subject: [PATCH 11/19] update miniconda version

---
 .github/workflows/docs.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 556bf46..f1cf0d3 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -17,7 +17,7 @@ jobs:
       - uses: actions/setup-python@v4
 
       - name: Setup Miniconda
-        uses: conda-incubator/setup-miniconda@v2
+        uses: conda-incubator/setup-miniconda@v3
         with:
           auto-update-conda: true
           python-version: 3.8

From 327de620bce7985ace3a234a19ad604eacc9c60f Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Sat, 7 Jun 2025 20:29:15 +0300
Subject: [PATCH 12/19] update miniconda action version

---
 .github/workflows/tests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index da12601..0dfa6e3 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -19,7 +19,7 @@ jobs:
           fetch-depth: 1
 
       - name: Setup Miniconda
-        uses: conda-incubator/setup-miniconda@v2
+        uses: conda-incubator/setup-miniconda@v3
         with:
           auto-update-conda: true
           python-version: 3.8

From b84d7ed903dd819dc4ac85b5b8c2ce17a7ca4759 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Sat, 7 Jun 2025 22:00:34 +0300
Subject: [PATCH 13/19] miniforge instead of mambaforge

---
 .github/workflows/docs.yml  | 2 +-
 .github/workflows/tests.yml | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index f1cf0d3..4601221 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -21,7 +21,7 @@ jobs:
         with:
           auto-update-conda: true
           python-version: 3.8
-          miniforge-variant: Mambaforge
+          miniforge-variant: Miniforge
           miniforge-version: latest
 
       - name: Conda info
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 0dfa6e3..a8f8078 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -23,7 +23,7 @@ jobs:
         with:
           auto-update-conda: true
           python-version: 3.8
-          miniforge-variant: Mambaforge
+          miniforge-variant: Miniforge
           miniforge-version: latest
 
       - name: Conda info

From acbcc28a4b1427014ad172664a3ea441348e05e6 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Sat, 7 Jun 2025 22:10:37 +0300
Subject: [PATCH 14/19] update env

---
 .github/workflows/tests.yml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index a8f8078..a170e33 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -23,7 +23,7 @@ jobs:
         with:
           auto-update-conda: true
           python-version: 3.8
-          miniforge-variant: Miniforge
+          miniforge-variant: Miniforge3
           miniforge-version: latest
 
       - name: Conda info
@@ -32,12 +32,12 @@ jobs:
           conda info
           which python
 
-      - name: Mamba install FEniCS
+      - name: Conda install FEniCS
         shell: bash -l {0}
         run: |
           conda config --set always_yes yes
           conda config --add channels conda-forge
-          mamba create -n fenicsproject -c conda-forge fenics 
+          conda create -n fenicsproject -c conda-forge fenics 
           conda activate fenicsproject
           which python
           python -c "from dolfin import *"

From 02366a36f23dc1ba043ccb9234f3f979d0eb9dcc Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Sat, 7 Jun 2025 22:15:16 +0300
Subject: [PATCH 15/19] update miniforge

---
 .github/workflows/docs.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 4601221..b090dc7 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -21,7 +21,7 @@ jobs:
         with:
           auto-update-conda: true
           python-version: 3.8
-          miniforge-variant: Miniforge
+          miniforge-variant: Miniforge3
           miniforge-version: latest
 
       - name: Conda info
@@ -30,7 +30,7 @@ jobs:
           conda info
           which python
 
-      - name: Mamba install FEniCS
+      - name: conda install FEniCS
         shell: bash -l {0}
         run: |
           conda config --set always_yes yes

From 327083f1042b6b427260d1aae88a4814dfacbe90 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Sat, 7 Jun 2025 22:43:21 +0300
Subject: [PATCH 16/19] few updates

---
 tests/test_PDEModel.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/test_PDEModel.py b/tests/test_PDEModel.py
index 107ce97..3a8a35e 100644
--- a/tests/test_PDEModel.py
+++ b/tests/test_PDEModel.py
@@ -301,7 +301,7 @@ class Poisson:
     def __init__(self, mesh):
 
         # Set the mesh
-        self._mesh =mesh
+        self._mesh = mesh
 
         # Define the boundary condition
         self.bc_value = dl.Constant(0.0)
@@ -543,6 +543,7 @@ def test_form_multiple_inputs():
 
     # Solve the PDE
     PDE_with_full_form.assemble(m, source_term)
+    # test also assembling with keyword arguments
     PDE_with_full_form.assemble(m=m, source_term=source_term)    
     u1, info = PDE_with_full_form.solve()
 
@@ -553,6 +554,7 @@ def test_form_multiple_inputs():
         poisson.solution_function_space,
         poisson.parameter_function_space,
         poisson.bcs)
+
     # Solve the PDE
     PDE_with_lhs_rhs_forms.assemble(m, source_term)
     u2, info = PDE_with_lhs_rhs_forms.solve()
@@ -602,7 +604,7 @@ def test_gradient_poisson_multiple_inputs():
             poisson.solution_function_space)
     )
 
-    # Create a prior distribution
+    # Create prior distributions
     m_prior = cuqi.distribution.Gaussian(
         mean=np.zeros(domain_geom1.par_dim), cov=1, geometry=domain_geom1
     )
@@ -615,7 +617,7 @@ def test_gradient_poisson_multiple_inputs():
         mean=PDE_model(m_prior, source_term_prior), cov=np.ones(PDE_model.range_dim)*.1**2, geometry=PDE_model.range_geometry)
     y = y(y=PDE_model(m.vector().get_local(), source_term.vector().get_local()))
 
-    # Evaluate the adjoint based gradient at value m2
+    # Evaluate the adjoint based gradient at value m2 and source_term2
     m2 = dl.Function(poisson.parameter_function_space[0])
     m2.vector()[:] = np.random.randn(domain_geom1.par_dim)
     source_term2 = dl.Function(poisson.parameter_function_space[1])

From 2576605f2ea9a77f4fb38916c706db0fb3a45a6b Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Sun, 8 Jun 2025 10:54:17 +0300
Subject: [PATCH 17/19] update documentation

---
 cuqipy_fenics/pde.py | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/cuqipy_fenics/pde.py b/cuqipy_fenics/pde.py
index a233342..ec74f29 100644
--- a/cuqipy_fenics/pde.py
+++ b/cuqipy_fenics/pde.py
@@ -25,16 +25,18 @@ class FEniCSPDE(PDE,ABC):
     ----------
     PDE_form : callable or tuple of two callables
         If passed as a callable: the callable returns the weak form of the PDE.
-        The callable should take three arguments, the first argument is the 
-        parameter (input of the forward model), the second argument is the state
-        variable (solution variable), and the third argument is the adjoint
+        The callable should take three or more arguments, the first arguments
+        are the unknown parameters (inputs of the forward model, e.g.
+        `parameter1`, `parameter2`), the second to last argument is the state
+        variable (solution variable), and the last argument is the adjoint
         variable (the test variable in the weak formulation).
 
         If passed as a tuple of two callables: the first callable returns the 
         weak form of the PDE left hand side, and the second callable returns the
         weak form of the PDE right hand side. The left hand side callable takes
-        the same three arguments as described above. The right hand side 
-        callable takes only the parameter and the adjoint variable as arguments.
+        the same arguments as described above. The right hand side callable
+        takes only the unknown parameters and the adjoint variable (the latter 
+        being the last argument) as arguments.
         See the example below.
 
     mesh : FEniCS mesh
@@ -43,8 +45,9 @@ class FEniCSPDE(PDE,ABC):
     solution_function_space : FEniCS function space
         FEniCS function space object that defines the function space of the state variable (solution variable).
 
-    parameter_function_space : FEniCS function space
-        FEniCS function space object that defines the function space of the Bayesian parameter (input of the forward model).
+    parameter_function_space : FEniCS function space or a list of them
+        FEniCS function space object or a list of them that defines the function space of the unknown parameters (inputs of the forward model).
+        If multiple parameters are passed, the function space should be a list of FEniCS function spaces, one for each parameter.
 
     dirichlet_bcs: FEniCS Dirichlet boundary condition object or a list of them
         FEniCS Dirichlet boundary condition object(s) that define the Dirichlet boundary conditions of the PDE.
@@ -54,7 +57,7 @@ class FEniCSPDE(PDE,ABC):
 
     observation_operator : python function handle, optional
         Function handle of a python function that returns the observed quantity from the PDE solution. If not provided, the identity operator is assumed (i.e. the entire solution is observed).
-        This python function takes as input the Bayesian parameter (input of the forward model) and the state variable (solution variable) as first and second inputs, respectively.
+        This python function takes as input the unknown parameters, e.g. `parameter1`, `parameter2`, and the state variable (solution variable) in that order.
 
         The returned observed quantity can be a ufl.algebra.Operator, FEniCS Function, np.ndarray, int, or float.
 
@@ -165,8 +168,8 @@ def parameter(self):
     
     @parameter.setter
     def parameter(self, value):
-        """ Set the parameter of the PDE. Since the PDE solution depends on the 
-        parameter, this will set the PDE solution to None. """
+        """ Set the parameters of the PDE. Since the PDE solution depends on the 
+        parameters, this will set the PDE solution to None. """
         if value is None:
             raise ValueError('Parameter cannot be None.')
 
@@ -426,7 +429,7 @@ def observe(self,PDE_solution_fun):
             return self._apply_obs_op(*self.parameter_args, PDE_solution_fun)
 
     def gradient_wrt_parameter(self, direction, *args, **kwargs):
-        """ Compute the gradient of the PDE with respect to the parameter
+        """ Compute the gradient of the PDE with respect to the parameters
 
         Note: This implementation is largely based on the code:
         https://github.com/hippylib/hippylib/blob/master/hippylib/modeling/PDEProblem.py

From 315711034aedac2d83435d700073fe0fdec97824 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Fri, 4 Jul 2025 09:29:01 +0300
Subject: [PATCH 18/19] update comments about extracting non_default_args

---
 cuqipy_fenics/pde.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/cuqipy_fenics/pde.py b/cuqipy_fenics/pde.py
index ec74f29..d38821e 100644
--- a/cuqipy_fenics/pde.py
+++ b/cuqipy_fenics/pde.py
@@ -201,9 +201,13 @@ def parameter_args(self):
     def _non_default_args(self):
         form = self._form
         if isinstance(self._form, tuple):
-            # extract non-default args from the lhs first form
+            # extract non-default args from the lhs form and exclude the last
+            # two arguments (u and p) from the list of non-default args since
+            # they are provided automatically within the PDE-type class and are
+            # not arguments to be inferred in Bayesian inference setting.
             form = self._form[0]
-        return get_non_default_args(form)[:-2] # Exclude the last two arguments (u and p) from the list of non-default args
+            non_default_args = get_non_default_args(form)[:-2]
+        return non_default_args
 
     @property
     def forward_solution(self):

From afe09a4c708624fe3cdbb7cec4d7eff2988fbde8 Mon Sep 17 00:00:00 2001
From: amal-ghamdi <amal.m.alghamdi@gmail.com>
Date: Fri, 4 Jul 2025 09:36:27 +0300
Subject: [PATCH 19/19] update comments

---
 cuqipy_fenics/pde.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/cuqipy_fenics/pde.py b/cuqipy_fenics/pde.py
index d38821e..359d09a 100644
--- a/cuqipy_fenics/pde.py
+++ b/cuqipy_fenics/pde.py
@@ -201,12 +201,14 @@ def parameter_args(self):
     def _non_default_args(self):
         form = self._form
         if isinstance(self._form, tuple):
-            # extract non-default args from the lhs form and exclude the last
-            # two arguments (u and p) from the list of non-default args since
-            # they are provided automatically within the PDE-type class and are
-            # not arguments to be inferred in Bayesian inference setting.
+            # Use the lhs form only to determine the non-default args
             form = self._form[0]
-            non_default_args = get_non_default_args(form)[:-2]
+
+        # Extract the non_default_args and exclude the last
+        # two arguments (u and p) from the list of non-default args since
+        # they are provided automatically within the PDE-type class and are
+        # not arguments to be inferred in Bayesian inference setting.
+        non_default_args = get_non_default_args(form)[:-2]
         return non_default_args
 
     @property