Add default function for DI + option to use in data_domain gradients (#163)

Author: moyner

src/Jutul.jl

@@ -40,8 +40,8 @@ module Jutul
     # AD
     import ForwardDiff
 
-    import DifferentiationInterface: AutoSparse, prepare_jacobian, jacobian
-    import SparseConnectivityTracer: TracerLocalSparsityDetector
+    import DifferentiationInterface: AutoSparse, prepare_jacobian, jacobian, AutoForwardDiff
+    import SparseConnectivityTracer: TracerLocalSparsityDetector, GradientTracer, IndexSetGradientPattern
     import SparseMatrixColorings: GreedyColoringAlgorithm
 
     # Timing

src/ad/AdjointsDI/adjoints.jl

@@ -93,9 +93,9 @@ end
 
 function setup_adjoint_storage_generic(X, F, packed_steps::AdjointPackedResult, G;
         state0 = missing,
-        backend = missing,
         do_prep = true,
         di_sparse = true,
+        backend = Jutul.default_di_backend(sparse = di_sparse),
         info_level = 0,
         single_step_sparsity = true,
         use_sparsity = true
@@ -285,18 +285,7 @@ end
 
 function setup_jacobian_evaluation!(storage, X, F, G, packed_steps, case0, backend, do_prep, single_step_sparsity, di_sparse)
     if ismissing(backend)
-        if di_sparse
-            gt = SparseConnectivityTracer.GradientTracer{SparseConnectivityTracer.IndexSetGradientPattern{Int, Set{Int}}}
-            sparsity_detector = TracerLocalSparsityDetector(gradient_tracer_type=gt)
-            # sparsity_detector = TracerLocalSparsityDetector()
-            backend = AutoSparse(
-                AutoForwardDiff();
-                sparsity_detector = sparsity_detector,
-                coloring_algorithm = GreedyColoringAlgorithm(),
-            )
-        else
-            backend = AutoForwardDiff()
-        end
+        backend = Jutul.default_di_backend(sparse = di_sparse)
     end
 
     H = AdjointObjectiveHelper(F, G, packed_steps)

src/variables/vectorization.jl

@@ -221,34 +221,70 @@ Compute the (sparse) Jacobian of parameters with respect to data_domain values
 (i.e. floating point values). Optionally, `config` can be passed to allow
 `vectorize_variables` to only include a subset of the parameters.
 """
-function parameters_jacobian_wrt_data_domain(model; copy = true, config = nothing)
-    if copy
-        model = deepcopy(model)
-    end
+function parameters_jacobian_wrt_data_domain(model;
+        copy = true,
+        config = nothing,
+        use_di = true,
+        backend = default_di_backend(),
+        prep = missing
+    )
     data_domain = model.data_domain
     x = vectorize_data_domain(data_domain)
-    x_ad = ST.create_advec(x)
-    devectorize_data_domain!(data_domain, x_ad)
-    prm = setup_parameters(model, perform_copy = false)
-    prm_flat = vectorize_variables(model, prm, :parameters, T = eltype(x_ad), config = config)
-    n_parameters = length(prm_flat)
-    n_data_domain_values = length(x)
-    # This is Jacobian of parameters with respect to data_domain
-    J = ST.jacobian(prm_flat, n_data_domain_values)
-    @assert size(J) == (n_parameters, n_data_domain_values)
+    if use_di
+        if copy
+            model = deepcopy(model)
+        end
+        function F(X)
+            model_tmp = deepcopy(model)
+            dd = model_tmp.data_domain
+            devectorize_data_domain!(dd, X)
+            prm = setup_parameters(model_tmp, perform_copy = false)
+            return vectorize_variables(model_tmp, prm, :parameters, T = eltype(X), config = config)
+        end
+        if ismissing(prep)
+            J = jacobian(F, backend, x)
+        else
+            J = jacobian(F, prep, backend, x)
+        end
+    else
+        x_ad = ST.create_advec(x)
+        devectorize_data_domain!(data_domain, x_ad)
+        prm = setup_parameters(model, perform_copy = false)
+        prm_flat = vectorize_variables(model, prm, :parameters, T = eltype(x_ad), config = config)
+        n_parameters = length(prm_flat)
+        n_data_domain_values = length(x)
+        # This is Jacobian of parameters with respect to data_domain
+        J = ST.jacobian(prm_flat, n_data_domain_values)
+        @assert size(J) == (n_parameters, n_data_domain_values)
+    end
     return J
 end
 
+function default_di_backend(; sparse = true)
+    if sparse
+        gt = GradientTracer{IndexSetGradientPattern{Int, Set{Int}}}
+        sparsity_detector = TracerLocalSparsityDetector(gradient_tracer_type=gt)
+        backend = AutoSparse(
+            AutoForwardDiff();
+            sparsity_detector = sparsity_detector,
+            coloring_algorithm = GreedyColoringAlgorithm(),
+        )
+    else
+        backend = AutoForwardDiff()
+    end
+    return backend
+end
+
 """
     data_domain_to_parameters_gradient(model, parameter_gradient; dp_dd = missing, config = nothing)
 
 Make a data_domain copy that contains the gradient of some objective with
 respect to the fields in the data_domain, assuming that the parameters were
 initialized directly from the data_domain via (`setup_parameters`)[@ref].
 """
-function data_domain_to_parameters_gradient(model, parameter_gradient; dp_dd = missing, config = nothing)
+function data_domain_to_parameters_gradient(model, parameter_gradient; dp_dd = missing, config = nothing, kwarg...)
     if ismissing(dp_dd)
-        dp_dd = parameters_jacobian_wrt_data_domain(model, copy = true, config = config)
+        dp_dd = parameters_jacobian_wrt_data_domain(model; copy = true, config = config, kwarg...)
     end
     do_dp = vectorize_variables(model, parameter_gradient, :parameters, config = config)
     # do/dd = do/dp * dp/dd