Fix bug in advection diffusion benchmark when computing adjoint solution with Robin boundary conditions.

Author: jdjakem

Diff for: pyapprox/benchmarks/pde_benchmarks.py

@@ -89,6 +89,8 @@ def raw_advection_diffusion_reaction_kle_dRdp(kle, residual, sol, param_vals):
         kle_vals = kle(param_vals[:, None])
         assert kle_vals.ndim == 2
         dkdp = kle_vals*kle.eig_vecs
+    else:
+        dkdp = kle.eig_vecs
     Du = [torch.linalg.multi_dot((dmats[dd], sol))
           for dd in range(mesh.nphys_vars)]
     kDu = [Du[dd][:, None]*dkdp for dd in range(mesh.nphys_vars)]
@@ -98,10 +100,32 @@ def raw_advection_diffusion_reaction_kle_dRdp(kle, residual, sol, param_vals):
 
 
 def advection_diffusion_reaction_kle_dRdp(
-        bndry_indices, kle, residual, sol, param_vals):
+        mesh, kle, bndry_conds, residual, sol, param_vals):
     dRdp = raw_advection_diffusion_reaction_kle_dRdp(
         kle, residual, sol, param_vals)
-    dRdp[np.hstack(bndry_indices)] = 0.0
+    for ii, bndry_cond in enumerate(bndry_conds):
+        idx = mesh._bndry_indices[ii]
+        if bndry_cond[1] == "D":
+            dRdp[idx] = 0.0
+        elif bndry_cond[1] == "R":
+            mesh_pts_idx = mesh._bndry_slice(mesh.mesh_pts, idx, 1)
+            normal_vals = mesh._bndrys[ii].normals(mesh_pts_idx)
+            if kle.use_log:
+                kle_vals = kle(param_vals[:, None])
+                dkdp = kle_vals*kle.eig_vecs
+            else:
+                dkdp = torch.as_tensor(kle.eig_vecs)
+            flux_vals = [
+                (torch.linalg.multi_dot(
+                    (mesh._bndry_slice(mesh._dmat(dd), idx, 0), sol))[:, None]
+                 * mesh._bndry_slice(dkdp, idx, 0))
+                for dd in range(mesh.nphys_vars)]
+            flux_normal_vals = [
+                normal_vals[:, dd:dd+1]*flux_vals[dd]
+                for dd in range(mesh.nphys_vars)]
+            dRdp[idx] = sum(flux_normal_vals)
+        else:
+            raise NotImplementedError()
     return dRdp
 
 
@@ -121,8 +145,7 @@ def __init__(self, mesh, bndry_conds, kle, vel_fun, react_funs, forc_fun,
         import inspect
         if "mesh" == inspect.getfullargspec(functional).args[0]:
             if "physics" == inspect.getfullargspec(functional).args[1]:
-                functional = partial(
-                    functional, mesh, self._fwd_solver.physics)
+                functional = partial(functional, mesh, self._fwd_solver.physics)
             else:
                 functional = partial(functional, mesh)
         for ii in range(len(functional_deriv_funs)):
@@ -141,7 +164,7 @@ def __init__(self, mesh, bndry_conds, kle, vel_fun, react_funs, forc_fun,
         if issubclass(type(self._fwd_solver), SteadyStatePDE):
             dqdu, dqdp = functional_deriv_funs
             dRdp = partial(advection_diffusion_reaction_kle_dRdp,
-                           mesh._bndry_indices, self._kle)
+                           mesh, self._kle, self._fwd_solver.physics._bndry_conds)
             # dRdp must be after boundary conditions are applied.
             # For now assume that parameters do not effect boundary conditions
             # so dRdp at boundary indices is zero

Diff for: pyapprox/benchmarks/tests/test_pde_benchmarks.py

@@ -68,21 +68,26 @@ def set_random_sample(physics, sample):
             physics._diff_fun = partial(
                 inv_model.base_model._fast_interpolate,
                 inv_model.base_model._kle(sample[:, None]))
-        sample.requires_grad_= True
+        sample.requires_grad_ = True
         dRdp_ad = torch.autograd.functional.jacobian(
             partial(inv_model.base_model._adj_solver._parameterized_residual,
                     fwd_sol, set_random_sample),
             sample, strict=True)
         dRdp = inv_model.base_model._adj_solver._dRdp(
             inv_model.base_model._fwd_solver.physics, fwd_sol.clone(),
             sample)
+        print(dRdp)
+        print(dRdp_ad)
         assert np.allclose(dRdp, dRdp_ad)
 
         # TODO add std to params list
         init_guess = true_params + np.random.normal(0, 0.01, true_params.shape)
         # init_guess = sample.numpy()[:, None]
-
-        # from pyapprox.util.utilities import approx_fprime
+        # from pyapprox.util.utilities import approx_fprime, approx_jacobian
+        # fd_jac = approx_jacobian(
+        #     lambda p: inv_model.base_model._adj_solver._parameterized_residual(
+        #         fwd_sol, set_random_sample, torch.as_tensor(p)), sample)
+        # print(fd_jac)
         # print(approx_fprime(init_guess, inv_model), 'fd')
         # print(inv_model(init_guess, return_grad=True), 'g')
         # assert False

Diff for: pyapprox/pde/autopde/mesh.py

@@ -400,10 +400,10 @@ def __init__(self, orders, basis_types=None):
         self._flux_islinear = False
         self._flux_normal_vals = [None for dd in range(2*self.nphys_vars)]
         self._normal_vals = [None for dd in range(2*self.nphys_vars)]
-        
+
     def _clear_flux_normal_vals(self):
         self._flux_normal_vals = [None for dd in range(2*self.nphys_vars)]
-        
+
     @staticmethod
     def _get_basis_types(nphys_vars, basis_types):
         if basis_types is None:
@@ -693,7 +693,7 @@ def _dmat(self, dd):
     def _apply_dirichlet_boundary_conditions_special_indexing(
             self, bndry_conds, residual, jac, sol):
         # special indexing copies data which slows down function
-        
+
         # needs to have indices as argument so this fucntion can be used
         # when setting boundary conditions for forward and adjoint solves
 
@@ -716,12 +716,11 @@ def _apply_dirichlet_boundary_conditions_special_indexing(
             residual[idx] = sol[idx]-bndry_vals
         return residual, jac
 
-    
     @staticmethod
     def _bndry_slice(vec, idx, axis):
         # avoid copying data
         if len(idx) == 1:
-            if axis == 0 :
+            if axis == 0:
                 return vec[idx]
             return vec[:, idx]
         stride = idx[1]-idx[0]
@@ -755,7 +754,7 @@ def _apply_neumann_and_robin_boundary_conditions(
             if bndry_cond[1] != "N" and bndry_cond[1] != "R":
                 continue
             idx = self._bndry_indices[ii]
-            mesh_pts_idx =  self._bndry_slice(self.mesh_pts, idx, 1)
+            mesh_pts_idx = self._bndry_slice(self.mesh_pts, idx, 1)
             if self._normal_vals[ii] is None:
                 self._normal_vals[ii] = self._bndrys[ii].normals(mesh_pts_idx)
             if not self._flux_islinear or self._flux_normal_vals[ii] is None:
@@ -770,9 +769,10 @@ def _apply_neumann_and_robin_boundary_conditions(
             # (D2*u)*n2+D2*u*n2
             jac[idx] = sum(flux_normal_vals)
             bndry_vals = bndry_cond[0](mesh_pts_idx)[:, 0]
-            
+
             # residual[idx] = torch.linalg.multi_dot((jac[idx], sol))-bndry_vals
-            residual[idx] = torch.linalg.multi_dot((self._bndry_slice(jac, idx, 0), sol))-bndry_vals
+            residual[idx] = torch.linalg.multi_dot(
+                (self._bndry_slice(jac, idx, 0), sol))-bndry_vals
             if bndry_cond[1] == "R":
                 jac[idx, idx] += bndry_cond[2]
                 # residual[idx] += bndry_cond[2]*sol[idx]
@@ -813,7 +813,7 @@ def _apply_boundary_conditions(self, bndry_conds, residual, jac, sol,
 class TransformedCollocationMesh(CanonicalCollocationMesh):
     # TODO need to changes weights of _get_quadrature_rule to account
     # for any scaling transformations
-    
+
     def __init__(self, orders, transform, basis_types=None):
 
         super().__init__(orders, basis_types)

Diff for: pyapprox/pde/autopde/physics.py

@@ -84,7 +84,8 @@ def _transient_residual(self, sol, time):
         return res, jac
 
     def _scalar_flux_jac(self, mesh, idx):
-        return [mesh._bndry_slice(mesh._dmat(dd), idx, 0) for dd in range(mesh.nphys_vars)]
+        return [mesh._bndry_slice(mesh._dmat(dd), idx, 0)
+                for dd in range(mesh.nphys_vars)]
 
     def _clear_data(self):
         # used for data that is the same for entire transient simulation