TSAdjoint

[salazardetroya]

Hi Ivan, this firedrake extension is awesome. Thanks for putting it together. I am interested in using TSAdjoint. I wonder if you have short-term plans to add it (less than two weeks or so). Otherwise, I'd be interested in making a contribution. I had the idea of writing a pyadjoint block so it can interact with pyadjoint as well. Although it is going to be tricky because TS/TSAdjoint also integrates the cost functions (somehow there are two blocks in one: the solve plus the cost function evaluation). Let me know what you think.

[IvanYashchuk]

Hello Miguel! I had plans to make TSAdjoint and TSForward work with pyadjoint, but I don't think I would do it in less than two weeks.
As 0th step in this process, we would need to make plain TSAdjoint work.
Here is an example file I was experimenting with:
https://github.com/IvanYashchuk/firedrake-ts/blob/adjoint/examples/adjoint.py
When adjoint solver needs to make more than one step it breaks (the output is zero or garbage).

Modifying a bit petsc4py's heat example also does not work, so probably the problem is not on the firedrake-petsc interface side.
https://gist.github.com/IvanYashchuk/f0a0474fae22bac1180e9e25fdccdd94

Do you have experience with implicit solver and adjoint stuff with petsc4py?

[salazardetroya]

Ivan, I ran into the same issue last week, and switching to Crank-Nicolson (CN) seemed to work. In your example, I also had to change the time step to something closer to the CFL condition:

ts.setTimeStep(ode.h ** 2 / 2.0 * 0.99)

setting the time step to exactly ode.h ** 2 / 2.0 does not seem to work in all cases. It also depends on whether I use a direct solver or the default solver and on the total number of time steps. This behavior seems a bit erratic unless there is some underlying theory that I am missing. It makes me wonder if it is the same reason why BEULER does not work. To sum up, here is a small table of what works with Crank Nicolson. By working I mean no divergence.

Time step	Number of time steps	Direct solver	Iterative solver	Works?
ode.h ** 2 / 2.0	20		X	Yes
ode.h ** 2 / 2.0	20	X		No
ode.h ** 2 / 2.0 * 0.99	20	X		Yes
ode.h ** 2 / 2.0 * 0.99	20		X	Yes
ode.h ** 2 / 2.0 * 0.99	100	X		Yes
ode.h ** 2 / 2.0	100	X		No
ode.h ** 2 / 2.0	100		X	No
ode.h ** 2 / 2.0 * 0.99	100		X	Yes

The stability issues are in the adjoint solve. Crank-Nicolson is unconditionally stable. The adjoint solve goes backwards in time so something in the stability might be different. I am going to share your code and these results to the petsc mail list if that is ok. I have experience with adjoint solvers, not so much using them in petsc4py.

[IvanYashchuk]

I think getting ADJOINT_DIVERGED_LINEAR_SOLVE is fine, as it indicates the problem with the ill-conditioning of the linear system.
That's great that it works with CN, for some reason I sticked with BEULER in my tests. Thanks for that!
I tested that further and found out that BEULER does not work at all. By that I mean that the adjoint vector is getting filled with zeros (and no error) or garbage large numbers (then ADJOINT_DIVERGED_LINEAR_SOLVE). I did test also with https://www.mcs.anl.gov/petsc/petsc-current/src/ts/tutorials/ex3.c.html, so the problem is not on petsc4py side.
I found out that the adjoint solve for THETA method works for all theta values except 1.0 (backward Euler case), even 0.999999 still works and everything breaks at 1.0.
Added those lines to the original example. Makefile is at $PETSC_DIR/src/ts/tutorials/.
Adjoint calculation is failing when run with
./ex3 -use_ifunc -ts_monitor -ts_adjoint_monitor -ts_type theta -ts_theta_theta 1.0
For other theta values it works (from 0.01 to 0.99999).

Could you please drop a link here for the mailing list thread once you send the message?

[IvanYashchuk]

I have found the bug in TS implementation in TSAdjointStepBEuler_Private. I will submit the fix. Now BEULER method also works as expected and gives the same result as pyadjoint-based example.

Now the future of pyadjoint block for TS Adjoint looks brighter to me.

[salazardetroya]

That's awesome! What are your thoughts for creating a pyadjoint block with TS Adjoint? Do you think it is ok to have one single block that both solves the transient problem and computes the cost function? I do no see another way.

[IvanYashchuk]

I think two separate cases need to be considered:

1. given initial conditions u0 and parameters m compute the solution state u at the final time. Then the task is to compute du/du0 and du/dm using TSAdjoint or TSForward. Once the pyadjoint block with evaluate_adj or evaluate_tlm is implemented for this function, u could be further used in transformations supported by pyadjoint.
2. given initial conditions u0 and parameters m compute the objective integral i. Then the task is to compute di/du0 and di/dm.

The second case requires additional wrappers for QuadratureTS to easily use expressions defined using Firedrake.
Once those two cases are implemented, performance optimizations can be considered so that we solve once the forward problem and then use the solution trajectory to solve two adjoint problems, one for du/d* and one for di/d*, if both the objective integral value and the solution state are needed in further differentiable computations.

[salazardetroya]

I have been thinking about how to implement the case for objective integrals i. TSAdjoint calculates the objective integral i using another TS and calculates the partial derivatives di/du0 and di/dm using matrices. These matrices have to be dense matrices (the number of columns is only as many as number of objective integrals to calculate). On the other hand, in Firedrake we will have a cost function defined by a 0-form, for instance i=u*dx and calculate its partial derivative with didu = derivative(i, u), when assembled, we will have a vector. Unfortunately, there is no way of assembling this into a dense matrix. I thought of using the Real space, but assembling with this space cannot return a Dense matrix. I think the way to go will be using MatDenseGetColumn to copy the result from assemble(didu) into a dense matrix. I ended up doing a PR in petsc4py https://bitbucket.org/petsc/petsc4py/pull-requests/150 to expose this function.

[IvanYashchuk]

Do these matrices have to be exactly of type dense though? I think it can be any Mat.
In Firedrake we can lift the cost function to 1-form using Real function space.

R = FunctionSpace(mesh, "R", 0)
vr = TestFunction(R)
i = u*vr*dx # now it is 1-form
# so we will get matrices when we assemble derivatives now
didu = assemble(derivative(i, u))

Assembled i will be a function on Real function space with size 1. Assembled didu is then a matrix of size 1 x u.function_space.dim() (R.dim() x u.function_space.dim())

I've started writing an example of using QuadratureTS
https://github.com/IvanYashchuk/firedrake-ts/blob/cost-adjoint/examples/time-distributed-control.py

But there is a bug somewhere now and I get the error

Assertion failed: (!PyErr_Occurred()), function __Pyx_PyCFunction_FastCall, file src/petsc4py.PETSc.c, line 339679.
Abort trap: 6

Could you take a look? Maybe we'll get it working together.

[salazardetroya]

I am not getting this error, which python version are you using? I have only seeing assertion errors from python when using python debug version.

I am getting another error though

  File "time-distributed-control.py", line 59, in form_cost_integrand
    with ctx._x.dat.vec_wo as v:
AttributeError: 'NoneType' object has no attribute '_x'

I think it is because the TS called within form_cost_integrand is not the same than the one inside DAESolver and it is not properly set up with the dmhooks. Somehow we need to attach the TSContext to the DM used by the quad_ts.

Previous times that I have used a sparse matrix, I ran into an error caused by this line 303 in https://www.mcs.anl.gov/petsc/petsc-current/src/ts/impls/implicit/theta/theta.c.html#TSAdjointStepBEuler_Private

[caidao22]

I have been thinking about how to implement the case for objective integrals i. TSAdjoint calculates the objective integral i using another TS and calculates the partial derivatives di/du0 and di/dm using matrices. These matrices have to be dense matrices (the number of columns is only as many as number of objective integrals to calculate). On the other hand, in Firedrake we will have a cost function defined by a 0-form, for instance i=u*dx and calculate its partial derivative with didu = derivative(i, u), when assembled, we will have a vector. Unfortunately, there is no way of assembling this into a dense matrix. I thought of using the Real space, but assembling with this space cannot return a Dense matrix. I think the way to go will be using MatDenseGetColumn to copy the result from assemble(didu) into a dense matrix. I ended up doing a PR in petsc4py https://bitbucket.org/petsc/petsc4py/pull-requests/150 to expose this function.

In PETSc, a vector can be placed into a dense matrix without copying the values explicitly:

  ierr = MatDenseGetColumn(A,0,&xarr);CHKERRQ(ierr);
  ierr = VecPlaceArray(sp,xarr);CHKERRQ(ierr);
  ierr = VecResetArray(sp);CHKERRQ(ierr);
  ierr = MatDenseRestoreColumn(A,&xarr);CHKERRQ(ierr);

In your PR to petsc4py, you might want to test if VecPlaceArray() and VecResetArray work in python as well.

[caidao22]

I am not getting this error, which python version are you using? I have only seeing assertion errors from python when using python debug version.

I am getting another error though

  File "time-distributed-control.py", line 59, in form_cost_integrand
    with ctx._x.dat.vec_wo as v:
AttributeError: 'NoneType' object has no attribute '_x'

I think it is because the TS called within form_cost_integrand is not the same than the one inside DAESolver and it is not properly set up with the dmhooks. Somehow we need to attach the TSContext to the DM used by the quad_ts.

Previous times that I have used a sparse matrix, I ran into an error caused by this line 303 in https://www.mcs.anl.gov/petsc/petsc-current/src/ts/impls/implicit/theta/theta.c.html#TSAdjointStepBEuler_Private

You are right that QuadratureTS does not inherit everything from the original TS. It is supposed to be a light-weight container for holding the call-back functions. Can you elaborate a bit about the problem with dmhooks? I can add them to QuadratureTS if needed.

[salazardetroya]

In PETSc, a vector can be placed into a dense matrix without copying the values explicitly:

  ierr = MatDenseGetColumn(A,0,&xarr);CHKERRQ(ierr);
  ierr = VecPlaceArray(sp,xarr);CHKERRQ(ierr);
  ierr = VecResetArray(sp);CHKERRQ(ierr);
  ierr = MatDenseRestoreColumn(A,&xarr);CHKERRQ(ierr);

In your PR to petsc4py, you might want to test if VecPlaceArray() and VecResetArray work in python as well.

What do you think of using MatDenseGetArray() directly instead of MatDenseGetColumn() as Lisandro suggested in the PR

You are right that QuadratureTS does not inherit everything from the original TS. It is supposed to be a light-weight container for holding the call-back functions. Can you elaborate a bit about the problem with dmhooks? I can add them to QuadratureTS if needed.

I am new at the DM management that Firedrake does to be honest. I think that they are attaching a context data structure so the TS has access to other data structures from Firedrake. Maybe @IvanYashchuk can ellaborate on this. I do wonder if it would be possible at all to set the same DM for both the QuadratureTS and the normal TS.

[caidao22]

In PETSc, a vector can be placed into a dense matrix without copying the values explicitly:

  ierr = MatDenseGetColumn(A,0,&xarr);CHKERRQ(ierr);
  ierr = VecPlaceArray(sp,xarr);CHKERRQ(ierr);
  ierr = VecResetArray(sp);CHKERRQ(ierr);
  ierr = MatDenseRestoreColumn(A,&xarr);CHKERRQ(ierr);

In your PR to petsc4py, you might want to test if VecPlaceArray() and VecResetArray work in python as well.

What do you think of using MatDenseGetArray() directly instead of MatDenseGetColumn() as Lisandro suggested in the PR

For dense matrices, 'MatDenseGetColumn()' is just a convenience function that does 'MatDenseGetArray' and a bit point manipulation. See
https://gitlab.com/petsc/petsc/-/blob/master/src/mat/impls/dense/seq/dense.c#L2498
But I think you still need VecPlaceArray() to work.

You are right that QuadratureTS does not inherit everything from the original TS. It is supposed to be a light-weight container for holding the call-back functions. Can you elaborate a bit about the problem with dmhooks? I can add them to QuadratureTS if needed.

I am new at the DM management that Firedrake does to be honest. I think that they are attaching a context data structure so the TS has access to other data structures from Firedrake. Maybe @IvanYashchuk can ellaborate on this. I do wonder if it would be possible at all to set the same DM for both the QuadratureTS and the normal TS.

Yes, it is not difficult to inherit any context from the original TS.

[IvanYashchuk]

I think it is because the TS called within form_cost_integrand is not the same than the one inside DAESolver and it is not properly set up with the dmhooks. Somehow we need to attach the TSContext to the DM used by the quad_ts.

We can simply set the DM for the QuadratureTS with quad_ts.setDM(problem.dm). Then the Firedrake level information can be accessed properly from the callbacks attached to quad_ts.

I now hit another error.

PETSc/TS.pyx in petsc4py.PETSc.TS.solve()

Error: error code 75
[0] TSSolve() line 4127 in /Users/yashchi1/devdir/firedrake/src/petsc/src/ts/interface/ts.c
[0] TSStep() line 3721 in /Users/yashchi1/devdir/firedrake/src/petsc/src/ts/interface/ts.c
[0] TSStep_Theta() line 223 in /Users/yashchi1/devdir/firedrake/src/petsc/src/ts/impls/implicit/theta/theta.c
[0] TSTheta_SNESSolve() line 185 in /Users/yashchi1/devdir/firedrake/src/petsc/src/ts/impls/implicit/theta/theta.c
[0] SNESSolve() line 4516 in /Users/yashchi1/devdir/firedrake/src/petsc/src/snes/interface/snes.c
[0] SNESSolve_NEWTONLS() line 175 in /Users/yashchi1/devdir/firedrake/src/petsc/src/snes/impls/ls/ls.c
[0] SNESComputeFunction() line 2379 in /Users/yashchi1/devdir/firedrake/src/petsc/src/snes/interface/snes.c
[0] SNESTSFormFunction() line 4983 in /Users/yashchi1/devdir/firedrake/src/petsc/src/ts/interface/ts.c
[0] SNESTSFormFunction_Theta() line 973 in /Users/yashchi1/devdir/firedrake/src/petsc/src/ts/impls/implicit/theta/theta.c
[0] VecAXPBYPCZ() line 684 in /Users/yashchi1/devdir/firedrake/src/petsc/src/vec/vec/interface/rvector.c
[0] Arguments are incompatible
[0] Incompatible vector global lengths parameter # 1 global size 1 != parameter # 5 global size 81

https://gitlab.com/petsc/petsc/-/blob/master/src/ts/impls/implicit/theta/theta.c#L942
Xdot has size 1 (size of the cost integrand vector) and X0 has size 81 (size of the state vector).

[salazardetroya]

It seems that calling quad_ts.setRHSFunction(form_cost_integrand) is actually setting the RHS function on the DM https://gitlab.com/petsc/petsc/-/blob/master/src/ts/utils/dmts.c#L523 and it is messing up with the function evaluations from the main TS. Calling quad_ts.setRHSFunction(form_cost_integrand) after solver.solve() runs without errors, but of course, it does not solve the problem.

[IvanYashchuk]

Thank you, that was very useful! I didn't know that the RHS function is saved in the DM and not the TS object.

I've updated the example. Now it seems to work, and something is computed. The correctness still needs to be verified though. And it doesn't work in parallel (form_djduand form_djdf need some more work).

The idea to use Real function space to get the matrices currently fails as @salazardetroya mentioned (when assembled the matrices are of type python)

Previous times that I have used a sparse matrix, I ran into an error caused by this line 303 in https://www.mcs.anl.gov/petsc/petsc-current/src/ts/impls/implicit/theta/theta.c.html#TSAdjointStepBEuler_Private

Also, petsc4py is missing a wrapper for TSSetIJacobianP.

[salazardetroya]

I've updated the example. Now it seems to work, and something is computed. The correctness still needs to be verified though. And it doesn't work in parallel (form_djduand form_djdf need some more work).

Awesome! I am going to play with it and try to verify it. I coded up a taylor_test like the one in pyadjoint for that.

Also, petsc4py is missing a wrapper for TSSetIJacobianP.

One can use TSSetRHSJacobianP instead. Check out the petsc manual beginning of page 154 https://www.mcs.anl.gov/petsc/petsc-current/docs/manual.pdf

If F() is a function of p one needs to also provide the Jacobian −Fp with

[salazardetroya]

Hi Ivan, I have worked out an analytical example for which I can obtain the correct numerical sensitivities. It is in https://gist.github.com/salazardetroya/a1e7b834ab6386a0ee6c9206e4549d7f

The problem is just a one dof ODE u_t = f* u; u(0) = a with known analytical solution of u(t) = a exp(f * t). The cost function time integrand is j = f * f * u * dx. Therefore the analytical cost function is just a * f * (exp(b * T) - 1.0) with T as final time. Note that in the code there is b which f=b, but f is a Function whereas b is a double. I did this to have f as a Function. The analytical sensitivities are easy to calculate from there.

In the PETSc Manual, they explain an additional step to calculate the sensitivities after the adjoint solve (https://www.mcs.anl.gov/petsc/petsc-current/docs/manual.pdf pag 153). This additional step seems to be necessary with CN method, but not with THETA method with theta=0.5 (same method than CN) as the correct sensitivity is already given by dJdf from the adjoint solve. This is confusing to me. Maybe @caidao22 can chime in. I also did not obtain the right sensitivities when setting the terminal conditions for dJdu and dJdf as in the manual. I had to leave them as zero.

I decided to work out this analytical example to narrow down the ways my code could fail. Once we understand better how the correct way to calculate the final sensitivities, I will expand to PDE's with a taylor test for derivatives checking.

[salazardetroya]

I also did not obtain the right sensitivities when setting the terminal conditions for dJdu and dJdf as in the manual. I had to leave them as zero.

This makes sense to me now because those terminal conditions do not depend on the integrand, but only on a scalar function of just the solution at the final time step. This is clear in the manual.

In the PETSc Manual, they explain an additional step to calculate the sensitivities after the adjoint solve (https://www.mcs.anl.gov/petsc/petsc-current/docs/manual.pdf pag 153). This additional step seems to be necessary with CN method, but not with THETA method with theta=0.5 (same method than CN) as the correct sensitivity is already given by dJdf from the adjoint solve

I also understand this better. The additional step described in the manual (in the equation \frac{\partial \Psi}{\partial p}) refers to the derivative of the initial condition w.r.t the design parameter p, instead of the entire residual as I was doing. Therefore, the method with THETA is giving the actual results. I am going to see what is happening with CN though.

[salazardetroya]

Ok, giving that CN is just THETA with 0.5 and TSThetaSetEndpoint() set to True (was missing this last part), now I can obtain same results whether I use CN or THETA.

The question now is to see why TSThetaSetEndpoint() to True returns the wrong derivative. This should not be hard to figure out. My intuition is that this is just not implemented.

[caidao22]

Hi Ivan, I have worked out an analytical example for which I can obtain the correct numerical sensitivities. It is in https://gist.github.com/salazardetroya/a1e7b834ab6386a0ee6c9206e4549d7f

The problem is just a one dof ODE u_t = f* u; u(0) = a with known analytical solution of u(t) = a exp(f * t). The cost function time integrand is j = f * f * u * dx. Therefore the analytical cost function is just a * f * (exp(b * T) - 1.0) with T as final time. Note that in the code there is b which f=b, but f is a Function whereas b is a double. I did this to have f as a Function. The analytical sensitivities are easy to calculate from there.

In the PETSc Manual, they explain an additional step to calculate the sensitivities after the adjoint solve (https://www.mcs.anl.gov/petsc/petsc-current/docs/manual.pdf pag 153). This additional step seems to be necessary with CN method, but not with THETA method with theta=0.5 (same method than CN) as the correct sensitivity is already given by dJdf from the adjoint solve. This is confusing to me. Maybe @caidao22 can chime in. I also did not obtain the right sensitivities when setting the terminal conditions for dJdu and dJdf as in the manual. I had to leave them as zero.

I decided to work out this analytical example to narrow down the ways my code could fail. Once we understand better how the correct way to calculate the final sensitivities, I will expand to PDE's with a taylor test for derivatives checking.

It is generally needed regardless of the TS method you use. But for certain cases (actually many cases in practice) you can skip it. Note that dJdf (lambda) and dJdp (mu) calculated by TSAdjoint are partial derivatives. You need to apply the chain rule to compute the total derivative. If the initial condition does not depend on your control parameter p, the term \frac{d y_0} {d p} will be zero, so the total derivative will be equal to the partial derivative dJdp according to this equation.

[salazardetroya]

@IvanYashchuk looks like we're in good shape now. I understand what TSAdjoint is doing in terms of calculating the sensitivities and @caidao22 has pushed a PR to fix the issue with the CN solver (https://gitlab.com/petsc/petsc/-/merge_requests/2830) I'll take a crack at it this week and try to put together the interface for calculating sensitivities with TSAdjoint when a time integral is given.

[IvanYashchuk]

petsc4py's master branch now has a wrapper for TSSetIJacobianP.

I think for the pyadjoint Block partial derivatives are what we actually need. If there is a dependence of initial condition on the control parameter, other pyadjoint blocks will take care of applying the chain rule.

[salazardetroya]

Hi @IvanYashchuk , I have updated the time-distributed-example.py and now returns the correct sensitivities (verified with taylor test) and also works in parallel. It is in https://github.com/salazardetroya/firedrake-ts/blob/cost-adjoint/examples/time-distributed-control.py

Now to refactor everything inside the DAESolver and think of ways to interface with pyadjoint.

Would it be possible for me to push to this repo directly (but not to master)? It will be easier for me.

[IvanYashchuk]

Great progress!
I can try to write a pyadjoint block with the current petsc4py interface after the weekend. Now that you've done a major part of the job to verify the adjoints, I think it shouldn't take much effort.

Sure, I will give you access if you prefer that. You can always just open a pull request from your fork to this repository.

[salazardetroya]

Thanks! Would you mind if I am the one taking a crack at implementing the pyadjoint block? I want to gain experience with software development. I am pretty familiar with pyadjoint and its inner workings. I have dug through the code many times so it should not take me too much time. Then I can open a PR and we can review it.

[IvanYashchuk]

Sure, go ahead then.

[salazardetroya]

Hi @IvanYashchuk, I refactored the adjoint functionality into the DAESolver and I am going to start building the pyadjoint blocks. There are two things to bear in mind.

• TSAdjoint admits several cost functions. This is equivalent to having a pyadjoint block with several outputs. I have not seen an example in pyadjoint block with several outputs, but there seems to be some functionality that takes this into account here Do you have any thoughts on this?
• TSAdjoint puts together all the dependencies into a single vector (TSSetRHSJacobianP). Pyadjoint on the other hand breaks them apart and calculates the sensitivities separately here. I am thinking of using a MixedFunctionSpace to build the data structures necessary for the evaluations of the jacobians with respect to the parameters. What are your thoughts?

[IvanYashchuk]

I am actually not sure whether pyadjoint supports multiple output Blocks. Yes there is block.add_output but I don't understand what's the expected outcome of evaluate_adj_component for this case. There is overloaded numpy.ndarray, we can probably make use of it here, or a list of AdjFloat might work. It looks like pyadjoint is designed to work mainly through ReducedFunctional, that's always one real variable output cost function.
I think it's better to implement first only for one cost function case so the output of the custom Block will be AdjFloat.

About single vector parameter. You don't need to calculate the sensitivities separately. There is prepare_evaluate_adj, you can do all the PETSc Adjoint calculations in this method and save the result to any object you'd like (dict, list, whatever). You can assemble the full JacobianP there piece by piece manually on PETSc level, you have all inputs available. Then in evaluate_adj_component you use the information saved in prepared to just assign sensitivities to corresponding inputs using idx.

[salazardetroya]

The problem of doing the assembly manually on PETSc level is how to do the parallel partitioning. This is taken care of automatically by Firedrake when using a MixedFunctionSpacewith all the function spaces from the parameters. Doing it by hand with PETSc will be more difficult. Do you think there will be advantages that will make it worth it?

[IvanYashchuk]

I mean we still have to convert assembled derivative functions to the dense matrix for PETSc Adjoint.
Since Firedrake has already partitioned each component we can just use that and assemble one dense matrix for all parameters (stacking everything). If we have two parameters p1 and p2 and one functional J, for setRHSJacobianP we create one matrix of size ( sizeof(dJdp1) + sizeof(dJdp2) ) x 1.
For example (from cost-adjoint branch):

    # let's just stack the same djdf as for two different parameters
    _djdf = fd.assemble(fd.derivative(j, f))
    local_size = _djdf.dat.data.shape[0]
    global_size = _djdf.ufl_function_space().dim()
    djdf_transposed_mat = PETSc.Mat().createDense(
        [[local_size + local_size, global_size + global_size], [1, 1]]
    )
    djdf_transposed_mat.setUp()


    def form_djdf(ts, t, X, Amat):
        dm = ts.getDM()
        ctx = dmhooks.get_appctx(dm)
        with ctx._x.dat.vec_wo as v:
            X.copy(v)
        ctx._time.assign(t)

        fd.assemble(fd.derivative(j, f), tensor=_djdf)
        Amat_array = Amat.getDenseArray()
        with _djdf.dat.vec_ro as v:
            # here we assign local parts to dense matrix
            Amat_array[0:local_size, 0] = v.array[:]
            Amat_array[local_size:(local_size+local_size), 0] = v.array[:]
        Amat.assemble()

With this approach, the user would need to bother creating the MixedFunctionSpace and then split the mixed function into individual parameters and it all can be messed up with using mixed_function.split() instead of ufl.split(mixed_function).

[salazardetroya]

This sounds like a plan. What about Constant variables? I guess we can stack them at end or at the beginning of the dense matrix. I need to figure out a way to do that.

[IvanYashchuk]

In Firedrake the assembled derivative of a functional wrt to a Constant variable is not a special case since it's an instance of firedrake.Function on Real function space.