Handwritten adjoints and removal of Zygote #107

michel2323 · 2021-02-25T03:28:06Z

This PR introduces handwritten adjoints for all the constraints. The line flow constraint adjoints exhibit the same issues than the Zygote ones: They disagree with ForwardDiff and FD about 1e-4. This issue has to be kept open. It is either a bug or some unknown numerical issue.

frapac

Just some preliminary comments on this PR. I need to write down the math to better understand what we are doing 😬

frapac · 2021-02-25T14:17:01Z

src/polar/kernels.jl

@@ -244,6 +307,52 @@ function transfer!(polar::PolarForm, buffer::PolarNetworkState, u)
    wait(ev)
 end

+KA.@kernel function adj_transfer_kernel!(
+    adj_vmag, adj_vang, adj_pinj, adj_qinj, adj_u, pv, pq, ref, adj_pload, adj_qload


I guess here adj_pload is the adjoint of the power generation pg (adj_pg), as the load is constant?

I am not sure. In that case we could remove that argument. I am not yet sure how to handle the adjoint! signatures in general. What we also have to clarify is weither the adjoint routines return a primal value. I would not recommend that.

frapac · 2021-02-25T14:17:24Z

src/polar/kernels.jl

@@ -130,14 +133,14 @@ KA.@kernel function residual_adj_kernel!(F, adj_F, v_m, adj_v_m, v_a, adj_v_a,
    if i > npv
        F[i + npq] -= qinj[fr]
    end
-    @inbounds for c in ybus_re_colptr[fr]:ybus_re_colptr[fr+1]-1
+    @inbounds for c in colptr[fr]:colptr[fr+1]-1


Nice cleansing!

frapac · 2021-02-25T14:18:57Z

src/polar/kernels.jl

+    nvbus = length(vm)
+    nnz = length(ybus_re.nzval)
+    T = eltype(vm)
+    edge_vm_from = zeros(T, nnz)


Do you know if we could reuse these variables to compute the adjoint of the line flow constraints?

Yes we could. I guess this belongs into the cache together with the adjoints.

michel2323 · 2021-02-26T15:41:15Z

Just some preliminary comments on this PR. I need to write down the math to better understand what we are doing 😬

I take shortcuts in the code; like for example when to use += etc. This is not the best example to learn manual differentiation.

* currently implemented for `power_balance` and `reactive_power_constraints` * homogeneize signature for AutoDiff.Hessian

frapac

Ok to merge once minor comments have been addressed

frapac · 2021-02-26T17:42:22Z

.github/workflows/action.yml

@@ -20,6 +20,8 @@ jobs:
      - uses: julia-actions/julia-runtest@latest

  test-moonshot:
+    env:
+      CUDA_VISIBLE_DEVICES: 1


Why setting this value? (I am curious 😺 )

This makes the tests on moonshot run on device=1 to avoid a clash with most users due to #110.

frapac · 2021-02-26T17:43:24Z

src/Polar/Constraints/line_flow.jl

    if isa(buffer.vmag, Array)
-        kernel! = accumulate_view!(KA.CPU())
+        kernel! = adj_branch_flow_kernel!(KA.CPU())


When polar is available in the scope, I would suggest defining the kernel directly as

kernel! = adj_branch_flow_kernel!(polar.device)

oh nice! indeed!

frapac · 2021-02-26T17:45:14Z

src/Polar/kernels.jl

@@ -1,4 +1,5 @@
 import KernelAbstractions: @index
+using ForwardDiff


I don't think we need to import ForwardDiff here, as it is already loaded in src/ExaPF.jl

frapac · 2021-02-26T17:46:47Z

src/Polar/kernels.jl


-"""
-    accumulate_view(x, vx, indices)
+    # not needed in the reverse run


Do we need this dead comment?

I prefer to have the forward code as a comment in the adjoint function. That's just a preference, we can remove it if you don't like it. It's easier for readability of the adjoint code if someone looks at it.

Ok to keep the comment, then!

frapac · 2021-02-26T17:46:55Z

src/Polar/kernels.jl

-        end
-    end
+    # not needed in the reverse run
+    # to_flow = vmag[to_bus]^2 * (


and this one?

frapac · 2021-02-26T17:47:29Z

src/Polar/kernels.jl

+    # slines[ℓ + nlines] = to_flow
+
+    adj_to_flow = adj_slines[ℓ + nlines]
+    adj_vmag[to_bus] += (2 * vmag[to_bus] * ytf_abs * vmag[fr_bus]^2


I would replace 2 by 2.0 (same 4 and 6)

You are right. I tried that actually yesterday to make sure that's not the cause of the numerical issues.

frapac · 2021-02-26T17:47:52Z

src/Polar/kernels.jl

+    adj_Δθ =   cosθ * adj_sinθ
+    adj_Δθ -=  sinθ * adj_cosθ
+    adj_vang[fr_bus] += adj_Δθ
+    adj_vang[to_bus] -= adj_Δθ


Nice job! This adjoint does not look trivial

That's why it's off!

michel2323 marked this pull request as draft February 25, 2021 03:28

michel2323 changed the title ~~Resolving condition in adjoint~~ [WIP] Resolving condition in adjoint Feb 25, 2021

frapac reviewed Feb 25, 2021

View reviewed changes

frapac force-pushed the ms/adjoint branch from 7c1848c to 9f76d08 Compare February 25, 2021 21:32

michel2323 and others added 3 commits February 26, 2021 11:01

draft adjoint line flow and reactive power

3e4bad0

polar: add generic adjoint! function for constraints

03dcb8f

* currently implemented for `power_balance` and `reactive_power_constraints` * homogeneize signature for AutoDiff.Hessian

polar: refactor AutoDiff.Hessian

f805b52

michel2323 force-pushed the ms/adjoint branch from 2b322e8 to ba217b5 Compare February 26, 2021 17:07

michel2323 changed the title ~~[WIP] Resolving condition in adjoint~~ [WIP] Handwritten adjoints and removal of Zygote Feb 26, 2021

michel2323 changed the title ~~[WIP] Handwritten adjoints and removal of Zygote~~ Handwritten adjoints and removal of Zygote Feb 26, 2021

michel2323 requested a review from frapac February 26, 2021 17:13

michel2323 marked this pull request as ready for review February 26, 2021 17:19

frapac reviewed Feb 26, 2021

View reviewed changes

michel2323 added 2 commits February 26, 2021 12:11

replaced zygote with handwritten line flow adjoints

d03ac7f

Run on device=1 on moonshot

84730e6

michel2323 force-pushed the ms/adjoint branch from 3fb4a6d to 84730e6 Compare February 26, 2021 18:13

michel2323 merged commit baf2be0 into develop Feb 26, 2021

This was referenced Feb 26, 2021

Implement reduced Hessian on GPU #108

Closed

Efficient transpose Jacobian vector product for reactive power generation #95

Closed

michel2323 deleted the ms/adjoint branch March 2, 2021 22:18

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Handwritten adjoints and removal of Zygote #107

Handwritten adjoints and removal of Zygote #107

michel2323 commented Feb 25, 2021 •

edited

Loading

frapac left a comment

frapac Feb 25, 2021

michel2323 Feb 26, 2021

frapac Feb 25, 2021

frapac Feb 25, 2021

michel2323 Feb 26, 2021

michel2323 commented Feb 26, 2021

frapac left a comment

frapac Feb 26, 2021

michel2323 Feb 26, 2021 •

edited

Loading

frapac Feb 26, 2021

michel2323 Feb 26, 2021

frapac Feb 26, 2021

michel2323 Feb 26, 2021

frapac Feb 26, 2021

michel2323 Feb 26, 2021

frapac Feb 26, 2021

frapac Feb 26, 2021

frapac Feb 26, 2021

michel2323 Feb 26, 2021

frapac Feb 26, 2021

michel2323 Feb 26, 2021

		@@ -1,4 +1,5 @@
		import KernelAbstractions: @index
		using ForwardDiff

Handwritten adjoints and removal of Zygote #107

Handwritten adjoints and removal of Zygote #107

Conversation

michel2323 commented Feb 25, 2021 • edited Loading

frapac left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

michel2323 commented Feb 26, 2021

frapac left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

michel2323 Feb 26, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

michel2323 commented Feb 25, 2021 •

edited

Loading

michel2323 Feb 26, 2021 •

edited

Loading