def prepare_model(rv_single_epoch_variance, target=0.01, model_path="model.stan", S=None, design_matrix_function=None, mask=None, **source_params): """ Compile the Stan model, and prepare the data and initialisation dictionaries for optimization and sampling. :param rv_single_epoch_variance: The variance in single epoch measurements of radial velocity. :param target: [optional] The target radial velocity variance to use when initialising the model coefficients. :param model_path: [optional] The local path of the Stan model. :param S: [optional] If not `None`, draw a random `S` valid sources from the data rather than using the full data set. :Keyword Arguments: * *source_params* (``dict``) These are passed directly to the `_rvf_design_matrix`, so they should include all of the source labels needed to construct the design matrix (e.g., `phot_rp_mean_flux`). The array length for each source label should match that of `rv_single_epoch_variance`. :returns: The compiled Stan model, the data dictionary, the initialsiation dictionary, and a mask corresponding to which `sources` were used to construct the data dictionary. """ data, indices = prepare_data(rv_single_epoch_variance, S=S, mask=mask, design_matrix_function=design_matrix_function, **source_params) coeff = _rvf_initial_coefficients(data["design_matrix"].T, target=target) init = dict(theta=0.1, mu_coefficients=coeff, sigma_coefficients=coeff) model = stan.load_stan_model(model_path) return (model, data, init, indices)
def _hierarchically_fit_waveform_parameters(self, waveform_parameters, **kwargs): r""" Fit a hierarchical model to the waveform parameters. :param waveform_parameters: A table of parameters for each of the numerical relativity waveforms that has [n_waveforms, ] rows and contains at least the following properties (as columns): The compactness `C`, the calculated :math:`\kappa` value `Kappa_calc`, the mass of the primary `M1`, the frequency of the high-frequency peak `f2`, the dimensionless tidal deformability `Lambda` ..math: \Lambda = \left(\frac{\lambda}{M_{1}^{5}}\right)^{1/5} where ..math: \lambda = \frac{2}{3}\overbar{kappa}_{2}^{T}\overbar{R}^5 """ y = np.array([waveform_parameters[pn] \ for pn in self.hierarchical_parameter_names]).T N, D = y.shape M1 = waveform_parameters["M1"] eos_param = waveform_parameters[self.eos_parameter_name] data_dict = dict(M1=M1, eos_param=eos_param, N=N, D=D, y=y) kwds = dict(iter=10000, tol_param=1e-12, tol_obj=1e-12, tol_grad=1e-12, tol_rel_grad=1e8) kwds.update(kwargs) p_opt = stan.load_stan_model("hpm.stan").optimizing(data_dict, **kwds) return (p_opt, data_dict)
N, D = X.shape F = finite_indices.size L = 4 * len(config["predictor_label_names"]) + 1 # + 1 if using mu_multiple_uv C = config["share_optimised_result_with_nearest"] kdt, scales, offsets = npm.build_kdtree( X[finite], relative_scales=config["kdtree_relative_scales"]) kdt_kwds = dict(offsets=offsets, scales=scales, full_output=False) kdt_kwds.update( minimum_radius=config["kdtree_minimum_radius"], minimum_points=config["kdtree_minimum_points"], maximum_points=config["kdtree_maximum_points"], minimum_density=config.get("kdtree_minimum_density", None)) model = stan.load_stan_model(config["model_path"], verbose=False) default_opt_kwds = config.get("optimisation_kwds", {}) # Make sure that some entries have the right units. for key in ("tol_obj", "tol_grad", "tol_rel_grad", "tol_rel_obj"): if key in default_opt_kwds: default_opt_kwds[key] = float(default_opt_kwds[key]) logging.info("k-d tree keywords: {}".format(kdt_kwds)) logging.info("optimization keywords: {}".format(default_opt_kwds)) done = np.zeros(N, dtype=bool) queued = np.zeros(N, dtype=bool) results = np.nan * np.ones((N, L), dtype=float)
import numpy as np import matplotlib.pyplot as plt import stan_utils as stan x, y, x_err, y_err = np.loadtxt("data.txt") x_err, y_err = (np.sqrt(np.abs(x_err)), np.sqrt(np.abs(y_err))) model = stan.load_stan_model("line.stan") data_dict = dict(x=x, y=y, x_err=x_err, y_err=y_err, N=x.size) # Don't really need an initialisation but f**k it. DM = np.vstack((np.ones_like(x), x)).T C = np.diag(y_err**2) cov = np.linalg.inv(DM.T @ np.linalg.solve(C, DM)) c, m = cov @ (DM.T @ np.linalg.solve(C, y)) init_dict = dict(c=c, m=m, x_t=x) # Optimize p_opt = model.optimizing(data=data_dict, init=init_dict, iter=2000) # Sample samples = model.sampling( **stan.sampling_kwds(data=data_dict, chains=2, init=p_opt, iter=2000)) fig = samples.traceplot() fig.savefig("trace.png")
binary_sigma_mta[i] = v_std binary_sigma_rv[i] = sigma_mta_to_sigma_rv(v_std, O) assert np.isfinite(v_std) print( f"Number of single stars with finite values: {np.isfinite(single_sigma_rv).sum()}" ) print( f"Number of binary stars with finite values: {np.isfinite(binary_sigma_rv).sum()}" ) # Now let us imagine that we were fitting these data with a mixture model in order to estimate the # intrinsic uncertainty. model = stan.load_stan_model("nlnmm-fixed.stan") y_rv = np.hstack([single_sigma_rv, binary_sigma_rv]) y_rv = y_rv[np.isfinite(y_rv)] data_dict = dict(y=y_rv, N=y_rv.size) data_dict.update(data_kwds) init_dict = dict(mu_single=np.min([np.median(y_rv, axis=0), 10])) init_dict.update(init_kwds) p_opt_rv = model.optimizing(data=data_dict, init=init_dict) fig, ax = plt.subplots() ax.hist(y_rv, bins=np.linspace(0, 10, 100))
X_kdt = X_kdt[subset] X_scale = np.ptp(X_kdt, axis=0) X_mean = np.mean(X_kdt, axis=0) _scale = lambda a: (a - X_mean) / X_scale _descale = lambda a: a * X_scale + X_mean # Normalise the array for the KD-tree X_norm = _scale(X_kdt) # Construct the KD-Tree kdt = npm_utils.build_kdt(X_norm) data = data[subset] model = stan.load_stan_model("npm.stan") # Calculate the total number of parameters M = len(data) L = len(predictor_label_names) K = 1 + 4 * L opt_params = np.empty((M, K)) #subset_points = np.random.choice(M, size=1000, replace=False) #for i in range(M): #for j, i in enumerate(subset_points): for j, i in enumerate(range(M)): print("At point {}/{}: {}".format(j, M, i))
def fit(self, waveform_parameters, frequencies, amplitudes, **kwargs): r""" Fit the model given the parameters of the waveforms, the common frequencies that those waveforms are calculated on, and the amplitudes of the Fourier spectrum for those waveforms. :param waveform_parameters: A table of parameters for each of the numerical relativity waveforms that has [n_waveforms, ] rows and may contain the following properties (as columns): The compactness `C`, the calculated :math:`\kappa` value `Kappa_calc`, the mass of the primary `M1`, the frequency of the high-frequency peak `f2`, the dimensionless tidal deformability `Lambda` ..math: \Lambda = \left(\frac{\lambda}{M_{1}^{5}}\right)^{1/5} where ..math: \lambda = \frac{2}{3}\overbar{kappa}_{2}^{T}\overbar{R}^5 :param frequencies: The common frequencies (in kHz) that each numerical relativity waveform is calculated on. This should be an array of shape [n_frequencies, ]. :param amplitudes: The amplitudes of the Fourier spectrum for the numerical relativity waveforms. This should be an array of shape [n_waveforms, n_frequencies]. """ self._frequencies = frequencies self._amplitudes = amplitudes self._waveform_parameters = waveform_parameters # Hierarchically fit the waveform parameters. p_opt_hpm, data_dict = self._hierarchically_fit_waveform_parameters( waveform_parameters) # Frequency-shift waveforms so that they are aligned at mean f2. shifted_amplitudes, f2_coeff, f2_mean = self._frequency_shift_waveforms( waveform_parameters, frequencies, amplitudes) # Fit the waveforms with a linear model of the whitened labels, and the # mass and dimensionless tidal deformability. labels = np.vstack([ waveform_parameters["M1"], waveform_parameters[self.eos_parameter_name], data_dict["y"].T ]).T mu, sigma = (np.mean(labels, axis=0), np.std(labels, axis=0)) whitened_labels = (labels - mu) / sigma model = stan.load_stan_model("waveform.stan") N, P = whitened_labels.shape F = frequencies.size data_dict = dict(F=frequencies.size, N=N, P=P, y=shifted_amplitudes, whitened_labels=whitened_labels) # TODO: move default op kwds to somewhere else. kwds = dict(iter=100000, tol_param=1e-12, tol_obj=1e-12, tol_grad=1e-12, tol_rel_grad=1e8) kwds.update(kwargs) p_opt_waveform = model.optimizing(data=data_dict, **kwds) # Save attributes so that we can make predictions. self._p_opt_waveform = p_opt_waveform self._p_opt_hpm = p_opt_hpm self._p_opt_f2 = (f2_mean, f2_coeff) self._p_opt_label_whiten = (mu, sigma) return self
axes[0, 0].set_ylabel(r"$J_{{{0}}}$".format(2)) axes[0, 1].set_visible(False) axes[1, 1].scatter(truths["theta"].T[2], truths["theta"].T[1]) axes[1, 1].set_xlabel(r"$J_{{{0}}}$".format(2)) axes[1, 1].set_ylabel(r"$J_{{{0}}}$".format(1)) fig.tight_layout() op_kwds = dict(init_alpha=1, tol_obj=1e-16, tol_rel_grad=1e-16, tol_rel_obj=1e-16) op_kwds = dict(data=data, seed=seed) model = stan.load_stan_model("mlf.stan") s_opt = model.optimizing(**op_kwds) fig, ax = plt.subplots() #ax.scatter(np.diag(psi), s_opt["psi"], facecolor="b") ax.scatter(np.diag(truths["psi"]), s_opt["psi"]) ax.set_title(r"$\psi$") limits = np.hstack([ax.get_xlim(), ax.get_ylim()]) limits = np.array([limits.min(), limits.max()]) ax.plot(limits, limits, c="#666666", zorder=-1) ax.set_xlim(limits) ax.set_ylim(limits) fig, axes = plt.subplots(truths["L"].shape[0], figsize=(4, 12)) for i, ax in enumerate(axes):
# Require finite entries for all predictors across all models. sources = data["sources"] M = config["number_of_sources_for_gaussian_process"] N = sources["source_id"].size # Create results file. with h5.File(results_path, "w") as results: # Add config. results.attrs.create("config", np.string_(yaml.dump(config))) results.attrs.create("config_path", np.string_(config_path)) # Load model and check optimization keywords model_path = os.path.join(pwd, config["model_path"]) model = stan.load_stan_model(model_path, verbose=False) # Make sure that some entries have the right type. default_opt_kwds = config.get("optimisation_kwds", {}) for key in ("tol_obj", "tol_grad", "tol_rel_grad", "tol_rel_obj"): if key in default_opt_kwds: default_opt_kwds[key] = float(default_opt_kwds[key]) logger.info(f"Optimization keywords:\n{utils.repr_dict(default_opt_kwds)}") default_bounds = dict(theta=[0.5, 1], mu_single=[0.5, 15], sigma_single=[0.05, 10], sigma_multiple=[0.2, 1.6]) # Plotting
y_err_intrinsic = 20 y_err = np.abs(np.random.normal(0, y_err_intrinsic, size=N)) x_true = np.random.uniform(0, 30, N) x = x_true + np.random.normal(0, 1, size=N) * x_err y_true = m_true * x + b_true y = y_true + np.random.normal(0, 1, size=N) * y_err fig, ax = plt.subplots() ax.scatter(x, y) ax.errorbar(x, y, yerr=y_err) data_dict = dict(N=N, x=x, y=y, y_err=y_err, x_err=x_err) init_dict = dict(m=0, b=0, x_true=x) sm = stan.load_stan_model("advanced_line.stan") opt = sm.optimizing(data=data_dict) sampling = sm.sampling(**stan.sampling_kwds(init=init_dict, data=data_dict, chains=2, iter=1000)) chains = sampling.extract() X = np.array([chains["m"], chains["b"]]).T fig = corner(X, truths=truths) plt.show()