Skip to content

Composable scalar transforms #128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 64 commits into
base: master
Choose a base branch
from
Open

Conversation

Ickaser
Copy link

@Ickaser Ickaser commented Mar 6, 2025

This is very much still WIP, but took a stab at constructing something like discussed in #95 and #126 . As of yet, I haven't even tested that this runs, much less correctly; will come back when I have more time.

One major question: for things like mapping to negative reals, I am imagining Scale(-1) ∘ Exp(). In this case, should the log-Jacobian of Scale be log(abs(scale))?

@Ickaser
Copy link
Author

Ickaser commented Mar 14, 2025

I've added some scalar consistency tests now, which pass, with this new API. If there are some other obvious tests, I am open to suggestions. Still missing is any treatment of non-Real number types, as would be necessary for Unitful, but if you want to merge this and get feedback first before tackling the Unitful-friendly part this might be ready for something like that.

New questions

  • Namespace: to avoid name collisions, I decided to rename Exp, Scale, and Shift to TVExp, TVScale, and TVShift. That way they can be exported and maintain a shorter name than TransformVariables.Exp. I'm of course open to other prefixes or ways to keep it readable.
  • Should unit transforming use the Scale operator (so allow Scale to take non-Real types), or give it its own separate transformation (which I guess could accommodate other kinds of exotic number types)?

@Ickaser
Copy link
Author

Ickaser commented Mar 14, 2025

Some other ideas I had but which haven't tried out yet (note to self as much as anything else)

  • We could try replacing the existing ShiftedExp and ScaledShiftedLogistic with these elemental transforms, but there may be some error/sanity checks we miss out on because a CompositeScalarTransform doesn't yet have end-to-end error checking like these manually-written ones do.
  • I haven't yet tried composing prior transformations with the new ones. If I'm lucky it will just work
  • It may make sense to introduce a TVNeg() transformation for a simple sign change, rather than using e.g. TVScale(-1) which encodes a number type in the -1.

@Ickaser
Copy link
Author

Ickaser commented Mar 20, 2025

I decided to add a TVNeg transformation, which if I am understanding correctly has a log-Jacobian of 0, and enforce that TVScale takes only positive numbers. Some tests are also in place and pass on my machine. I haven't gone so far as to replace existing transforms with the refactored form, which would put the new version through the full test suite.

I haven't thought deeply about what this current implementation will do with real number types, e.g. Float32, either in parameters or inputs.

As far as using this to introduce Unitful support, I see two main options:

  1. Create a Unitful extension with a single scalar transform, e.g. TVUnit(u::Unitful.Units). This transform would construct a Quantity in the forward direction, and apply ustrip for the inverse transform. This approach would have to be custom-implemented for each exotic number types (e.g. units provided by DynamicQuantities), but I think there would be a natural way to do so for any given type.
  2. Generalize the TVScale transform to allow any object; in this case, you multiply by units in the forward direction, then the inverse divides by the units. In the case of Unitful, this should produce the right types on the inverse. For other number types, multiplying should yield new types with the help of promotion, but the inverse would not be so simple.

In case 1, new number types each need their own extension. In case 2, inverse transforms will not yield original types.

@tpapp
Copy link
Owner

tpapp commented Mar 20, 2025

I think this is shaping up nicely, thanks for the work! Sorry I have been very busy with a project and failed to provide any feedback so far, I will try to make up for it now.

I am imagining Scale(-1) ∘ Exp(). In this case, should the log-Jacobian of Scale be log(abs(scale))?

Chaining just multiplies the Jacobians, so abs(log(...)) Jacobians are just added up. It should be something like log(abs(scale)) + x, I think your implementation is correct.

In case 2, inverse transforms will not yield original types.

I think that is OK, I think we should just clarify the API for inverse. In particular, the best the user can expect is a vector of equal length, approximately equal elementwise. When this PR is done I will make that clear.

@Ickaser
Copy link
Author

Ickaser commented Mar 25, 2025

I went ahead and tried out replacing the ShiftedExp and ScaledShiftedLogistic with the corresponding refactored transformations, and by and large the refactor works with the existing test suite. Exceptions that I marked with @test_broken:

  • as(Real, 1, 4.0) == as(Real, 1.0, 4.0)
    • The new transforms each hold their own number types (e.g. TVShift{Int64} vs TVShift{Float64}. Since my goal is to eventually accommodate new (unitful) number types, I don't know whether to continue to enforce this. I can define something like ==(TVShift(a), TVShift(b)) = (a==b) and run with that
  • Scalar show tests.
    • This is a valid question: I decided to show a CompositeScalarTransform by printing a valid expression (TVShift(5) ∘ TVExp()); would you prefer to represent it differently? Should there be special dispatches for e.g. asℝ₋, which is TVNeg() ∘ TVExp(), like there were for ShiftedExp and ShiftedScaledLogistic? In particular, if we have a special show for asℝ₊ == TVExp(), then we get results like TVExp() ∘ TVScale(2) printed like asℝ₊ ∘ TVScale(2), which I think I don't like.

@Ickaser
Copy link
Author

Ickaser commented Mar 25, 2025

In case 2, inverse transforms will not yield original types.

I think that is OK, I think we should just clarify the API for inverse. In particular, the best the user can expect is a vector of equal length, approximately equal elementwise. When this PR is done I will make that clear.

I guess if we take this approach and have a TVScale{T} with T that can take any type, it is always possible to make an extension with dispatches like transform(t::TVScale{T}, x) where T <: DynamicQuantities.Quantity. So the two options are not mutually exclusive, and option 2 provides a hopefully-sane fallback anyway.

@Ickaser
Copy link
Author

Ickaser commented Apr 9, 2025

With respect to the inverse methods all accepting Number My motivation for getting involved and building this PR is that I want to have transforms that add units by scaling, e.g. TVScale(5u"m"). So at the very least, the inverse of TVScale needs to accept a broad enough type to allow Unitful quantities, which are <:Number but not <:Real; in principle this can mean that on the inverse transform we get numbers that are not pure floats (although in the case of Unitful, when the units exactly cancel we get a plain float back).

@devmotion
Copy link
Collaborator

If we have asℝ₊ = TVExp(), then the string output of e.g. TVScale(5.0) ∘ TVExp() is TVScale(5.0) ∘ asℝ₊

That's the expected behaviour IMO, it should print TVScale(5.0) ∘ asℝ₊. It happens everywhere in the Julia ecosystem, e.g. Vector{T} is not a separate type but just an alias for Array{T,1}.

@tpapp
Copy link
Owner

tpapp commented Apr 10, 2025

@devmotion, thanks for the review. As for the following:

What impact does this change have on compilation time?

I did not benchmark, but I am OK with a minor hit in compilation time in exchange for more flexibility. Eg I want to introduce other R->R+ transforms, as exp is not always the best option.

What impact does this change have on performance?

Benchmarks above suggest it is innocuous.

What impact does this change have on numerical accuracy?

It just does what those transformations did previously, only in composable steps.

@Ickaser
Copy link
Author

Ickaser commented Apr 16, 2025

  • With the change to make asR+ = TVExp() and asI = TVLogistic(), this leaves asR- isa CompositeScalarTransform{Tuple{TVNeg, TVExp}} . So that won't get printed as a single element transform. This seems fine to me.
  • TransformVariables.compose is now just an alias to CompositionsBase.compose; I have not exported compose but we could mark it with public.
  • The old implementation of ShiftedExp and ScaledShiftedLogistic is gone, including from the tests and docs
  • I added some docs explaining how one would add new methods for TVScale that create number types where multiplying/dividing isn't desired
  • We could conceivably add some sort of convenience constructor like as(Real, 5u"m", \infty) that would automatically add units. But at the moment I can't think of a non-awkward syntax for that so I'll let that go for now, maybe make a future PR if I hit on something good.

@Ickaser
Copy link
Author

Ickaser commented Apr 18, 2025

After trying to use this in my code base, the two things I found myself wanting were:

  • as(Real, 1u"m", 5u"m")
  • A ScalarTransform that allows for a variable to be wrapped in some other function (or output type), e.g.
struct TVConstWrap <: ScalarTransform end
TransformVariables.transform(::TVConstWrap, x) = MyConstType(x)
TransformVariables.inverse(::TVConstWrap, x) = x.value

(although I did already add docs indicating this could be done by extending the TVScale transformation...)

That said, those can go into a future PR, and I'm content with this PR.

Copy link
Owner

@tpapp tpapp left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just have a cosmetic docs change, if you are OK with that please commit and I will merge. Thanks for all the great work! I especially like how Unitful etc does not require special treatment.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants