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