def init_path(arch_string, init_type, init_params):
"""Deduces initialization file path from initialization type and parameters.
:param arch_string: Architecture string of normalizing flow.
:type arch_string: str
:param init_type: Initialization type \in ['iso_gauss']
:type init_type: str
:param init_param: init_type dependent parameters for initialization (more deets)
:type dict:
:return: Initialization save path.
:rtype: str
"""
if type(arch_string) is not str:
raise TypeError(
format_type_err_msg("epi.util.init_path", "arch_string",
arch_string, str))
if type(init_type) is not str:
raise TypeError(
format_type_err_msg("epi.util.init_path", "init_type", init_type,
str))
path = "./data/" + arch_string + "/"
if init_type == "iso_gauss":
if "loc" in init_params:
loc = init_params["loc"]
else:
raise ValueError("'loc' field not in init_param for %s." %
init_type)
if "scale" in init_params:
scale = init_params["scale"]
else:
raise ValueError("'scale' field not in init_param for %s." %
init_type)
path += init_type + "_loc=%.2E_scale=%.2E/" % (loc, scale)
elif init_type == "gaussian":
if "mu" in init_params:
mu = np_column_vec(init_params["mu"])[:, 0]
else:
raise ValueError("'mu' field not in init_param for %s." %
init_type)
if "Sigma" in init_params:
Sigma = init_params["Sigma"]
else:
raise ValueError("'Sigma' field not in init_param for %s." %
init_type)
D = mu.shape[0]
mu_str = array_str(mu)
Sigma_str = array_str(Sigma[np.triu_indices(D, 0)])
path += init_type + "_mu=%s_Sigma=%s/" % (mu_str, Sigma_str)
if not os.path.exists(path):
os.makedirs(path)
return path
def __init__(self, lb, ub):
"""Constructor method."""
super().__init__(forward_min_event_ndims=1, inverse_min_event_ndims=1)
# Check types.
if type(lb) not in [list, np.ndarray]:
raise TypeError(format_type_err_msg(self, "lb", lb, np.ndarray))
if type(ub) not in [list, np.ndarray]:
raise TypeError(format_type_err_msg(self, "ub", ub, np.ndarray))
# Handle list input.
if type(lb) is list:
lb = np.array(lb)
if type(ub) is list:
ub = np.array(ub)
# Make sure we have 1-D np vec
self.lb = np_column_vec(lb)[:, 0]
self.ub = np_column_vec(ub)[:, 0]
if self.lb.shape[0] != self.ub.shape[0]:
raise ValueError("lb and ub have different lengths.")
self.D = self.lb.shape[0]
for lb_i, ub_i in zip(self.lb, self.ub):
if lb_i >= ub_i:
raise ValueError("Lower bound %.2E > upper bound %.2E." %
(lb_i, ub_i))
sigmoid_flg, softplus_flg = self.D * [0], self.D * [0]
sigmoid_m, sigmoid_c = self.D * [1.0], self.D * [0.0]
softplus_m, softplus_c = self.D * [1.0], self.D * [0.0]
for i in range(self.D):
lb_i, ub_i = self.lb[i], self.ub[i]
has_lb = not np.isneginf(lb_i)
has_ub = not np.isposinf(ub_i)
if has_lb and has_ub:
sigmoid_flg[i] = 1
sigmoid_m[i] = ub_i - lb_i
sigmoid_c[i] = lb_i
elif has_lb:
softplus_flg[i] = 1
softplus_m[i] = 1.0
softplus_c[i] = lb_i
elif has_ub:
softplus_flg[i] = 1
softplus_m[i] = -1.0
softplus_c[i] = ub_i
self.sigmoid_flg = tf.constant(sigmoid_flg, dtype=DTYPE)
self.softplus_flg = tf.constant(softplus_flg, dtype=DTYPE)
self.sigmoid_m = tf.constant(sigmoid_m, dtype=DTYPE)
self.sigmoid_c = tf.constant(sigmoid_c, dtype=DTYPE)
self.softplus_m = tf.constant(softplus_m, dtype=DTYPE)
self.softplus_c = tf.constant(softplus_c, dtype=DTYPE)
def _set_parameters(self, parameters):
if type(parameters) is not list:
raise TypeError(
format_type_err_msg(self, parameters, "parameters", list))
for parameter in parameters:
if not parameter.__class__.__name__ == "Parameter":
raise TypeError(
format_type_err_msg(self, "parameter", parameter,
Parameter))
if not self.parameter_check(parameters, verbose=True):
raise ValueError("Invalid parameter list.")
self.parameters = parameters
self.D = sum([param.D for param in parameters])
def _set_parameters(self, parameters):
if parameters is None:
parameters = [Parameter("z%d" % (i + 1), 1) for i in range(self.D)]
self.parameters = parameters
elif type(parameters) is not list:
raise TypeError(
format_type_err_msg(self, "parameters", parameters, list))
else:
for parameter in parameters:
if not parameter.__class__.__name__ == "Parameter":
raise TypeError(
format_type_err_msg(self, "parameter", parameter,
Parameter))
self.parameters = parameters
def _set_bounds(self, bounds):
if bounds is not None:
_type = type(bounds)
if _type in [list, tuple]:
len_bounds = len(bounds)
if _type is list:
bounds = tuple(bounds)
else:
raise TypeError(
"NormalizingFlow argument bounds must be tuple or list not %s."
% _type.__name__)
if len_bounds != 2:
raise ValueError(
"NormalizingFlow bounds arg must be length 2.")
for i, bound in enumerate(bounds):
if type(bound) is not np.ndarray:
raise TypeError(
format_type_err_msg(self, "bounds[%d]" % i, bound,
np.ndarray))
self.lb, self.ub = bounds[0], bounds[1]
else:
self.lb, self.ub = None, None
def _set_D(self, D):
if type(D) is not int:
raise TypeError(format_type_err_msg(self, "D", D, int))
elif D < 2:
raise ValueError("NormalizingFlow D %d must be greater than 0." %
D)
self.D = D
def _set_num_units(self, num_units):
if type(num_units) is not int:
raise TypeError(
format_type_err_msg(self, "num_units", num_units, int))
elif num_units < 1:
raise ValueError(
"NormalizingFlow num_units %d must be greater than 0." %
num_units)
self.num_units = num_units
def _set_num_bins(self, num_bins):
if type(num_bins) is not int:
raise TypeError(
format_type_err_msg(self, "num_bins", num_bins, int))
elif num_bins < 2:
raise ValueError(
"NormalizingFlow spline num_bins %d must be greater than 1." %
num_units)
self.num_bins = num_bins
def _set_elemwise_fn(self, elemwise_fn):
elemwise_fns = ["affine", "spline"]
if type(elemwise_fn) is not str:
raise TypeError(
format_type_err_msg(self, "elemwise_fn", elemwise_fn, str))
if elemwise_fn not in elemwise_fns:
raise ValueError(
'NormalizingFlow elemwise_fn must be "affine" or "spline"')
self.elemwise_fn = elemwise_fn
def _set_bn_momentum(self, bn_momentum):
if type(bn_momentum) is not float:
raise TypeError(
format_type_err_msg(self, "bn_momentum", bn_momentum, float))
self.bn_momentum = bn_momentum
bijectors = self.trans_dist.bijector.bijectors
for bijector in bijectors:
if type(bijector).__name__ == "BatchNormalization":
bijector.batchnorm.momentum = bn_momentum
return None
def _set_arch_type(self, arch_type): # Make this noninherited?
arch_types = ["coupling", "autoregressive"]
if type(arch_type) is not str:
raise TypeError(
format_type_err_msg(self, "arch_type", arch_type, str))
if arch_type not in arch_types:
raise ValueError(
'NormalizingFlow arch_type must be "coupling" or "autoregressive"'
)
self.arch_type = arch_type
def _set_bounds(self, lb, ub):
if lb is None:
lb = np.NINF * np.ones(self.D)
elif isinstance(lb, REAL_NUMERIC_TYPES):
lb = np.array([lb])
if ub is None:
ub = np.PINF * np.ones(self.D)
elif isinstance(ub, REAL_NUMERIC_TYPES):
ub = np.array([ub])
if type(lb) is not np.ndarray:
raise TypeError(format_type_err_msg(self, "lb", lb, np.ndarray))
if type(ub) is not np.ndarray:
raise TypeError(format_type_err_msg(self, "ub", ub, np.ndarray))
lb_shape = lb.shape
if len(lb_shape) != 1:
raise ValueError("Lower bound lb must be vector.")
if lb_shape[0] != self.D:
raise ValueError("Lower bound lb does not have dimension D = %d." %
self.D)
ub_shape = ub.shape
if len(ub_shape) != 1:
raise ValueError("Upper bound ub must be vector.")
if ub_shape[0] != self.D:
raise ValueError("Upper bound ub does not have dimension D = %d." %
self.D)
for i in range(self.D):
if lb[i] > ub[i]:
raise ValueError(
"Parameter %s lower bound is greater than upper bound." %
self.name)
elif lb[i] == ub[i]:
raise ValueError(
"Parameter %s lower bound is equal to upper bound." %
self.name)
self.lb = lb
self.ub = ub
def check_bound_param(bounds, param_name):
if type(bounds) not in [list, tuple]:
raise TypeError(
format_type_err_msg("sample_aug_lag_hps", param_name, bounds,
list))
if len(bounds) != 2:
raise ValueError("Bound should be length 2.")
if bounds[1] < bounds[0]:
raise ValueError(
"Bounds are not ordered correctly: bounds[1] < bounds[0].")
return None
def gaussian_backward_mapping(mu, Sigma):
"""Calculates natural parameter of multivaraite gaussian from mean and cov.
:param mu: Mean of gaussian
:type mu: np.ndarray
:param Sigma: Covariance of gaussian.
:type Sigma: np.ndarray
:return: Natural parameter of gaussian.
:rtype: np.ndarray
"""
if type(mu) is not np.ndarray:
raise TypeError(
format_type_err_msg("epi.util.gaussian_backward_mapping", "mu", mu,
np.ndarray))
elif type(Sigma) is not np.ndarray:
raise TypeError(
format_type_err_msg("epi.util.gaussian_backward_mapping", "Sigma",
Sigma, np.ndarray))
mu = np_column_vec(mu)
Sigma_shape = Sigma.shape
if len(Sigma_shape) != 2:
raise ValueError("Sigma must be 2D matrix, shape ", Sigma_shape, ".")
if Sigma_shape[0] != Sigma_shape[1]:
raise ValueError("Sigma must be square matrix, shape ", Sigma_shape,
".")
if not np.allclose(Sigma, Sigma.T, atol=1e-10):
raise ValueError("Sigma must be symmetric. shape.")
if Sigma_shape[1] != mu.shape[0]:
raise ValueError("mu and Sigma must have same dimensionality.")
D = mu.shape[0]
Sigma_inv = np.linalg.inv(Sigma)
x = np.dot(Sigma_inv, mu)
y = np.reshape(-0.5 * Sigma_inv, (D**2))
eta = np.concatenate((x[:, 0], y), axis=0)
return eta
def array_str(a):
"""Returns a compressed string from a 1-D numpy array.
:param a: A 1-D numpy array.
:type a: str
:return: A string compressed via scientific notation and repeated elements.
:rtype: str
"""
if type(a) is not np.ndarray:
raise TypeError(
format_type_err_msg("epi.util.array_str", "a", a, np.ndarray))
if len(a.shape) > 1:
raise ValueError("epi.util.array_str takes 1-D arrays not %d." %
len(a.shape))
def repeats_str(num, mult):
if mult == 1:
return "%.2E" % num
else:
return "%dx%.2E" % (mult, num)
d = a.shape[0]
if d == 1:
nums = [a[0]]
mults = [1]
else:
mults = []
nums = []
prev_num = a[0]
mult = 1
for i in range(1, d):
if a[i] == prev_num:
mult += 1
else:
nums.append(prev_num)
prev_num = a[i]
mults.append(mult)
mult = 1
if i == d - 1:
nums.append(prev_num)
mults.append(mult)
array_str = repeats_str(nums[0], mults[0])
for i in range(1, len(nums)):
array_str += "_" + repeats_str(nums[i], mults[i])
return array_str
def np_column_vec(x):
""" Takes numpy vector and orients it as a n x 1 column vec.
:param x: Vector of length n
:type x: np.ndarray
:return: n x 1 numpy column vector
:rtype: np.ndarray
"""
if type(x) is not np.ndarray:
raise (TypeError(
format_type_err_msg("epi.util.np_column_vec", "x", x, np.ndarray)))
x_shape = x.shape
if len(x_shape) == 1:
x = np.expand_dims(x, 1)
elif len(x_shape) == 2:
if x_shape[1] != 1:
if x_shape[0] > 1:
raise ValueError("x is matrix.")
else:
x = x.T
elif len(x_shape) > 2:
raise ValueError("x dimensions > 2.")
return x
def _set_name(self, name):
if type(name) is not str:
raise TypeError(format_type_err_msg(self, "name", name, str))
self.name = name
def _set_batch_norm(self, batch_norm):
if type(batch_norm) is not bool:
raise TypeError(
format_type_err_msg(self, "batch_norm", batch_norm, bool))
self.batch_norm = batch_norm
def _set_beta(self, beta):
if type(beta) not in [float, np.float32, np.float64]:
raise TypeError(format_type_err_msg(self, "beta", beta, float))
elif beta < 0.0:
raise ValueError("beta %.2E must be greater than 0." % beta)
self.beta = beta
def _set_gamma(self, gamma):
if type(gamma) not in [float, np.float32, np.float64]:
raise TypeError(format_type_err_msg(self, "gamma", gamma, float))
elif gamma < 0.0:
raise ValueError("gamma %.2E must be greater than 0." % gamma)
self.gamma = gamma
def _set_c0(self, c0):
if type(c0) not in [float, np.float32, np.float64]:
raise TypeError(format_type_err_msg(self, "c0", c0, float))
elif c0 < 0.0:
raise ValueError("c0 %.2E must be greater than 0." % c0)
self.c0 = c0
def _set_lr(self, lr):
if type(lr) not in [float, np.float32, np.float64]:
raise TypeError(format_type_err_msg(self, "lr", lr, float))
elif lr < 0.0:
raise ValueError("lr %.2E must be greater than 0." % lr)
self.lr = lr
def _set_N(self, N):
if type(N) is not int:
raise TypeError(format_type_err_msg(self, "N", N, int))
elif N < 2:
raise ValueError("N %d must be greater than 1." % N)
self.N = N
def _set_post_affine(self, post_affine):
if type(post_affine) is not bool:
raise TypeError(
format_type_err_msg(self, "post_affine", post_affine, bool))
self.post_affine = post_affine
def _set_D(self, D):
if type(D) is not int:
raise TypeError(format_type_err_msg(self, "D", D, int))
if D < 1:
raise ValueError("Dimension of parameter must be positive.")
self.D = D
def _set_random_seed(self, random_seed):
if type(random_seed) is not int:
raise TypeError(
format_type_err_msg(self, "random_seed", random_seed, int))
self.random_seed = random_seed
def load_epi_dist(
self,
mu,
k=None,
alpha=None,
nu=0.1,
arch_type="coupling",
num_stages=3,
num_layers=2,
num_units=None,
batch_norm=True,
bn_momentum=0.99,
post_affine=False,
random_seed=1,
init_type="iso_gauss",
init_params={
"loc": 0.0,
"scale": 1.0
},
N=500,
lr=1e-3,
c0=1.0,
gamma=0.25,
beta=4.0,
):
if k is not None:
if type(k) is not int:
raise TypeError(
format_type_err_msg("Model.load_epi_dist", "k", k, int))
if k < 0:
raise ValueError(
"k must be augmented Lagrangian iteration index.")
if num_units is None:
num_units = max(2 * self.D, 15)
nf = NormalizingFlow(
arch_type=arch_type,
D=self.D,
num_stages=num_stages,
num_layers=num_layers,
num_units=num_units,
batch_norm=batch_norm,
bn_momentum=bn_momentum,
post_affine=post_affine,
bounds=self._get_bounds(),
random_seed=random_seed,
)
aug_lag_hps = AugLagHPs(N, lr, c0, gamma, beta)
optimizer = tf.keras.optimizers.Adam(lr)
checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=nf)
ckpt_dir = self.get_save_path(mu, nf, aug_lag_hps)
ckpt_state = tf.train.get_checkpoint_state(ckpt_dir)
if ckpt_state is not None:
ckpts = ckpt_state.all_model_checkpoint_paths
else:
raise ValueError("No checkpoints found.")
if k is not None:
if k >= len(ckpts):
raise ValueError("Index of checkpoint 'k' too large.")
status = checkpoint.restore(ckpts[k])
status.expect_partial()
q_theta = Distribution(nf, self.parameters)
return q_theta
