Clean up unit tests for benchmarkable code for 1st order

[jlperla]

The broad goals: (1) : remove the DSSM dependency from the tests after proving there is no weird interaction in the zygote gradients; (2) ensure simple unit tests exist so that if people tweak performance it won't accidentally break anything; and (3) prepare the exact setup we want to benchmark so that @jlperla can write a simple benchmark script.

At that point, this should be self-contained for anyone to look at performance of individual peices of the code and to run on julia 1.6 and 1.7.

The biggest here is to get the unit tests decoupled from DSSM which has way too many upstream interactions and dependencies (especially things holding back Julia 1.7 testing). But I don't think we need to add in new ones yet. We just don't want other people working on performance to accidentally break functionality.

So basically we need to get rid of https://github.com/SciML/DifferenceEquations.jl/blob/main/test/dssm.jl and convert it into first_order.jl or something along those lines. We can add in 2nd order later after we are done with first order. Part of doing first order is to see if a custom adjoint for the linear simulations and joint-likelihood is helpful relative to a gerneric one. If not, then 2nd order will follow relatively quickly and would be implemented in a very different way.

To get rid of the DSSM, then things I see as likely important are:

• [ ] Put https://github.com/HighDimensionalEconLab/DifferentiableStateSpaceModels.jl/blob/main/.JuliaFormatter.toml in the root of the repo. This will help automatically format literals in the unit tetss and we can always change the formatting later.
• [ ] Save off the A, B, C and the Diagonal(D), observables and whatever necessary for RBC example into literals in a unit test. I don't think we need to save them in files because they aren't that big.
  • I think we can use zeros for u0 and noise so I don't think those are necessary to save. But we donj't want anything randomized in the unit tests.
  • Just a reminder on this: YOu can get a copy-pasteable matrix doing something like the following, which menas you do not need to transcribe antyhing manually. The low precision in that output shouldn't be too important for these tests.

    A = [1.0 4.0; 6.0 7.8]
    @show A;

• [ ] Then we should add in a simple unit test that only checks the primal of the log likelihood and that it infers. e.g. something like
• [ ]

  
  const A_rbc = ...
  const B_rbc = ...
  const n_rbc = ...
  const n_ϵ_rbc = ...
  const T_rbc = ...
  const noise_rbc = [zeros(randn(n_rbc, n_ϵ_rbc) for _ in 1:T_rbc]
  const D_rbc = ... # The diagonal matrix of the noise
  #etc.

General function to evaluate joint likeliihoods. Can't pass in Distributions directly for unit tests

function joint_likelihood_1(A,B,C,u0,noise,observables,D) problem = LinearStateSpaceProblem(A, B, C, u0, (0,length(noise)); noise, observables, obj_noise = MvNormal(zeros(x.D), Diagonal(x.D))) return solve(problem, NoiseConditionalFilter(); save_everystep = false).loglikelihood end

@testset "linear rbc joint likelihood" begin @test joint_likelihood_1(A_rbc, B_rbc, C_rbc, u0_rbc, observables_rbc, D_rbc) ≈ WHATEVER IT IS @inferred joint_likelihood_1(x) # would this catch inference problems in the solve? test_rrule(Zygote.ZygoteRuleConfig(), joint_likelihood_1, A_rbc, B_rbc, u0_rbc, observables_rbc, D_rbc; rrule_f = rrule_via_ad) end

- [ ] Next we would want to make this work to find the gradient for a mock version of our likelihood.  See the above code as well.
- [x] Later we can be mroe careful on the inference issues of the solv itself and the constructor.  Maybe code like the following added to the above one, but also might not be necessary.
```julia
    problem = LinearStateSpaceProblem(x.A, x.B, x.C, x.u0, (0,length(x.noise)); noise = x.noise, obj_noise = MvNormal(SoMETHING WITH x.D), observables = x.observables)
    @ inferred LinearStateSpaceProblem(x.A, x.B, x.C, x.u0, (0,length(x.noise)); noise = x.noise, obj_noise = MvNormal(SoMETHING WITH x.D), observables = x.observables)
    @inferred solve(problem, NoiseConditionalFilter(); save_everystep = false)

• [ ] After this, we should do a sanity check that the timing of the joint_likelihood_1 for the rbc makes sense relative to our sampling speed.
  • This will require some thinking. Basically, we want to verify that this part of the gradient is the dominating problem for our performance after we do a sanity check on the DSSM timing for the gradient.
  • To do that, we can time the gradient of the existing code with the built-in DSSM and then compare it to one without. Or just look at the time to make one sample in our code and see if this is a significant part of it.
  • Assuming that we have verified that this is really the performance bottleneck right now, we can continue and just optimize this function on its own.
• [ ] Then with this, we do the same thing for rbc_kalman and then FVGQ joint 1 and FVGQ kalman
  • For FVGQ if the matrices are massive when we format them, it may make sense to save these as files in the tests and load them. THe easiest path is to do jointpath(pkgdir(DifferenceEquations), "tests/data/FVGQ_A.csv") or something along those lines.
  • But I think that with the @show A_rbc; and then copy paste, and then the code formatter these might be more compact than we would expect.
• [ ] With this @jlperla can set up more formal benchmarks next week, using these setups and we can discuss how to carefully check custom adjoints, etc.

Otherwise, I think a few important cleanup things before we ask anyone else to look at this (because many might be on Julia 1.7) are:

• [ ] Bump versions in https://github.com/SciML/DifferenceEquations.jl/blob/main/Project.toml#L17-L24 especially ChainRulesCore, CommonSovle, SciMLBase. Zygote is temporary until we can get proper @requires in there and see if we will implement custom adjoints, but that can wait.
• [ ] Strip out anything we can in the test dependencies that we are not using. DifferentiableStateSpaceModels in particular, etc. in https://github.com/SciML/DifferenceEquations.jl/blob/main/test/Project.toml I don't see why we would need anything more that Test, FiniteDiff and the AD packages after we move out Zygote properly, but I could be missing something.

[jlperla]

@wupeifan They told us the trick for the CRTU. Use the rrule_f argument.

using Zygote, LinearAlgebra, ChainRulesTestUtils
N = 50
A = rand(N, N)
b = rand(N)
h(A, b) = sum(A \ b)
test_rrule(Zygote.ZygoteRuleConfig(), h, A, b; rrule_f = rrule_via_ad)

It didn't work with a Distribution as an argument but I think that was because of type stability issues rather than CRTU having problems itself.

If we need to do things manually, we can always do

@test gradient(h, A,b)[1] ≈ FiniteDiff.finite_difference_gradient(h, A,b)

or something like that.