Skip to content

FlowDensity

Normalizing flow networks for density estimation.

Overview

FlowDensity is the internal nn.Module that wraps various normalizing flow architectures for use in posterior estimation. It is used internally by the Flow estimator.

Supported Flow Types

zuko_* types are available by default. Types backed by sbi/nflows require the optional sbi extra: pip install falcon-sbi[sbi].

Type Library Description
nsf sbi/nflows Neural Spline Flow (recommended when sbi is installed)
maf sbi/nflows Masked Autoregressive Flow
made sbi Masked Autoencoder for Distribution Estimation
maf_rqs sbi MAF with Rational Quadratic Splines
zuko_nice Zuko NICE architecture (default — no extra required)
zuko_maf Zuko Masked Autoregressive Flow (Zuko)
zuko_nsf Zuko Neural Spline Flow (Zuko)
zuko_ncsf Zuko Neural Circular Spline Flow
zuko_sospf Zuko Sum-of-Squares Polynomial Flow
zuko_naf Zuko Neural Autoregressive Flow
zuko_unaf Zuko Unconstrained NAF
zuko_gf Zuko Glow Flow
zuko_bpf Zuko Bernstein Polynomial Flow

Class Reference

FlowDensity 

FlowDensity(theta, s, theta_norm=False, norm_momentum=0.003, net_type='nsf', use_log_update=False, adaptive_momentum=False)

Bases: Module

Normalizing flow network with optional parameter normalization.

Source code in src/falcon/estimators/flow_density.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def __init__(
    self,
    theta,
    s,
    theta_norm=False,
    norm_momentum=3e-3,
    net_type="nsf",
    use_log_update=False,
    adaptive_momentum=False,
):
    super().__init__()
    self.theta_norm = (
        LazyOnlineNorm(
            momentum=norm_momentum,
            use_log_update=use_log_update,
            adaptive_momentum=adaptive_momentum,
            log_prefix = "FlowDensity:OnlineNorm",
        )
        if theta_norm
        else None
    )

    net_builders = _get_net_builders()
    builder = net_builders.get(net_type)
    if builder is None:
        raise ValueError(f"Unknown net_type: {net_type}. Available: {list(net_builders.keys())}")
    self.net = builder(theta.float(), s.float(), z_score_x=None, z_score_y=None)

    if self.theta_norm is not None:
        self.theta_norm(theta)  # Initialize normalization stats
    self.scale = 0.2

loss 

loss(theta, s)

Compute negative log-likelihood loss.

Source code in src/falcon/estimators/flow_density.py 
73
74
75
76
77
78
79
80
81
82
def loss(self, theta, s):
    """Compute negative log-likelihood loss."""
    if self.theta_norm is not None:
        theta = self.theta_norm(theta)
    theta = theta.float() * self.scale
    loss = self.net.loss(theta, condition=s.float())
    loss = loss - np.log(self.scale) * theta.shape[-1]
    if self.theta_norm is not None:
        loss = loss + torch.log(self.theta_norm.volume())
    return loss

sample 

sample(num_samples, s)

Sample from the flow.

Source code in src/falcon/estimators/flow_density.py 
84
85
86
87
88
89
90
def sample(self, num_samples, s):
    """Sample from the flow."""
    samples = self.net.sample((num_samples,), condition=s).detach()
    samples = samples / self.scale
    if self.theta_norm is not None:
        samples = self.theta_norm.inverse(samples).detach()
    return samples

log_prob 

log_prob(theta, s)

Compute log probability.

Source code in src/falcon/estimators/flow_density.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def log_prob(self, theta, s):
    """Compute log probability."""
    if self.theta_norm is not None:
        theta = self.theta_norm(theta).detach()
    theta = theta * self.scale
    log_prob = self.net.log_prob(theta.float(), condition=s.float())
    log_prob = log_prob + np.log(self.scale) * theta.shape[-1]
    if self.theta_norm is not None:
        log_prob = log_prob - torch.log(self.theta_norm.volume().detach())
    return log_prob