Repository: CUQI-DTU/CUQIpy-FEniCS (language: python)

Issue:
Update FEniCS plug in to use the new Gaussian distribution 



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 2aa6000..c8d28f7 100644
--- a/README.md
+++ b/README.md
@@ -30,7 +30,7 @@ model, y_data, info = cuqipy_fenics.testproblem.FEniCSDiffusion1D.get_components
 x = cuqi.distribution.GMRF(np.zeros(model.domain_dim),
                            25, 1, 'zero', geometry=model.domain_geometry)
 # y ~ N(model(x), 0.01^2)
-y = cuqi.distribution.GaussianCov(model(x), 0.05**2)
+y = cuqi.distribution.Gaussian(mean=model(x), cov=0.05**2)
 
 # Set up Bayesian Problem object
 BP = cuqi.problem.BayesianProblem(y, x).set_data(y=y_data)
diff --git a/cuqipy_fenics/geometry.py b/cuqipy_fenics/geometry.py
index 3ed6b2f..648dc72 100644
--- a/cuqipy_fenics/geometry.py
+++ b/cuqipy_fenics/geometry.py
@@ -158,7 +158,7 @@ class MaternExpansion(_WrappedGeometry):
         import numpy as np
         import matplotlib.pyplot as plt
         from cuqi.fenics.geometry import MaternExpansion, FEniCSContinuous
-        from cuqi.distribution import GaussianCov
+        from cuqi.distribution import Gaussian
         import dolfin as dl
         
         mesh = dl.UnitSquareMesh(20,20)
@@ -168,7 +168,7 @@ class MaternExpansion(_WrappedGeometry):
                                         length_scale = .2,
                                         num_terms=128)
         
-        MaternField = GaussianCov(np.zeros(MaternGeometry.dim),
+        MaternField = Gaussian(mean=np.zeros(MaternGeometry.dim),
                         cov=np.eye(MaternGeometry.dim),
                         geometry= MaternGeometry)
         
diff --git a/demos/demo01_1D_Diffusion.ipynb b/demos/demo01_1D_Diffusion.ipynb
index c95f8f9..b5d654d 100644
--- a/demos/demo01_1D_Diffusion.ipynb
+++ b/demos/demo01_1D_Diffusion.ipynb
@@ -261,7 +261,7 @@
    "outputs": [],
    "source": [
     "sigma = np.linalg.norm(problemInfo.exactData)/SNR \n",
-    "likelihood = cuqi.distribution.GaussianCov(model, sigma**2*np.eye(model.range_dim)).to_likelihood(data)"
+    "likelihood = cuqi.distribution.Gaussian(mean=model, cov=sigma**2*np.eye(model.range_dim)).to_likelihood(data)"
    ]
   },
   {
diff --git a/demos/demo01_1D_Diffusion.py b/demos/demo01_1D_Diffusion.py
index 3560a03..4ec4987 100644
--- a/demos/demo01_1D_Diffusion.py
+++ b/demos/demo01_1D_Diffusion.py
@@ -28,7 +28,7 @@
 
 
 sigma = np.linalg.norm(problemInfo.exactData)/SNR 
-likelihood = cuqi.distribution.GaussianCov(model, sigma**2*np.eye(model.range_dim)).to_likelihood(data)
+likelihood = cuqi.distribution.Gaussian(mean=model, cov=sigma**2*np.eye(model.range_dim)).to_likelihood(data)
 posterior = cuqi.distribution.Posterior(likelihood, prior)
 # %%
 sampler = cuqi.sampler.pCN(posterior)
diff --git a/demos/demo02_Matern.py b/demos/demo02_Matern.py
index 9245727..46fbdda 100644
--- a/demos/demo02_Matern.py
+++ b/demos/demo02_Matern.py
@@ -4,7 +4,7 @@
 import matplotlib.pyplot as plt
 sys.path.append("../")
 from cuqipy_fenics.geometry import MaternExpansion, FEniCSContinuous
-from cuqi.distribution import GaussianCov
+from cuqi.distribution import Gaussian
 import dolfin as dl
 
 mesh = dl.UnitSquareMesh(20,20)
@@ -14,9 +14,9 @@
                                 length_scale = .2,
                                 num_terms=128)
 
-MaternField = GaussianCov(np.zeros(MaternGeometry.par_dim),
-                cov=np.eye(MaternGeometry.par_dim),
-                geometry= MaternGeometry)
+MaternField = Gaussian(mean=np.zeros(MaternGeometry.par_dim),
+                       cov=np.eye(MaternGeometry.par_dim),
+                       geometry=MaternGeometry)
 
 samples = MaternField.sample()
 samples.plot()
diff --git a/demos/demo03_poisson_circular.py b/demos/demo03_poisson_circular.py
index c81f6a3..dc6cdc1 100644
--- a/demos/demo03_poisson_circular.py
+++ b/demos/demo03_poisson_circular.py
@@ -91,7 +91,9 @@ def heavy_map(func):
 
 #%% 2.5 Create a prior
 pr_mean = np.zeros(domain_geometry.par_dim)
-prior = cuqi.distribution.GaussianCov(pr_mean, cov=np.eye(domain_geometry.par_dim), geometry= domain_geometry)
+prior = cuqi.distribution.Gaussian(mean=pr_mean,
+                                   cov=np.eye(domain_geometry.par_dim),
+                                   geometry=domain_geometry)
 
 #%% 2.6 Define the exact solution
 exactSolution = prior.sample()
@@ -103,7 +105,9 @@ def heavy_map(func):
 SNR = 100
 sigma = np.linalg.norm(b_exact)/SNR
 sigma2 = sigma*sigma # variance of the observation Gaussian noise
-data_distribution = cuqi.distribution.GaussianCov(model, sigma2*np.ones(range_geometry.par_dim), geometry=range_geometry)
+data_distribution = cuqi.distribution.Gaussian(mean=model,
+                                               cov=sigma2*np.ones(range_geometry.par_dim),
+                                               geometry=range_geometry)
 
 #%% 2.9 Generate noisy data
 data = data_distribution(x=exactSolution).sample()


Test patch:
diff --git a/cuqipy_fenics/testproblem.py b/cuqipy_fenics/testproblem.py
index 66a91fc..a98c02d 100644
--- a/cuqipy_fenics/testproblem.py
+++ b/cuqipy_fenics/testproblem.py
@@ -136,8 +136,8 @@ def form(m,u,p):
             b_exact + alpha*noise, geometry=range_geometry)
 
         # Create likelihood
-        y = cuqi.distribution.GaussianCov(
-            model(x), alpha*np.eye(range_geometry.par_dim))
+        y = cuqi.distribution.Gaussian(
+            mean=model(x), cov=alpha*np.eye(range_geometry.par_dim))
 
         # Initialize FEniCSDiffusion1D as BayesianProblem problem
         super().__init__(y, x, y=data)
diff --git a/tests/test_PDEModel.py b/tests/test_PDEModel.py
new file mode 100644
index 0000000..fc210ff
--- /dev/null
+++ b/tests/test_PDEModel.py
@@ -0,0 +1,37 @@
+import dolfin as dl
+import cuqi
+import cuqipy_fenics
+import numpy as np
+import pytest
+
+
+def test_model_input():
+    """Test passing different data structures for PDEModel input"""
+    model = cuqipy_fenics.testproblem.FEniCSDiffusion1D().model
+    V = model.pde.parameter_function_space
+
+    # Test passing a CUQIarray containing parameters
+    u = dl.Function(V)
+    u_CUQIarray = cuqi.samples.CUQIarray(u.vector().get_local(), is_par=True, geometry=model.domain_geometry)
+    y = model(u_CUQIarray)
+
+    # Test passing parameters as a numpy array
+    u = dl.Function(V)
+    u_numpy = u.vector().get_local()
+    y = model(u_numpy, is_par=True)
+
+    # Test passing a CUQIarray containing dolfin function
+    u = dl.Function(V)
+    u_CUQIarray = cuqi.samples.CUQIarray(u, is_par=False, geometry=model.domain_geometry)
+    y = model(u_CUQIarray)
+
+    # Test passing a CUQIarray containing dolfin function wrapped in np.array
+    # This is not recommended, but should work
+    u = dl.Function(V)
+    u_CUQIarray = cuqi.samples.CUQIarray(np.array(u, dtype='O'), is_par=False, geometry=model.domain_geometry)
+    y = model(u_CUQIarray)  
+
+    # Test passing a dolfin function (should fail)
+    u = dl.Function(V)
+    with pytest.raises(AttributeError):
+        y = model(u)