Skip to content

Commit ebcd68e

Browse files
anastasiakrouglovaNastya KrouglovaNastya KrouglovaNastya KrouglovaNastya Krouglova
authored
feat: logit transform for zuko training in unconstrained space (#1485)
* logit transform changes * add new ZScoreTypes * add new ZScoreTypes * resolving bug z_scoring last year * resolving bug z_scoring last year * revert z_score parser * adjusted logit structure in build_zuko_flow * resolve pyright error * revert flow as a test * add x_dist variable * add logit to sbiutils_test.py * add logit if statement * add logit if statement * add logit if statement * add logit if statement * remove logit if statement * resolve pyright issues * cover logit in tests * cover tests for logit in flow.py * cover tests for CNF * adding faq for logit transformation * adding faq for logit transformation * stash changes * feedback guy adjustments * add documentation if statements * resolve comment 1 and 2 of Jan * ruff linted push * cleanup density_estimator_test.py * cleanup density_estimator_test.py and ruff check * adjusted docstrings * add tests convergence unconstrained space * adjust faq * add test snle * adjust linear gaussian and estimate c2st * adjust documentation * Update sbi/neural_nets/net_builders/flow.py Co-authored-by: Jan <janfb@users.noreply.github.com> * Update sbi/utils/sbiutils.py Co-authored-by: Jan <janfb@users.noreply.github.com> * Update sbi/utils/sbiutils.py Co-authored-by: Jan <janfb@users.noreply.github.com> * add literal import to sbi utils * adjust literals and add get_transform_to_unconstrained * adjust literals and format * add new line for ruff * stying ruff * stying ruff * fix flow builder z-score defaults. * refactor zuko flow build functions * re-use y-embedding helper function. * fix typing * small fixes. * fix unconstrained nle test * refactor z-score-parser test --------- Co-authored-by: Nastya Krouglova <nastyakrouglova@webauth217.uni-tuebingen.de> Co-authored-by: Nastya Krouglova <nastyakrouglova@webauth494.uni-tuebingen.de> Co-authored-by: Nastya Krouglova <nastyakrouglova@webauth446.uni-tuebingen.de> Co-authored-by: Nastya Krouglova <nastyakrouglova@webauth233.uni-tuebingen.de> Co-authored-by: Nastya Krouglova <nastyakrouglova@Nastyas-MacBook-Pro.local> Co-authored-by: Jan <janfb@users.noreply.github.com> Co-authored-by: Jan <jan.boelts@mailbox.org>
1 parent 65ae720 commit ebcd68e

File tree

7 files changed

+629
-145
lines changed

7 files changed

+629
-145
lines changed
Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# Using the logit transformation
2+
If you've ruled out simulator issues, you can try
3+
training your density or ratio estimator in an unbounded space
4+
using a logit transformation:
5+
6+
- **For NPE**: The transformation maps bounded parameters θ
7+
to unbounded space before training, then applies the inverse (sigmoid)
8+
after training to ensure posterior samples stay within prior bounds.
9+
10+
- **For NLE/NRE**: The transformation would need to map bounded
11+
data x to unbounded space, which requires estimating data bounds
12+
from simulations (more complex).
13+
14+
To enable this for NPE:
15+
16+
```python
17+
density_estimator_build_fun = posterior_nn(
18+
model="zuko_nsf",
19+
hidden_features=60,
20+
num_transforms=3,
21+
z_score_theta="transform_to_unconstrained" # Transforms parameters to unconstrained space
22+
x_dist=prior # For NPE, this specifies bounds for parameters (internally called 'x')
23+
)
24+
inference = NPE(prior, density_estimator=density_estimator_build_fun)
25+
```
26+
27+
This ensures that your density estimator operates in a
28+
transformed space where it respects prior bounds,
29+
improving the efficiency of rejection sampling.
30+
31+
Note: The `x_dist=prior` might seem confusing - internally,
32+
sbi uses generic `x,y` notation where for NPE, `x` represents
33+
parameters (θ) and `y` represents data.
34+
This is why we pass the prior as `x_dist`.
35+
36+
Important:
37+
38+
- This transformation is currently only supported for zuko density estimators.
39+
- For **NLE/NRE**, setting up this transformation is more
40+
complex as it requires estimating bounds for the simulated data
41+
rather than using prior bounds.

0 commit comments

Comments
 (0)