google / tangent

Source-to-Source Debuggable Derivatives in Pure Python
Apache License 2.0
2.31k stars 434 forks source link

Incorrect gradient calculation #64

Closed thearn closed 6 years ago

thearn commented 6 years ago

Hi, I've been experimenting with this library for a few months now and really like the capabilities present. I'm working to develop an AD capability via code generation for a set of aerospace engineering codes in Python.

However, I think I've run into either a bug or a usage misunderstanding. Consider the following stand-alone function, which takes several parameters and returns a scalar:

import tangent

BTU_s2HP, HP_per_RPM_to_FT_LBF = 1.4148532, 5252.11

def enthalpyandpower(W_in, W_out, ht_in, ht_out_ideal, eff, Nmech, b1_W, b1_ht, b1_ht_ideal):

    ht_out = W_in/W_out * (ht_in * (1.0 - eff) + ht_out_ideal * eff)
    power = W_in * eff * (ht_in - ht_out_ideal) * BTU_s2HP

    ht_out += b1_W / W_out * \
        (b1_ht * (1.0 - eff) +
         b1_ht_ideal * eff)
    power += b1_W * eff * \
        (b1_ht - b1_ht_ideal) * BTU_s2HP

    # calculate torque based on revised power and shaft speed
    trq = power / \
        Nmech * HP_per_RPM_to_FT_LBF

    return power

If I generate the partial derivative of power with respect to the first parameter W_in

dpower_dwin = tangent.autodiff(enthalpyandpower, wrt=(0,), verbose=1)

then I get:

def denthalpyandpowerdW_in(W_in, W_out, ht_in, ht_out_ideal, eff, Nmech,
    b1_W, b1_ht, b1_ht_ideal, bpower):
    # Initialize the tape
    _stack = tangent.Stack()
    _ht_out3 = ht_out_ideal * eff
    _1_0_minus_eff = 1.0 - eff
    _ht_out2 = ht_in * _1_0_minus_eff
    _ht_out = _ht_out2 + _ht_out3
    W_in_over_W_out = W_in / W_out
    ht_out = W_in_over_W_out * _ht_out
    _power2 = ht_in - ht_out_ideal
    W_in_times_eff = W_in * eff
    _power = W_in_times_eff * _power2
    power = _power * BTU_s2HP
    tangent.push(_stack, ht_out, '_1c132dd6')
    _3285 = b1_ht - b1_ht_ideal
    b1_W_times_eff = b1_W * eff
    _3244 = b1_W_times_eff * _3285
    _eb3b = _3244 * BTU_s2HP
    tangent.push(_stack, power, '_b65b4e60')
    power = power + _eb3b
    assert tangent.shapes_match(power, bpower
        ), 'Shape mismatch between return value (%s) and seed derivative (%s)' % (
        numpy.shape(power), numpy.shape(bpower))
    power = tangent.pop(_stack, '_b65b4e60')
    bpower = tangent.init_grad(power, allow_lazy_initializer=True)
    ht_out = tangent.pop(_stack, '_1c132dd6')
    bht_out = tangent.init_grad(ht_out, allow_lazy_initializer=True)

    # Grad of: power = W_in * eff * (ht_in - ht_out_ideal) * BTU_s2HP
    _b_power = tangent.unbroadcast(bpower * BTU_s2HP, _power)
    b_power = _b_power
    _3f78 = tangent.unbroadcast(b_power * _power2, W_in_times_eff)
    bW_in_times_eff = _3f78
    _bW_in2 = tangent.unbroadcast(bW_in_times_eff * eff, W_in)
    bW_in = _bW_in2

    # Grad of: ht_out = W_in / W_out * (ht_in * (1.0 - eff) + ht_out_ideal * eff)
    _a32f = tangent.unbroadcast(bht_out * _ht_out, W_in_over_W_out)
    _9f3c = _a32f
    _bW_in = _9f3c / W_out
    bW_in = tangent.add_grad(bW_in, _bW_in)
    return bW_in

Running this seems to give me a partial derivative of 0.0 regardless of the evaluation point. However, the partial is certainly non-zero, e.g. at (30., 30., 10., 9.5, 0.95, 1000., 1000., 1000., 999.) it would be about 0.67206, but instead

x = denthalpyandpowerdW_in(30., 30., 10., 9.5, 0.95, 1000., 1000., 1000., 999., 1.0)
print(x)

returns 0.0.

The correct derivative of can be found analytically with some work, or confirmed roughly by finite difference:

x0 = enthalpyandpower(30., 30., 10., 9.5, 0.95, 1000., 1000., 1000., 999.)
x1 = enthalpyandpower(30.001, 30., 10., 9.5, 0.95, 1000., 1000., 1000., 999.)

print((x1 - x0) / (0.001))

Have I made a user error, or is this an unexpected bug? Thanks!

mdanatg commented 6 years ago

It looks like a bug. I was able to get a value that looks similar to the finite difference by replacing the += operators, e.g. ht_out = ht_out + .... Could you try to see if that fixes things for you? This would be just a temporary workaround, until we get to fix the bug.

thearn commented 6 years ago

Sure, I can use the workaround for now, I appreciate the feedback.

alexbw commented 6 years ago

Definitely a bug, thanks for reporting it, and great catch! We were handling augmented assignments incorrectly in one of our pre-processing passes (anf.py). Should be fixed by #65 , which I hope to merge to master if tests pass.

alexbw commented 6 years ago

One other nit, just an FYI for myself: it looks like trq and the aug assign to ht_power are unused in this calculation (it was the power += ... statement that actually caused the gradient bug). Is that intentional, or an artifact of yanking code out of your codebase to give us a running example?

thearn commented 6 years ago

Yeah, It was originally a multi-output function that I tweaked just to get a running example. I may have more questions about best practices surrounding more specific use cases some time soon, but figured I would at least report this bug. Thanks again for the feedback!