Optimizing boundary condition for given 2D solution

Hello, I am trying to use torch-fenics to compute gradients with respect to input array which is imposed as a Dirichlet boundary condition on the top surface of the domain. I have written a simple solver in FEniCS which extrudes the given 1 dimensional input to a 2D domain by solving du/dy = 0, subject to boundary condition u(x) = some input array on the top boundary. To achieve this, I create a 1D function space using SubMesh on top boundary of the 2D mesh and I use this 1D function space to represent the input function in input_templates() for torch_fenics module. Then I pass a 1D torch tensor as input to the solver. Here is a minimal example of my implementation:

import numpy as np
import dolfin as dl
import torch
from fenics import *
from fenics_adjoint import *
import torch_fenics

class Top(SubDomain):
    def inside(self, x, on_boundary):
        return near(x[1], 1.0, DOLFIN_EPS)

class Extrude(torch_fenics.FEniCSModule):
    def __init__(self):
        super().__init__()
        # define the mesh and mesh function
        self.mesh = UnitSquareMesh(nx=10, ny=10)
        self.mesh_function = MeshFunction("size_t", self.mesh, 1)
        self.mesh_function.set_all(0)

        # mark top boundary
        self.top = Top()
        self.top.mark(self.mesh_function, 1)

        self.scalar_function_space = FunctionSpace(self.mesh, "Lagrange", 1)

        # create 1D function space corresponding to top boundary top used for input_templates
        bm = BoundaryMesh(self.mesh, "exterior")
        self.mesh_1d = SubMesh(bm, self.top)
        self.function_space_1d = FunctionSpace(self.mesh_1d, "Lagrange", 1)
        self.u = TrialFunction(self.scalar_function_space)
        self.v = TestFunction(self.scalar_function_space)
        self.u_sol = Function(self.scalar_function_space)

    def solve(self, T_in):
        T_in.set_allow_extrapolation(True)
        self.T_in = T_in
        self.bc1 = DirichletBC(self.scalar_function_space, self.T_in, self.top, "pointwise")

        # define the weak for du/dy = 0
        weak_form = grad(self.u)[1] * self.v * dx
        weak_lhs = lhs(weak_form)
        weak_rhs = rhs(weak_form)

        solve(weak_lhs == weak_rhs, self.u_sol, [self.bc1])

        return self.u_sol

    def input_templates(self):
        return Function(self.function_space_1d)

extrude = Extrude()

# set the input array (top boundary values)
u_1d = np.linspace(1.0, 0.0, 11, dtype=np.float64) + 1.0
u_1d = torch.tensor(u_1d, requires_grad=True, dtype=torch.float64).reshape(1,-1)

u_2d = extrude(u_1d)
J = u_2d.sum()
print(f"J = {J}")
J.backward()
print(u_1d.grad)

The solver seems to work fine and it extrudes the input array to the 2D domain as expected, but when calling J.backward(), I get the following error saying FunctionSpaces are expected to be on same mesh:

J = 181.5
Traceback (most recent call last):
  File "test.py", line 60, in <module>
    J.backward()
  File "~/anaconda3/envs/fenics/lib/python3.7/site-packages/torch/_tensor.py", line 396, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "~/anaconda3/envs/fenics/lib/python3.7/site-packages/torch/autograd/__init__.py", line 175, in backward
    allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward pass
  File "~/anaconda3/envs/fenics/lib/python3.7/site-packages/torch/autograd/function.py", line 253, in apply
    return user_fn(self, *args)
  File "~/anaconda3/envs/fenics/lib/python3.7/site-packages/torch_fenics/torch_fenics.py", line 96, in backward
    tape=ctx.tape, adj_value=adj_value)
  File "~/anaconda3/envs/fenics/lib/python3.7/site-packages/pyadjoint/drivers.py", line 28, in compute_gradient
    tape.evaluate_adj(markings=True)
  File "~/anaconda3/envs/fenics/lib/python3.7/site-packages/pyadjoint/tape.py", line 140, in evaluate_adj
    self._blocks[i].evaluate_adj(markings=markings)
  File "~/anaconda3/envs/fenics/lib/python3.7/site-packages/pyadjoint/tape.py", line 47, in wrapper
    return function(*args, **kwargs)
  File "~/anaconda3/envs/fenics/lib/python3.7/site-packages/pyadjoint/block.py", line 134, in evaluate_adj
    prepared)
  File "~/anaconda3/envs/fenics/lib/python3.7/site-packages/fenics_adjoint/types/dirichletbc.py", line 143, in evaluate_adj_component
    r = compat.extract_bc_subvector(adj_value, c.function_space(), bc)
  File "~/anaconda3/envs/fenics/lib/python3.7/site-packages/fenics_adjoint/types/compat.py", line 218, in extract_bc_subvector
    assigner = backend.FunctionAssigner(Vtarget, backend.FunctionSpace(bc.function_space()))
RuntimeError: 

*** -------------------------------------------------------------------------
*** DOLFIN encountered an error. If you are not able to resolve this issue
*** using the information listed below, you can ask for help at
***
***     fenics-support@googlegroups.com
***
*** Remember to include the error message listed below and, if possible,
*** include a *minimal* running example to reproduce the error.
***
*** -------------------------------------------------------------------------
*** Error:   Unable to create function assigner.
*** Reason:  Expected all FunctionSpaces to be defined over the same Mesh.
*** Where:   This error was encountered inside FunctionAssigner.cpp.
*** Process: 0
*** 
*** DOLFIN version: 2019.1.0
*** Git changeset:  c5b9b269f4a6455a739109e3a66e036b5b8412f5
*** -------------------------------------------------------------------------

I am quite new to both FEniCS and Torch-FEniCS, and I would greatly appreciate any help in solving this issue. Also, please let me know if there is a better way to achieve this kind of optimization with respect to the boundary condition than creating 1D SubMesh and FunctionSpace as input.

Thanks in advance, Shashwat

barkm / torch-fenics

Optimizing boundary condition for given 2D solution #5