Derive logprob for hyperbolic and error transformations #6664

LukeLB · 2023-04-10T22:09:41Z

What is this PR about?
I have implemented additional Elemwise transformations as suggested in issue #6631. Specifically, this pull request adds cosh, sinh, tanh, erf, erfc, and erfcx functions. I plan to address the other suggested transformations in a separate pull request, as they require a more significant rewrite of existing functions. However, if it is preferred to include them all in one pull request, I'm happy to do so.

Please note that this is still a work in progress, and I have not yet written any tests for the new Transforms. I would appreciate some guidance on how to design these tests as its not clear to me what I should be testing them against.

Also for the erfcx transform it would be great double check my math is correct, for the backward I have rewrote a matlab function and for the log jacobian determinant I used wolfram alpha to get the derivative of erfcx.
...

Checklist

Explain important implementation details 👆
Make sure that the pre-commit linting/style checks pass.
Link relevant issues (preferably in nice commit messages)
Are the changes covered by tests and docstrings?
Fill out the short summary sections 👇

Major / Breaking Changes

New elemwise transforms
Cleaned up the if block in find_measureable_transforms() as it was getting quite large

New features

Transforms for:

cosh
sinh
tanh
erf
erfc
erfcx

Bugfixes

NA

Documentation

Haven't added any in this PR

Maintenance

NA

📚 Documentation preview 📚: https://pymc--6664.org.readthedocs.build/en/6664/

codecov · 2023-04-10T22:23:37Z

Codecov Report

Merging #6664 (11d41db) into main (b7764dd) will increase coverage by 0.03%.
The diff coverage is 88.73%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6664      +/-   ##
==========================================
+ Coverage   91.96%   92.00%   +0.03%     
==========================================
  Files          94       95       +1     
  Lines       15927    16101     +174     
==========================================
+ Hits        14647    14813     +166     
- Misses       1280     1288       +8

Impacted Files	Coverage Δ
pymc/logprob/transforms.py	`94.95% <88.73%> (-0.93%)`	⬇️

... and 12 files with indirect coverage changes

pymc/logprob/transforms.py

ricardoV94 · 2023-04-11T06:23:03Z

Please note that this is still a work in progress, and I have not yet written any tests for the new Transforms. I would appreciate some guidance on how to design these tests as its not clear to me what I should be testing them against.

Before I have just tested the logp matches against equivalent RVs forms such as abs(normal) == halfnormal, but that won't work here AFAICT :)

So then it boils down to:

testing the new transforms have the right log_jac_det, which you can do with resource to

pymc/tests/distributions/test_transform.py

Line 81 in 5d68bf3

def check_jacobian_det(
Test the logprob derivation is working, something like:

def test_erf_logp():
  base_rv = pt.random.normal(0.5, 1, name="base_rv")  # Something not centered around 0 is usually better
  rv = pt.erf(base_rv)
  vv = rv.clone()

  rv_logp = logp(rv, vv)
  assert_no_rvs(rv_logp)

  transform = ErfTransform
  expected_logp = logp(rv, transform.backward(vv)) + transform.log_jac_det(vv)

  vv_test = np.array(0.25)  # Arbitrary test value
  np.testing.assert_almost_equal(
    rv_logp.eval({vv: vv_test}),
    expected_logp.eval({vv: vv_test}),
  )

You can probably parametrize and test all new functions with the same test.

Alternatively you can try to hijack

pymc/tests/logprob/test_transforms.py

Line 212 in 5d68bf3

def test_transformed_logprob(at_dist, dist_params, sp_dist, size):

The test now assumes you are testing only _default_transform, but you could make it accept a non default transform. Everything else should work the same?

LukeLB · 2023-04-15T15:27:08Z

@ricardoV94 I've been working on 2. and the test now runs however its throwing an assertion error. The test is:

@pytest.mark.parametrize("transform", [ErfTransform])
def test_erf_logp(transform):
    base_rv = pt.random.normal(0.5, 1, name="base_rv")  # Something not centered around 0 is usually better
    rv = pt.erf(base_rv)
    vv = rv.clone()
    rv_logp = joint_logprob({rv: vv})

    transform = transform()
    expected_logp = joint_logprob({rv: transform.backward(vv)}) + transform.log_jac_det(vv)

    vv_test = np.array(0.25)  # Arbitrary test value
    np.testing.assert_almost_equal(
    rv_logp.eval({vv: vv_test}),
    expected_logp.eval({vv: vv_test}),
    )

This gives the assertion error:

E       AssertionError: 
E       Arrays are not almost equal to 7 decimals
E       
E       Mismatched elements: 1 / 1 (100%)
E       Max absolute difference: 0.06346299
E       Max relative difference: 0.07601085
E        x: array(-0.898383)
E        y: array(-0.83492)

They're close but still quite a difference. I'm not sure if this is the way I've written the logp in the test or that the internal transform functions are wrong. Any ideas?

ricardoV94 · 2023-04-18T07:13:12Z

Oh my example was wrong. You want to compare with the base_rv + jacobian, not rv + jacobian:

This passes locally:

import numpy as np
import pytensor.tensor as pt
from pymc.logprob.basic import logp
from pymc.logprob.transforms import ErfTransform

base_rv = pt.random.normal(0.5, 1, name="base_rv")  # Something not centered around 0 is usually better
rv = pt.erf(base_rv)

vv = rv.clone()
rv_logp = logp(rv, vv)

transform = ErfTransform()
expected_logp = logp(base_rv, transform.backward(vv)) + transform.log_jac_det(vv)

vv_test = np.array(0.25)  # Arbitrary test value
np.testing.assert_almost_equal(
    rv_logp.eval({vv: vv_test}),
    expected_logp.eval({vv: vv_test}),
)

LukeLB · 2023-04-18T20:54:41Z

Cheers will make the change!

LukeLB · 2023-04-19T18:33:58Z

Oh my example was wrong. You want to compare with the base_rv + jacobian, not rv + jacobian:

This passes locally:

import numpy as np
import pytensor.tensor as pt
from pymc.logprob.basic import logp
from pymc.logprob.transforms import ErfTransform

base_rv = pt.random.normal(0.5, 1, name="base_rv")  # Something not centered around 0 is usually better
rv = pt.erf(base_rv)

vv = rv.clone()
rv_logp = logp(rv, vv)

transform = ErfTransform()
expected_logp = logp(base_rv, transform.backward(vv)) + transform.log_jac_det(vv)

vv_test = np.array(0.25)  # Arbitrary test value
np.testing.assert_almost_equal(
    rv_logp.eval({vv: vv_test}),
    expected_logp.eval({vv: vv_test}),
)

LukeLB · 2023-04-19T18:45:36Z

Woops clicked the wrong button didn't mean to close!

So test 2. now works for all Transforms however I'm having an issue with the check_jacobian_det which gives an assertion error for all transforms. Does this suggest that the math is wrong?

Note for test 2. I had to make changes to the test by adding a switch statement and editing the switch statement on line 416 in transforms.py to take input_logprob AND jacobian because of the descrepency of returning nans vs. -infs as if input_logprob is nan then this also returns nan and not -inf.

…sformations.

ricardoV94 · 2023-04-20T08:24:13Z

Woops clicked the wrong button didn't mean to close!

No problem :)

So test 2. now works for all Transforms however I'm having an issue with the check_jacobian_det which gives an assertion error for all transforms. Does this suggest that the math is wrong?

I think there must have been an error in your log_jac_det expression. I tweaked the default implementation in RVTransform so that it works for both elemwise and vector transforms, and (after allowing the test to accept nan) it passes, whereas with your hand-written implementation it did not.

I think it's fine to use the default implementation (for the cases where it works). I didn't try to find what was the error.

LukeLB · 2023-04-20T19:54:05Z

Great, I'll take a look at what you did and see if I can try and implement it with the other transforms.

LukeLB · 2023-04-20T20:44:19Z

OK all done, all tests pass.

pymc/logprob/transforms.py

ricardoV94

Looks good. I don't know why the coverage shows some of the new transforms not being covered, just a fluke?

I just have a question about a change below.

ricardoV94 · 2023-04-27T15:11:54Z

pymc/logprob/transforms.py

@@ -391,7 +419,7 @@ def measurable_transform_logprob(op: MeasurableTransform, values, *inputs, **kwa
        jacobian = jacobian.sum(axis=tuple(range(-ndim_supp, 0)))

    # The jacobian is used to ensure a value in the supported domain was provided
-    return pt.switch(pt.isnan(jacobian), -np.inf, input_logprob + jacobian)
+    return pt.switch(pt.isnan(input_logprob + jacobian), -np.inf, input_logprob + jacobian)


Can we revert this change?

Potentially let me check. The reason I did that was because it meant that we return -np.inf consistently when input_logprob = nan, which is the case for some of the transforms.

Okay seems that isn't the case anymore and tests pass with the reverted change

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

ricardoV94 · 2023-04-28T09:44:42Z

Awesome work @LukeLB! Looking forward to your next PR :)

LukeLB · 2023-04-28T10:10:08Z

Thanks @ricardoV94 it's been a pleasure! Thanks for reviewing :)

Luke LB added 5 commits April 4, 2023 21:02

added tranform classes for sinh cosh and tanh

f5ee4de

cleaned up if block in find_measurable_transform

d80d8ec

added an erfcx transform

592068f

added erf, erfc, erfcx now working in notebook

236e7d7

added a comment to erfcx backward fn and simplified the iteration code

05b8d55

ricardoV94 reviewed Apr 11, 2023

View reviewed changes

pymc/logprob/transforms.py Outdated Show resolved Hide resolved

used scan for erfcx and got test running

1825315

test now running and passing on all but erfc and erfcx, will fix

f847786

LukeLB closed this Apr 19, 2023

LukeLB reopened this Apr 19, 2023

test 2. now passing

b5e1517

ricardoV94 added 4 commits April 20, 2023 10:18

Simplify nan to ninf in test

3c5c733

Adapt default RVTransform.log_jac_det to univariate and vector tran…

b8bf86f

…sformations.

Use np.testing in check_jacobian_det

bcc1eb9

Use default log_jac_det in ErfTransform

43f4150

ricardoV94 marked this pull request as draft April 20, 2023 08:26

tests fixed, required removing handwritten log_jac_det

c429b2d

LukeLB marked this pull request as ready for review April 27, 2023 07:48

ricardoV94 reviewed Apr 27, 2023

View reviewed changes

pymc/logprob/transforms.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Apr 27, 2023

View reviewed changes

ricardoV94 changed the title ~~Additional Elemwise Transformations WIP~~ Derive logprob for hyperbolic and error transformations Apr 28, 2023

LukeLB and others added 2 commits April 28, 2023 09:36

Update pymc/logprob/transforms.py

2f3087c

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

reverted change to measurable_transform_logprob

11d41db

ricardoV94 approved these changes Apr 28, 2023

View reviewed changes

ricardoV94 merged commit d4bb701 into pymc-devs:main Apr 28, 2023
21 checks passed

ricardoV94 added enhancements logprob labels Apr 28, 2023

LukeLB mentioned this pull request Jul 13, 2023

Derive logprob for exp2, log2, log10, log1p, expm1, log1mexp, log1pexp (softplus), and sigmoid transformations #6826

Merged

5 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Derive logprob for hyperbolic and error transformations #6664

Derive logprob for hyperbolic and error transformations #6664

LukeLB commented Apr 10, 2023 •

edited

Loading

codecov bot commented Apr 10, 2023 •

edited

Loading

ricardoV94 commented Apr 11, 2023 •

edited

Loading

LukeLB commented Apr 15, 2023

ricardoV94 commented Apr 18, 2023

LukeLB commented Apr 18, 2023

LukeLB commented Apr 19, 2023

LukeLB commented Apr 19, 2023

ricardoV94 commented Apr 20, 2023 •

edited

Loading

LukeLB commented Apr 20, 2023

LukeLB commented Apr 20, 2023

ricardoV94 left a comment

ricardoV94 Apr 27, 2023

LukeLB Apr 28, 2023

LukeLB Apr 28, 2023

ricardoV94 commented Apr 28, 2023

LukeLB commented Apr 28, 2023

Derive logprob for hyperbolic and error transformations #6664

Derive logprob for hyperbolic and error transformations #6664

Conversation

LukeLB commented Apr 10, 2023 • edited Loading

Major / Breaking Changes

New features

Bugfixes

Documentation

Maintenance

codecov bot commented Apr 10, 2023 • edited Loading

Codecov Report

ricardoV94 commented Apr 11, 2023 • edited Loading

LukeLB commented Apr 15, 2023

ricardoV94 commented Apr 18, 2023

LukeLB commented Apr 18, 2023

LukeLB commented Apr 19, 2023

LukeLB commented Apr 19, 2023

ricardoV94 commented Apr 20, 2023 • edited Loading

LukeLB commented Apr 20, 2023

LukeLB commented Apr 20, 2023

ricardoV94 left a comment

Choose a reason for hiding this comment

ricardoV94 Apr 27, 2023

Choose a reason for hiding this comment

LukeLB Apr 28, 2023

Choose a reason for hiding this comment

LukeLB Apr 28, 2023

Choose a reason for hiding this comment

ricardoV94 commented Apr 28, 2023

LukeLB commented Apr 28, 2023

LukeLB commented Apr 10, 2023 •

edited

Loading

codecov bot commented Apr 10, 2023 •

edited

Loading

ricardoV94 commented Apr 11, 2023 •

edited

Loading

ricardoV94 commented Apr 20, 2023 •

edited

Loading