Skip to content
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

More work on VarNameVector #637

Merged
merged 76 commits into from
Sep 3, 2024
Merged

Conversation

mhauru
Copy link
Member

@mhauru mhauru commented Aug 8, 2024

The diff is currently a bit messy because I merged #575 into this. Will clean it up once #575 has been merged into master.

torfjelde and others added 30 commits April 16, 2024 08:48
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
latter can result in type-instability since transformed variables
typically result in non-view even if input is a view
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Tor Erlend Fjelde <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@mhauru
Copy link
Member Author

mhauru commented Aug 23, 2024

What exactly is the meaning of the return value of istrans(vi)? Almost all values in a vi are transformed in some way, even if it's only a trivial reshape, or wrapping a scalar in a singleton vector, so in some sense they are all istrans. Does istrans rather mean that they've been transformed into the full Euclidean space?

If yes, I might look into changing the data stored in a VarNamedVector to be something like a domain where the variable lives. Or if that seems like overkill, maybe rename the is_transformed field into something like is_constrained.

@yebai
Copy link
Member

yebai commented Aug 23, 2024

Does istrans rather mean that they've been transformed into the full Euclidean space?

yes

Or if that seems like overkill, maybe rename the is_transformed field into something like is_constrained.

I think constrained is more accurate, but need to be careful whether this interacts with #575

@mhauru
Copy link
Member Author

mhauru commented Aug 27, 2024

@torfjelde, is there a reason you want to keep the container types in VarNamedVector somewhat free, allowing any subtypes of AbstractVector? I'm somewhat tempted to make it be

struct VarNamedVector{K<:VarName,V,Trans}
    varname_to_index::OrderedDict{K,Int}
    varnames::Vector{K}
    ranges::Vector{UnitRange{Int}}
    vals::Vector{V}
    transforms::Vector{Trans}
    is_transformed::BitVector
    num_inactive::OrderedDict{Int,Int}
end

@mhauru
Copy link
Member Author

mhauru commented Aug 29, 2024

SimpleVarInfo{VarNamedVector} is now as well tested as SimpleVarInfo{Dict}. There are cases that it can't handle though, around some complicated use of lenses that probably won't occur often.

Tests pass except for

Now that things mostly work, I'll take another read through the code and this time hopefully actually understand it, and clean things up as I go.

Copy link
Member

@torfjelde torfjelde left a comment

Choose a reason for hiding this comment

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

Had a brief look, but will have to have a closer look tomorrow / over the weekend:) But looking neat! Good stuff @mhauru :)

src/varnamedvector.jl Show resolved Hide resolved
Comment on lines 80 to 92
# The various constructor(orig...) calls are to create concrete element types. E.g. if
# vnv.vals is a Vector{Real} but all elements are Float64s, [vnv.vals...] will be a
# Vector{Float64}
function VarNamedVector(vnv::VarNamedVector)
return VarNamedVector(
OrderedDict(vnv.varname_to_index...),
[vnv.varnames...],
copy(vnv.ranges),
[vnv.vals...],
[vnv.transforms...],
copy(vnv.is_transformed),
copy(vnv.num_inactive),
)
Copy link
Member

Choose a reason for hiding this comment

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

This seems somewhat controversial 😬 What's the motivation here?

At least, we should do map(identity, vnv.varnames), etc. instead of splatting.

Copy link
Member Author

Choose a reason for hiding this comment

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

I'm thinking of renaming this, maybe not make it be a constructor method. But the motivation for its existence is the need to be able to collect variables gradually, e.g. start with a VarNamedVector(), which has element type Real, and call push!! repeatedly, resulting in a VNV with abstract container element types, and then split it into multiple VNVs grouped by varname symbol to achieve type concreteness. This concretises the eltypes.

Is map(identity, x) better than [x...] in some way, or is this a style preference?

Copy link
Member Author

Choose a reason for hiding this comment

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

This has been renamed to tighten_types, to complement loosen_types!!.

Copy link
Member

Choose a reason for hiding this comment

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

Is map(identity, x) better than [x...] in some way, or is this a style preference?

I believe strictly better. IIRC [x...] will, AFAIK, result in [x[1], x[2], ...], i.e. as if you manually filled in the array with entries. If this is so, then this will put quite the burden on the compiler for large arrays.

Copy link
Member Author

Choose a reason for hiding this comment

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

Cool, thanks, will fix that.

src/varnamedvector.jl Outdated Show resolved Hide resolved
src/varnamedvector.jl Outdated Show resolved Hide resolved
src/varnamedvector.jl Outdated Show resolved Hide resolved
@mhauru
Copy link
Member Author

mhauru commented Sep 3, 2024

I think this branch, as a way to compare my changes against @torfjelde's, has outlived its usefulness. I'll merge this into Tor's branch, and continue development in #555, with master as the base branch for diffs.

@mhauru mhauru merged commit 19978ec into torfjelde/varnamevector Sep 3, 2024
@mhauru mhauru deleted the mhauru/varnamevector branch September 3, 2024 09:01
github-merge-queue bot pushed a commit that referenced this pull request Oct 8, 2024
#555)

* initial implementation of VarNameVector

* added some hacky getval and getdist get things to work for VarInfo

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* added arbitrary metadata field as discussed

* renamed idcs to varname_to_index

* renamed vns to varnames for VarNameVector

* added keys impl for Metadata

* added push! and update! for VarNameVector

* added getindex_raw! and setindex_raw! for VarNameVector

* added `iterate` and `convert` (for `AbstractDict) impls for `VarNameVector`

* make the key and eltype part of the `VarNameVector` type

* added more tests for VarNameVector

* formatting

* more testing for VarNameVector

* minor changes to some comments

* added a bunch more tests for VarNameVector + several bugfixes in the process

* formatting

* added `similar` implementation for `VarNameVector`

* formatting

* removed debug statement

* made VarInfo slighly more generic wrt. underlying metadata

* fixed incorrect behavior in `keys` for `Metadata`

* minor style changes to VarNameVector tests

* style

* added testing of `update!` with smaller sizes and fixed bug related to this

* formatting

* move functionality related to `push!` for `VarNameVector` into `push!`

* Update src/varnamevector.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* several fixes to make sampling with VarNameVector + initiall tests for
sampling with VarNameVector

* VarInfo + VarNameVector tests for all demo models

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* added docs on the design of `VarNameVector`

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* added note on `update!`

* further elaboration of the design of `VarInfo` and `VarNameVector`

* more writing improvements

* added docstring to `has_inactive_ranges` and `inactive_ranges_sweep!`

* moved docs on `VarInfo` design to a separate internals section

* writing improvements for internal docs

* further motivation of the design choices made in `VarNameVector`

* improved writing

* VarNameVector is now grown as much as needed

* updated `delete!`

* Significant changes to implementation of `VarNameVector`:
- "delete-by-mark" is now replaced by proper deletion.
- `inactive_ranges` replaced by `num_inactive`, which only keeps track
of the number of inactive entries for a given `VarName.
- `VarNameVector` is now a "grow-as-needed" structure where the
underlying also mimics the order that the user experiences.`

* added `copy` when constructing `VectorVarInfo` from `VarInfo`

* added missing `isempty` impl

* remove impl of `iterate` and instead implemented `pairs` and `values` iterators

* added missing `empty!` for `num_inactive`

* removed redundant `shift_left!` methd

* fixed `delete!` for `VarNameVector`

* added `is_contiguous` as an alterantive to `!has_inactive`

* updates to internal docs

* renamed `sweep_inactive_ranges!` to `contiguify!`

* improvements to internal docs

* more improvements to internal docs

* moved additional methods description in internals to earlier in the doc

* moved internals docs to a separate directory and split into files

* more improvements to internals doc

* formatting

* added tests for `delete!` and fixed reference to old method

* addition to `delete!` test

* added `values_as` impls for `VarNameVector`

* added docs for `replace_valus` and `values_as` for `VarNameVector`

* fixed doctest

* formatting

* temporarily disable doctests so we can build docs

* added missing compat entry for ForwardDiff in docs

* moved some shared code into methods to make things a bit cleaner

* added impl of `merge` for `VarNameVector`

* renamed a few variables in `merge` impl for `VarNameVector`

* forgot to include some changes in previous commit

* added impl of `subset` for `VarNameVector`

* fixed `pairs` impl for `VarNameVector`

* added missing impl of `subset` for `VectorVarInfo`

* added missing impl of `merge_metadata` for `VarNameVector`

* added a bunch of `from_vec_transform` and `tovec` impls to make
`VarNameVector` work with `Cholesky`, etc.

* make default args use `from_vec_transform` rather than `FromVec`

* fixed `values_as` fro `VarInfo` with `VarNameVector` as `metadata`

* fixed impl of `getindex_raw` when using integer index for `VarNameVector`

* added tests for `getindex` with `Int` index for `VarNameVector`

* fix for `setindex!` and `setindex_raw!` for `VarNameVector`

* introduction of `from_vec_transform` and `tovec` and its usage in `VarInfo`

* moved definition of `is_splat_symbol` to the file where it's used

* added `VarInfo` constructor with vector input for `VectorVarInfo`

* make `extract_priors` take the `rng` as an argument

* added `replace_values` for `Metadata`

* make link and invlink act on the `metadata` field for `VarInfo` +
implementations of these for `Metadata` and `VarNameVector`

* added temporary defs of `with_logabsdet_jacobian` and `inverse` for
`transpose` and `Bijectors.vec_to_triu`

* added invlink_with_logpdf overload for `ThreadSafeVarInfo`

* added `is_transformed` field to `VarNameVector`

* removed unnecessary defintions of `with_logabsdet_jacobian` and
`inverse` for `transpose`

* fixed issue where we were storing the wrong transformations in `VarNameVector`

* make sure `extract_priors` doesn't mutate the `varinfo`

* updated `similar` for `VarNameVector` and fixed `invlink` for `VarNameVector`

* added handling of `is_transformed` in `merge` for `VarNameVector`

* removed unnecesasry `deepcopy` from outer `link`

* updated `push!` to also `push!` on `is_transformed`

* skip tests for mutating linking when using VarNameVector

* use same projection for `Cholesky` in `VarNameVector` as in `VarInfo`

* fixed `settrans!!` for `VarInfo` with `VarNameVector`

* fixed bug in `set_flag!`

* fixed another typo

* fixed return values of `settrans!!`

* updated static transformation tests

* Update test/simple_varinfo.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* removed unnecessary impl of `extract_priors`

* make `short_varinfo_name` of `TypedVarInfo` a bit more informative

* moved impl of `has_varnamevector` for `ThreadSafeVarInfo`

* added back `extract_priors` impl as we do need it

* forgot to include tests for `VarNameVector` in `runtests.jl`

* fix for `relax_container_types` in `test/varnamevector.jl`

* fixed `need_transforms_relaxation`

* updated some tests to not refer directly to `FromVec`

* introduce `from_internal_transform` and its siblings

* remove `with_logabsdet_jacobian_and_reconstruct` in favour of
`with_logabsdet_jacobian` with `from_linked_internal_transform`, etc.

* added `internal_to_linked_internal_transform` + fixed a few bugs in
the linking as a resultt

* added `linked_internal_to_internal_transform` as a complement to `interanl_to_linked_interanl_transform`

* fixed bugs in `invlink` for `VarInfo` using `linked_internal_to_internal_transform`

* more work on removing calls to `reconstruct`

* removed redundant comment

* added `from_linked_vec_transform` specialization for `LKJ`

* more work on removing references to `reconstruct`

* added `copy` in `values_from_metadata` to preserve behavior and avoid
refs to internal representation

* remove `reconstruct_and_link` and `invlink_and_reconstruct`

* replaced references to `link_and_reconstruct` and `invlink_and_reconstruct`

* introduced `recombine` and replaced calls to `reconstruct` with `n` samples

* completely removed `reconstruct`

* renamed `maybe_reconstruct_and_link` to `to_maybe_linked_internal` and
`maybe_invlink_and_reconstruct` to `from_maybe_linked_internal`

* added impls of `from_*_internal_transform` for `ThreadSafeVarInfo`

* removed `reconstruct` from docs and from exports

* renamed `getval` to `getindex_internal` and made `dist` an optional
argument for all the transform-related methods

* updated docs + added description of how internals of transforms work

* added a bunch of illustrations for the transforms docs + dot files used to generated

* temporarily removed `VarNameVector` completely

* formatting

* Update docs/src/internals/transformations.md

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update docs/src/internals/transformations.md

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* removed refs to VectorVarInfo

* added impls of `from_internal_transform` for `ThreadSafeVarInfo`

* reverted accidental removal of old `VarInfo` constructor

* fixed incorrect `recombine` call

* removed undefined refs to `VarNameVector` stuff in `setup_varinfos`

* bump minior version because Turing breaks

* fix: was using `from_linked_internal_transform` in
`from_internal_transform` for `ThreadSafeVarInfo`

* removed `getindex_raw`

* removed redundant docstrings

* fixed tests

* fixed comparisons in tests

* try relative references for images in transformation docs

* another attempt at fixing asset-references

* fixed getindex diagrams in docs

* minor changes to comments

* remove Combinatorics as a test dep, as it's not needed for this PR

* reverted unnecessary change

* disable type-stability tests for models on older Julia versions

* removed seemingly completely unused impl of `setval!`

* Revert "temporarily removed `VarNameVector` completely"

This reverts commit 95dc8e3.

* Revert "remove Combinatorics as a test dep, as it's not needed for this PR"

This reverts commit 071bebf.

* More work on `VarNameVector` (#637)

* Update test/model.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Apply suggestions from code review

Co-authored-by: Markus Hauru <[email protected]>

* Type-stability tests are now correctly using `rand_prior_true` instead
of `rand`

* `getindex_internal` now calls `getindex` instead of `view`, as the
latter can result in type-instability since transformed variables
typically result in non-view even if input is a view

* Removed seemingly unnecessary definition of `getindex_internal`

* Fixed references to `newmetadata` which has been replaced by `replace_values`

* Made implementation of `recombine` more explicit

* Added docstrings for `untyped_varinfo` and `typed_varinfo`

* Added TODO comment about implementing `view` for `VarInfo`

* Fixed potential infinite recursion as suggested by @mhauru

* added docstring to `from_vec_trnasform_for_size

* Replaced references to `vectorize(dist, x)` with `tovec(x)`

* Fixed docstring

* Update src/extract_priors.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Bump minor version since this is a breaking change

* Apply suggestions from code review

Co-authored-by: Markus Hauru <[email protected]>

* Update src/varinfo.jl

Co-authored-by: Tor Erlend Fjelde <[email protected]>

* Apply suggestions from code review

* Apply suggestions from code review

* Update src/extract_priors.jl

Co-authored-by: Xianda Sun <[email protected]>

* Added fix for product distributions of targets with changing support + tests

* Addeed tests for product of distributions with dynamic support

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Fix typos, improve docstrings

* Use Accessors rather than Setfield

* Simplify group_by_symbol

* Add short_varinfo_name(::VectorVarInfo)

* Add tests for subset

* Export VectorVarInfo

* Tighter type bound for has_varnamevector

* Add some VectorVarName methods

* Add todo notes, remove dead code, fix a typo.

* Bug fixes and small improvements

* VarNameVector improvements

* Improve generated_quantities and its tests

* Improvement to VarNameVector

* Fix a test to work with VectorVarName

* Fix generated_quantities

* Fix type stability issues

* Various VarNameVector fixes and improvements

* Bump version number

* Improvements to generated_quantities

* Code formatting

* Code style

* Add fallback implementation of findinds for VarNameVector

* Rename VarNameVector to VarNamedVector

* More renaming of VNV. Remove unused VarNamedVector.metadata field.

* Rename FromVec to ReshapeTransform

* Progress towards having VarNamedVector as storage for SimpleVarInfo

* Fix unflatten(vnv::VarNamedVector, vals)

* More work on SimpleVarInfo{VarNamedVector}

* More tests for SimpleVarInfo{VarNamedVector}

* More tests for SimpleVarInfo{VarNamedVector}

* Respond to review feedback

* Add float_type_with_fallback(::Type{Union{}})

* Move some VNV functions to the correct file

* Fix push! for VNV

* Rename VNV.is_transformed to VNV.is_unconstrained

* Improve VNV docstring

* Add VNV inner constructor checks

* Reorganise parts of VNV code

* Documentation and small fixes for VNV

* Rename loosen_types!! and tighten_types, add docstrings and doctests

* Rename VarNameVector to VarNamedVector in docs

* Documentation and small fixes to VNV

* Fix subset(::VarNamedVector, args...) for unconstrained variables.

---------

Co-authored-by: Tor Erlend Fjelde <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Xianda Sun <[email protected]>

* Bump Bijectors dependecy

* Remove dead TODO note

* Remove old TODOs, improve VNV invlinking

* Fix from_vec_transform for 0-dim arrays

* Fix unflatten for VarInfo

* Fix some VarInfo index getters

* Change how VNV handles transformations, and other VNV stuff

* Small docs fixes

* Small fixes all over for VNV

* Add comments

* Fix some tests

* Change long string formatting to support Julia 1.6

* Small changes to ReshapeTransformation

* Revert unrelated changes to ReverseDiff extension

* Improve VarNamedVector VarInfo testing

* Fix some depwarns

* Improvements to test/simple_varinfo.jl

* Fix for unset_flag!, better docstring

* Add a comment about hasvalue/getvalue

* Add @non_differentiable calls to work around Zygote limitations

* Fix docs, workaround Zygote issue

* Remove outdated workaround

* Move has_varnamedvector(varinfo::VarInfo) to abstract_varinfo.jl

* Make copies of logp and num_produce in subset

* Rename getindex_raw to getindex_internal

* Add push!(::VarNamedVector, ::Pair)

* Improve VarNamedVector docs

* Simplify VarNamedVector constructors

* Change how VNV setindex! et al work

* More improvements to VNV setters and their tests

* Fix style issues in VNV

* Update VNV docs. Add haskey to VarInfo

* Fix VarInfo docs

* Disable a test that only works for VectorVarInfo

* Fix bug in isempty(::TypedVarInfo)

* Make some doctests platform independent

* Better implementation of haskey(::VarInfo, ::VarName)

Co-authored-by: Tor Erlend Fjelde <[email protected]>

* Improve haskey for VarInfo

* Make a VNV doctest more robust

* Remote IndexStyle for VNV

* Clean up an old comment

* Fix haskey(::VarInfo, ::VarName)

* Clarify a TODO note in varinfo.jl

* Reintroduce Int indexing to VNV

* Stop exporting any VNV stuff

* Fix docs

* Revert default VarInfo metadata type to Metadata

* Fix a few trivial issues with Metadata

* Docs bug fix

* Fix type constraint

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Markus Hauru <[email protected]>
Co-authored-by: Markus Hauru <[email protected]>
Co-authored-by: Xianda Sun <[email protected]>
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.

4 participants