def set_up_networks(seed=10, dim=2):
torch.manual_seed(seed)
base_dist_lik = StandardNormal(shape=[2])
num_layers = 5
transforms = []
for _ in range(num_layers):
transforms.append(ReversePermutation(features=2))
transforms.append(
MaskedAffineAutoregressiveTransform(features=2,
hidden_features=50,
context_features=dim,
num_blocks=1))
transform = CompositeTransform(transforms)
flow_lik = Flow(transform, base_dist_lik)
base_dist_post = StandardNormal(
shape=[dim]
) # BoxUniform(low=-2*torch.ones(2), high=2*torch.ones(2)) #StandardNormal(shape=[dim])
# base_dist_post = BoxUniform(low=-2*torch.ones(2), high=2*torch.ones(2))
num_layers = 5
transforms = []
num_off_set = 0.0001 # numerical offset since the prior is on the open space
#shift, scale = calc_scale_and_shift(-1, 1)
#print(shift)
#print(scale)
transforms.append(PointwiseAffineTransform(shift=0.5, scale=1 / 4.0))
#transforms.append(PointwiseAffineTransform(shift=shift, scale=scale))
transforms.append(InvSigmoid.InvSigmoid()) # this should be inv sigmoide!
for _ in range(num_layers):
transforms.append(ReversePermutation(features=dim))
transforms.append(
MaskedAffineAutoregressiveTransform(features=dim,
hidden_features=50,
context_features=2,
num_blocks=1))
transform = CompositeTransform(transforms)
flow_post = Flow(transform, base_dist_post)
return flow_lik, flow_post
def make_model(num_layers, num_features, num_hidden_features, device):
#context_encoder = nn.Sequential(
# nn.Linear(num_features, 4*num_features),
# nn.ReLU(),
# nn.Linear(4*num_features, 4*num_features),
# nn.ReLU(),
# nn.Linear(4*num_features, 2*num_features)
# )
#context_encoder = nn.Sequential(
# nn.Linear(num_features, 2*num_features),
# nn.ReLU(),
# nn.Linear(2*num_features, 2*num_features),
# nn.ReLU(),
# nn.Linear(2*num_features, 2*num_features)
# )
context_encoder = nn.Sequential(nn.Linear(num_features, 256), nn.ReLU(),
nn.Linear(256, 256), nn.ReLU(),
nn.Linear(256, 2 * num_features))
base_dist = ConditionalDiagonalNormal(shape=[num_features],
context_encoder=context_encoder)
#base_dist = StandardNormal(shape=[num_features])
#base_dist = DiagonalNormal(shape=[num_features])
transforms = []
for _ in range(num_layers):
transforms.append(ReversePermutation(features=num_features))
#UMNN
#transforms.append(MaskedUMNNAutoregressiveTransform(features=num_features,
# hidden_features=num_hidden_features))
#Conditional MAA
#transforms.append(MaskedAffineAutoregressiveTransform(features=num_features,
# hidden_features=num_hidden_features,
# context_features=num_features))
#Conditional UMNN
transforms.append(
MaskedUMNNAutoregressiveTransform(
features=num_features,
hidden_features=num_hidden_features,
context_features=num_features))
transform = CompositeTransform(transforms)
#Uncomment the below if float64
#flow = Flow(transform, base_dist).double().to(device)
#Uncomment the below if float32
flow = Flow(transform, base_dist).to(device)
optimizer = optim.Adam(flow.parameters())
return (flow, optimizer)
def conditional_distribution(self,
context: Union[np.array, Tensor] = None,
data_normalization=True,
context_one_hot_encoding=True) -> Flow:
_, context, _ = self._preprocess_data(None, context,
data_normalization,
context_one_hot_encoding, False)
transforms_list = torch.nn.ModuleList()
if self.inputs_normalization:
# Inverse normalisation
transforms_list.append(
PointwiseAffineTransform(shift=-self.inputs_mean /
self.inputs_std,
scale=1 / self.inputs_std))
# Forward pass, to init conditional parameters
with torch.no_grad():
_ = super().log_prob(torch.zeros(len(context), 1), context)
transforms_list.extend(deepcopy(self._transform._transforms))
cond_dist = Flow(CompositeTransform(transforms_list),
self._distribution)
return cond_dist
def make_model(num_layers,num_features,num_hidden_features,device):
base_dist = StandardNormal(shape=[num_features])
#base_dist = DiagonalNormal(shape=[num_features])
transforms = []
for _ in range(num_layers):
transforms.append(ReversePermutation(features=num_features))
transforms.append(MaskedUMNNAutoregressiveTransform(features=num_features,
hidden_features=num_hidden_features))
transform = CompositeTransform(transforms)
#Uncomment the below if float64
#flow = Flow(transform, base_dist).double().to(device)
flow = Flow(transform, base_dist).to(device)
optimizer = optim.Adam(flow.parameters())
return (flow,optimizer)
def set_up_networks(seed=10, dim=4):
torch.manual_seed(seed)
base_dist_lik = StandardNormal(shape=[9])
num_layers = 5
transforms = []
for _ in range(num_layers): # TODO add inv sigmoide fnunc
transforms.append(ReversePermutation(features=9))
transforms.append(
MaskedAffineAutoregressiveTransform(features=9,
hidden_features=10,
context_features=dim,
num_blocks=2))
transform = CompositeTransform(transforms)
flow_lik = Flow(transform, base_dist_lik)
base_dist_post = StandardNormal(shape=[dim])
num_layers = 4
transforms = []
num_off_set = 0.0001 # numerical offset since the prior is on the open space
shift, scale = calc_scale_and_shift(-5, 2)
#transforms.append(PointwiseAffineTransform(shift=5 / 7.0, scale=1 / 7.0))
transforms.append(PointwiseAffineTransform(shift=shift, scale=scale))
transforms.append(InvSigmoid.InvSigmoid()) # this should be inv sigmoide!
for _ in range(num_layers):
transforms.append(ReversePermutation(features=dim))
transforms.append(
MaskedAffineAutoregressiveTransform(features=dim,
hidden_features=10,
context_features=9,
num_blocks=2))
transform = CompositeTransform(transforms)
flow_post = Flow(transform, base_dist_post)
return flow_lik, flow_post
def toy_flow(args, n_blocks, input_dim, hidden_dim, num_layers):
base_dist = StandardNormal(shape=[input_dim])
transforms = []
for _ in range(num_layers):
transforms.append(ReversePermutation(features=input_dim))
transforms.append(
MaskedAffineAutoregressiveTransform(features=input_dim,
hidden_features=hidden_dim))
transform = CompositeTransform(transforms)
flow = Flow(transform, base_dist)
return flow
def set_up_networks(seed=10, dim=2):
torch.manual_seed(seed)
base_dist = StandardNormal(shape=[10])
num_layers = 4
transforms = []
for _ in range(num_layers):
transforms.append(ReversePermutation(features=10))
transforms.append(MaskedAffineAutoregressiveTransform(features=10,
hidden_features=40,
context_features=dim,
num_blocks=1))
transform = CompositeTransform(transforms)
flow_lik = Flow(transform, base_dist)
base_dist_post = StandardNormal(shape=[dim])
num_layers = 4
transforms = []
for _ in range(num_layers):
transforms.append(ReversePermutation(features=dim))
transforms.append(MaskedAffineAutoregressiveTransform(features=dim,
hidden_features=40,
context_features=5, # it sort of makes more sense that
# context_features is the nbr of features returned by the
# summary net
num_blocks=1))
transform = CompositeTransform(transforms)
# here we add the SummaryNet as an embedded network to the flow model for the posterior
flow_post = Flow(transform, base_dist_post, embedding_net=SummaryNet())
return flow_lik, flow_post
def set_up_networks(seed=10, dim=2):
torch.manual_seed(seed)
base_dist_lik = StandardNormal(shape=[5])
num_layers = 5
transforms = []
for _ in range(num_layers):
transforms.append(ReversePermutation(features=5))
transforms.append(
MaskedAffineAutoregressiveTransform(features=5,
hidden_features=10,
context_features=dim,
num_blocks=2))
transform = CompositeTransform(transforms)
flow_lik = Flow(transform, base_dist_lik)
base_dist_post = StandardNormal(shape=[dim])
num_layers = 4
transforms = []
for _ in range(num_layers):
transforms.append(ReversePermutation(features=dim))
transforms.append(
MaskedAffineAutoregressiveTransform(features=dim,
hidden_features=10,
context_features=5,
num_blocks=2))
transform = CompositeTransform(transforms)
flow_post = Flow(transform, base_dist_post)
return flow_lik, flow_post
def __init__(self,
context_size: int,
inputs_size: int = 1,
transform_classes: Tuple[Type] = 2 *
(ContextualInvertableRadialTransform, ) +
(ContextualPointwiseAffineTransform, ),
hidden_units: Tuple[int] = (16, ),
base_distribution: Distribution = None,
n_epochs: int = 1000,
lr: float = 0.001,
weight_decay: float = 0.0,
batch_size: int = None,
input_noise_std: float = 0.05,
context_noise_std: float = 0.1,
cat_context: np.array = None,
context_normalization=True,
inputs_normalization=True,
inputs_noise_nonlinearity=torch.nn.Identity(),
device='cpu',
**kwargs):
"""
PyTorch implementation of Noise Regularization for Conditional Density Estimation. Also works unconditionally
(https://github.com/freelunchtheorem/Conditional_Density_Estimation)
[Rothfuss et al, 2020; https://arxiv.org/pdf/1907.08982.pdf]
Args:
context_size: Dimensionality of global_context
transform_classes: Contextual transformations list
hidden_units: Tuple of hidden sizes for global_context embedding network
n_epochs: Number of training epochs
lr: Learning rate
weight_decay: Weight decay (not applied to bias)
input_noise_std: Noise regularisation for input
context_noise_std: Noise regularisation for global_context
base_distribution: Base distribution of normalising flow
context_normalization: Mean-std normalisation of global_context
inputs_normalization: Mean-std normalisation of context_vars
device: cpu / cuda
"""
# Constructing composite transformation & Initialisation of global_context embedding network
self.hidden_units = hidden_units
if context_size > 0:
transform = ContextualCompositeTransform([
transform_cls(inputs_size=inputs_size)
for transform_cls in transform_classes
])
embedding_net = ContextEmbedding(transform,
input_units=context_size,
hidden_units=hidden_units)
else:
transform = CompositeTransform([
transform_cls(inputs_size=inputs_size, conditional=False)
for transform_cls in transform_classes
])
embedding_net = None
if base_distribution is None:
base_distribution = StandardNormal(shape=[inputs_size])
assert base_distribution._shape[0] == inputs_size
Flow.__init__(self, transform, base_distribution, embedding_net)
# Training params
self.lr = lr
self.weight_decay = weight_decay if weight_decay is not None else 0.0
self._init_optimizer(lr, weight_decay)
self.n_epochs = n_epochs
self.device = device
self.to(self.device)
self.batch_size = batch_size
# Regularisation
self.input_noise_std = input_noise_std
self.context_noise_std = context_noise_std
if cat_context is None or len(cat_context) == 0:
self.cat_context = []
self.cont_context = np.arange(0, context_size)
else:
self.cat_context = cat_context
self.cont_context = np.array(
[i for i in range(context_size) if i not in cat_context])
self.cont_context = [] if len(
self.cont_context) == 0 else self.cont_context
# Inputs / Context one-hot encoders
self.context_enc = OneHotEncoder(drop='if_binary', sparse=False)
self.inputs_enc = OneHotEncoder(drop='if_binary', sparse=False)
# Normalisation
self.context_normalization = context_normalization
self.inputs_normalization = inputs_normalization
self.inputs_noise_nonlinearity = inputs_noise_nonlinearity
self.inputs_mean, self.inputs_std = None, None
self.context_mean, self.context_std = None, None
self.inputs_size = inputs_size
self.context_size = context_size
def get_neural_posterior(model, parameter_dim, observation_dim, simulator):
# Everything is a flow because we need to normalize parameters based on prior.
mean, std = simulator.normalization_parameters
normalizing_transform = AffineTransform(shift=-mean / std, scale=1 / std)
if model == "mdn":
hidden_features = 50
neural_posterior = MultivariateGaussianMDN(
features=parameter_dim,
context_features=observation_dim,
hidden_features=hidden_features,
hidden_net=nn.Sequential(
nn.Linear(observation_dim, hidden_features),
nn.ReLU(),
nn.Dropout(p=0.0),
nn.Linear(hidden_features, hidden_features),
nn.ReLU(),
nn.Linear(hidden_features, hidden_features),
nn.ReLU(),
),
num_components=20,
custom_initialization=True,
)
elif model == "made":
num_mixture_components = 5
transform = normalizing_transform
distribution = MADEMoG(
features=parameter_dim,
hidden_features=50,
context_features=observation_dim,
num_blocks=2,
num_mixture_components=num_mixture_components,
use_residual_blocks=True,
random_mask=False,
activation=F.relu,
dropout_probability=0.0,
use_batch_norm=False,
custom_initialization=True,
)
neural_posterior = Flow(transform, distribution)
elif model == "maf":
transform = CompositeTransform(
[
CompositeTransform(
[
transforms.MaskedAffineAutoregressiveTransform(
features=parameter_dim,
hidden_features=50,
context_features=observation_dim,
num_blocks=2,
use_residual_blocks=False,
random_mask=False,
activation=F.tanh,
dropout_probability=0.0,
use_batch_norm=True,
),
transforms.RandomPermutation(features=parameter_dim),
]
)
for _ in range(5)
]
)
transform = CompositeTransform([normalizing_transform, transform,])
distribution = StandardNormal((parameter_dim,))
neural_posterior = Flow(transform, distribution)
elif model == "nsf":
transform = CompositeTransform(
[
CompositeTransform(
[
transforms.PiecewiseRationalQuadraticCouplingTransform(
mask=create_alternating_binary_mask(
features=parameter_dim, even=(i % 2 == 0)
),
transform_net_create_fn=lambda in_features, out_features: nets.ResidualNet(
in_features=in_features,
out_features=out_features,
hidden_features=50,
context_features=observation_dim,
num_blocks=2,
activation=F.relu,
dropout_probability=0.0,
use_batch_norm=False,
),
num_bins=10,
tails="linear",
tail_bound=3.0,
apply_unconditional_transform=False,
),
transforms.LULinear(parameter_dim, identity_init=True),
]
)
for i in range(5)
]
)
distribution = StandardNormal((parameter_dim,))
neural_posterior = Flow(transform, distribution)
else:
raise ValueError
return neural_posterior
def get_neural_likelihood(model, parameter_dim, observation_dim):
if model == "mdn":
hidden_features = 50
neural_likelihood = MultivariateGaussianMDN(
features=observation_dim,
context_features=parameter_dim,
hidden_features=hidden_features,
hidden_net=nn.Sequential(
nn.Linear(parameter_dim, hidden_features),
nn.BatchNorm1d(hidden_features),
nn.ReLU(),
nn.Dropout(p=0.0),
nn.Linear(hidden_features, hidden_features),
nn.BatchNorm1d(hidden_features),
nn.ReLU(),
nn.Linear(hidden_features, hidden_features),
nn.BatchNorm1d(hidden_features),
nn.ReLU(),
),
num_components=20,
custom_initialization=True,
)
elif model == "made":
neural_likelihood = MixtureOfGaussiansMADE(
features=observation_dim,
hidden_features=50,
context_features=parameter_dim,
num_blocks=4,
num_mixture_components=10,
use_residual_blocks=True,
random_mask=False,
activation=F.relu,
use_batch_norm=True,
dropout_probability=0.0,
custom_initialization=True,
)
elif model == "maf":
transform = CompositeTransform(
[
CompositeTransform(
[
MaskedAffineAutoregressiveTransform(
features=observation_dim,
hidden_features=50,
context_features=parameter_dim,
num_blocks=2,
use_residual_blocks=False,
random_mask=False,
activation=F.tanh,
dropout_probability=0.0,
use_batch_norm=True,
),
RandomPermutation(features=observation_dim),
]
)
for _ in range(5)
]
)
distribution = StandardNormal((observation_dim,))
neural_likelihood = Flow(transform, distribution)
elif model == "nsf":
transform = CompositeTransform(
[
CompositeTransform(
[
PiecewiseRationalQuadraticCouplingTransform(
mask=create_alternating_binary_mask(
features=observation_dim, even=(i % 2 == 0)
),
transform_net_create_fn=lambda in_features, out_features: nets.ResidualNet(
in_features=in_features,
out_features=out_features,
hidden_features=50,
context_features=parameter_dim,
num_blocks=2,
activation=F.relu,
dropout_probability=0.0,
use_batch_norm=False,
),
num_bins=10,
tails="linear",
tail_bound=3.0,
apply_unconditional_transform=False,
),
LULinear(observation_dim, identity_init=True),
]
)
for i in range(5)
]
)
distribution = StandardNormal((observation_dim,))
neural_likelihood = Flow(transform, distribution)
else:
raise ValueError
return neural_likelihood
#base_dist = DiagonalNormal(shape=[3])
transforms = []
for _ in range(num_layers):
transforms.append(ReversePermutation(features=num_features))
# transforms.append(MaskedAffineAutoregressiveTransform(features=num_features,
# hidden_features=100))
transforms.append(
MaskedAffineAutoregressiveTransform(features=num_features,
hidden_features=80,
context_features=num_features))
#transforms.append(MaskedUMNNAutoregressiveTransform(features=num_features,
# hidden_features=4))
transform = CompositeTransform(transforms)
flow = Flow(transform, base_dist).to(device)
optimizer = optim.Adam(flow.parameters())
print("number of params: ", sum(p.numel() for p in flow.parameters()))
def plot_histo_1D(real_vals,
gen_vals,
label_real="Physics Data",
label_gen="NFlow Model",
col2="blue",
title="Physics vs NFlow Models",
saveloc=None):
fig, axes = plt.subplots(1, 4, figsize=(4 * 5, 5))
for INDEX, ax in zip((0, 1, 2, 3), axes):
_, bins, _ = ax.hist(real_vals[:, INDEX],
bins=100,
base_dist = StandardNormal(shape=[num_features])
#base_dist = DiagonalNormal(shape=[3])
transforms = []
for _ in range(num_layers):
transforms.append(ReversePermutation(features=num_features))
transforms.append(
MaskedUMNNAutoregressiveTransform(features=num_features,
hidden_features=20))
#context_features=len(in_columns)))
#transforms.append(MaskedAffineAutoregressiveTransform(features=num_features,
# hidden_features=100))
transform = CompositeTransform(transforms)
flow = Flow(transform, base_dist).to(device)
#model = MaskedUMNNAutoregressiveTransform()
flow.load_state_dict(torch.load("models/trainedmodel_final_16.pkl"))
flow.eval()
zs = []
start = datetime.now()
start_time = start.strftime("%H:%M:%S")
print("Start Time =", start_time)
max_range = 200
for i in range(1, max_range + 1):
print("On set {}".format(i))
z0 = flow.sample(200).cpu().detach().numpy()
base_dist = StandardNormal(shape=[num_features])
#base_dist = DiagonalNormal(shape=[3])
transforms = []
for _ in range(num_layers):
transforms.append(ReversePermutation(features=num_features))
transforms.append(
MaskedUMNNAutoregressiveTransform(features=num_features,
hidden_features=20))
#context_features=len(in_columns)))
#transforms.append(MaskedAffineAutoregressiveTransform(features=num_features,
# hidden_features=100))
transform = CompositeTransform(transforms)
flow = Flow(transform, base_dist).to(device)
#model = MaskedUMNNAutoregressiveTransform()
flow.load_state_dict(torch.load("models/trainedmodel_799_100_-15.19.pkl"))
#flow.load_state_dict(torch.load("trainedmodel_99_100_-2.76.pkl"))
#flow.double().eval()
flow.eval()
zs = []
start = datetime.now()
start_time = start.strftime("%H:%M:%S")
print("Start Time =", start_time)
max_range = 20
sample_size = 200
def set_up_networks(lower_post_limits, upper_post_limits, dim_post=12, dim_summary_stat=19, seed=10):
torch.manual_seed(seed)
base_dist_lik = StandardNormal(shape=[dim_summary_stat])
num_layers = 4
transforms = []
for _ in range(num_layers): # TODO add inv sigmoide fnunc
transforms.append(ReversePermutation(features=dim_summary_stat))
transforms.append(MaskedAffineAutoregressiveTransform(features=dim_summary_stat,
hidden_features=40,
context_features=dim_post,
num_blocks=2))
transform = CompositeTransform(transforms)
flow_lik = Flow(transform, base_dist_lik)
base_dist_post = StandardNormal(shape=[dim_post])
#base_dist_post = ConditionalDiagonalNormal(shape=[dim_post], context_encoder=nn.Linear(dim_summary_stat, 2*dim_post))
#base_dist_post = UniformContext.UniformContext(low=lower_post_limits, high=upper_post_limits, shape=[dim_post])
num_layers = 4
transforms = []
# def post model in prior snplace
shift_vec = torch.zeros(dim_post)
scale_vec = torch.zeros(dim_post)
num_off_set = 0.001 # numerical offset since the prior is on the open space
print("set priors")
print(upper_post_limits)
print(lower_post_limits)
for i in range(dim_post):
shift_tmp, scale_tmp = calc_scale_and_shift(lower_post_limits[i], upper_post_limits[i])
shift_vec[i] = shift_tmp
scale_vec[i] = scale_tmp
#print(shift_vec)
#print(scale_vec)
print(shift_vec)
print(scale_vec)
# last transformation
transforms.append(PointwiseAffineTransform(shift=shift_vec, scale=scale_vec))
transforms.append(InvSigmoid.InvSigmoid()) # this should be inv sigmoide!
for _ in range(num_layers):
transforms.append(ReversePermutation(features=dim_post))
transforms.append(MaskedAffineAutoregressiveTransform(features=dim_post,
hidden_features=50,
context_features=dim_summary_stat,
num_blocks=2))
transform = CompositeTransform(transforms)
flow_post = Flow(transform, base_dist_post)
return flow_lik, flow_post
