diff --git a/sunbird/cosmology/growth_rate.py b/sunbird/cosmology/growth_rate.py index a59634d..3462375 100644 --- a/sunbird/cosmology/growth_rate.py +++ b/sunbird/cosmology/growth_rate.py @@ -18,11 +18,11 @@ class Growth: def __init__( self, - theta_star: float = AbacusSummit(0).theta_star, + theta_MC_100: float = AbacusSummit(0)['theta_MC_100'], emulate=False, emulator_data_dir=DEFAULT_PATH / "data/hemu/", ): - self.theta_star = theta_star + self.theta_MC_100 = theta_MC_100 self.emulate = emulate self.emulator_data_dir = emulator_data_dir if self.emulate: @@ -59,10 +59,13 @@ def generate_emulator_training_data( h_values, sample_parameters = [], [] for i, sample in enumerate(samples_matrix): try: - cosmology = self.get_cosmology_fixed_theta_star( + # print every 1000th sample + if i % 1000 == 0: + print(i) + cosmology = self.get_cosmology_fixed_theta_MC_100( DESI(engine="class"), dict( - theta_star=self.theta_star, + theta_MC_100=self.theta_MC_100, omega_b=sample[0], omega_cdm=sample[1], sigma8=sample[2], @@ -274,29 +277,30 @@ def get_emulated_h(self, omega_b, omega_cdm, sigma8, N_ur, n_s, w0_fld, wa_fld): x = jnp.vstack([omega_b, omega_cdm, sigma8, N_ur, n_s, w0_fld, wa_fld]).T return self.model.apply(self.params, x) - def get_cosmology_fixed_theta_star( + def get_cosmology_fixed_theta_MC_100( self, fiducial, params, h_limits=[0.4, 1.0], xtol=1.0e-6, ): - theta = params.pop("theta_star", None) + theta = params.pop("theta_MC_100", None) fiducial = fiducial.clone(base="input", **params) if theta is not None: if "h" in params: - raise ValueError("Cannot provide both theta_star and h") + raise ValueError("Cannot provide both theta_MC_100 and h") def f(h): cosmo = fiducial.clone(base="input", h=h) - return 100.0 * (theta - cosmo.get_thermodynamics().theta_star) + # return 100.0 * (theta - cosmo.get_thermodynamics().theta_MC_100) + return 100.0 * (theta - cosmo['theta_MC_100']) rtol = xtol try: h = optimize.bisect(f, *h_limits, xtol=xtol, rtol=rtol, disp=True) except ValueError as exc: raise ValueError( - "Could not find proper h value in the interval that matches theta_star = {:.4f} with [f({:.3f}), f({:.3f})] = [{:.4f}, {:.4f}]".format( + "Could not find proper h value in the interval that matches theta_MC_100 = {:.4f} with [f({:.3f}), f({:.3f})] = [{:.4f}, {:.4f}]".format( theta, *h_limits, *list(map(f, h_limits)) ) ) from exc @@ -335,10 +339,10 @@ def get_growth( z=z, ) else: - cosmology = self.get_cosmology_fixed_theta_star( + cosmology = self.get_cosmology_fixed_theta_MC_100( DESI(engine="class"), dict( - theta_star=self.theta_star, + theta_MC_100=self.theta_MC_100, omega_b=omega_b, omega_cdm=omega_cdm, sigma8=sigma8, @@ -393,10 +397,10 @@ def get_fsigma8( ) return growth_rate * sigma8_z else: - cosmology = self.get_cosmology_fixed_theta_star( + cosmology = self.get_cosmology_fixed_theta_MC_100( DESI(engine="class"), dict( - theta_star=self.theta_star, + theta_MC_100=self.theta_MC_100, omega_b=omega_b, omega_cdm=omega_cdm, sigma8=sigma8, @@ -414,6 +418,6 @@ def get_fsigma8( t0 = time.time() growth = Growth() - # growth.generate_emulator_training_data() + growth.generate_emulator_training_data(n_samples=100_000) growth.train_emulator() print(f"It took {time.time() - t0} seconds") diff --git a/sunbird/data/transforms_array.py b/sunbird/data/transforms_array.py index b2462b9..abb92b1 100644 --- a/sunbird/data/transforms_array.py +++ b/sunbird/data/transforms_array.py @@ -11,6 +11,27 @@ def transform(self, x): @abstractmethod def inverse_transform(self, x): pass + + @abstractmethod + def get_jacobian_diagonal(self, y): + """ + Get the diagonal of the Jacobian matrix df/dy for transforming covariance matrices. + + For an element-wise transformation f(y), the transformed covariance is: + Cov_transformed = diag(J) @ Cov @ diag(J) + where J = df/dy is the Jacobian diagonal. + + Parameters + ---------- + y : array_like + Data vector in the original (untransformed) space. + + Returns + ------- + array_like + Diagonal of the Jacobian matrix, same shape as y. + """ + pass class LogTransform(BaseTransform): @@ -22,6 +43,27 @@ def transform(self, x): def inverse_transform(self, x): return 10**x + + def get_jacobian_diagonal(self, y): + """ + Get Jacobian diagonal for log10 transform: d(log10(y))/dy = 1/(y * ln(10)) + + Parameters + ---------- + y : array_like + Data vector in the original (untransformed) space. + + Returns + ------- + array_like + Jacobian diagonal: 1/(y * ln(10)) + """ + if type(y) == torch.Tensor: + return 1.0 / (y * torch.log(torch.tensor(10.0))) + elif type(y) == np.ndarray: + return 1.0 / (y * np.log(10.0)) + else: + return 1.0 / (y * jnp.log(10.0)) class ArcsinhTransform(BaseTransform): @@ -40,6 +82,27 @@ def inverse_transform(self, x): return np.sinh(x) else: return jnp.sinh(x) + + def get_jacobian_diagonal(self, y): + """ + Get Jacobian diagonal for arcsinh transform: d(arcsinh(y))/dy = 1/sqrt(1 + y^2) + + Parameters + ---------- + y : array_like + Data vector in the original (untransformed) space. + + Returns + ------- + array_like + Jacobian diagonal: 1/sqrt(1 + y^2) + """ + if type(y) == torch.Tensor: + return 1.0 / torch.sqrt(1.0 + y**2) + elif type(y) == np.ndarray: + return 1.0 / np.sqrt(1.0 + y**2) + else: + return 1.0 / jnp.sqrt(1.0 + y**2) class WeiLiuOutputTransForm(BaseTransform): """Class to reconcile output the Minkowski functionals model @@ -56,6 +119,27 @@ def transform(self, x): def inverse_transform(self, x): return x * self.std + self.mean + + def get_jacobian_diagonal(self, y): + """ + Get Jacobian diagonal for affine transform: d(y * std + mean)/dy = std + + Parameters + ---------- + y : array_like + Data vector in the original (untransformed) space. + + Returns + ------- + array_like + Jacobian diagonal: std (broadcast to match y shape) + """ + if type(y) == torch.Tensor: + return torch.ones_like(y) * self.std + elif type(y) == np.ndarray: + return np.ones_like(y) * self.std.numpy() + else: + return jnp.ones_like(y) * self.std.numpy() class WeiLiuInputTransform(BaseTransform): """Class to reconcile input of the Minkowski functionals model @@ -72,4 +156,25 @@ def transform(self, x): def inverse_transform(self, x): return x + + def get_jacobian_diagonal(self, y): + """ + Get Jacobian diagonal for standardization: d((y - mean) / std)/dy = 1/std + + Parameters + ---------- + y : array_like + Data vector in the original (untransformed) space. + + Returns + ------- + array_like + Jacobian diagonal: 1/std (broadcast to match y shape) + """ + if type(y) == torch.Tensor: + return torch.ones_like(y) / self.std + elif type(y) == np.ndarray: + return np.ones_like(y) / self.std.numpy() + else: + return jnp.ones_like(y) / self.std.numpy() \ No newline at end of file diff --git a/sunbird/emulators/models/fcn.py b/sunbird/emulators/models/fcn.py index be93e4a..055a6a9 100644 --- a/sunbird/emulators/models/fcn.py +++ b/sunbird/emulators/models/fcn.py @@ -89,7 +89,7 @@ def __init__( ], ) self.compression_matrix = compression_matrix - + @staticmethod def add_model_specific_args(parent_parser): """Model arguments that could vary @@ -275,16 +275,28 @@ def forward(self, x: Tensor) -> Tensor: y_var = torch.zeros_like(y_pred) return y_pred, y_var - def get_prediction(self, x: Tensor, filters: Optional[dict] = None) -> Tensor: + def get_prediction(self, x: Tensor, filters: Optional[dict] = None, skip_output_inverse_transform: bool = False) -> Tensor: + """Get prediction from the model. + + Args: + x (Tensor): Input tensor + filters (dict, optional): Filters to apply. Defaults to None. + skip_output_inverse_transform (bool, optional): If True, skip the output inverse transformation, + keeping predictions in the transformed space. Useful when performing inference in transformed + space (requires transforming observations and covariance to match). Defaults to False. + + Returns: + Tensor: Model prediction + """ x = torch.Tensor(x) - if self.transform_input: + if self.transform_input is not None: x = self.transform_input.transform(x) y, _ = self.forward(x) if self.standarize_output: std_output = self.std_output.to(x.device) mean_output = self.mean_output.to(x.device) y = y * std_output + mean_output - if self.transform_output: + if self.transform_output is not None and not skip_output_inverse_transform: y = self.transform_output.inverse_transform(y) if self.compression_matrix is not None: y = y @ self.compression_matrix diff --git a/sunbird/inference/base.py b/sunbird/inference/base.py index c632796..3271453 100644 --- a/sunbird/inference/base.py +++ b/sunbird/inference/base.py @@ -1,3 +1,5 @@ +"""Base classes and utilities for inference samplers.""" + import logging import numpy as np from tabulate import tabulate @@ -5,7 +7,14 @@ from sunbird.inference.priors import AbacusSummitEllipsoid class BaseSampler: - def __init__(self, + """Base class for inference samplers. + + Handles parameter bookkeeping, optional transformed-space sampling, and + convenience utilities for saving chains and summary tables. + """ + + def __init__( + self, observation, precision_matrix, theory_model, @@ -13,33 +22,76 @@ def __init__(self, ranges: Optional[Dict[str, tuple]] = {}, labels: Dict[str, str] = {}, fixed_parameters: Dict[str, float] = {}, - slice_filters: Dict = {}, - select_filters: Dict = {}, - coordinates: list = [], ellipsoid: bool = False, markers: dict = {}, + sample_in_transformed_space: bool = False, **kwargs, ): + """Initialize the sampler base. + + Args: + observation: Observed data vector. + precision_matrix: Inverse covariance matrix. + theory_model: Callable model that maps parameters to predictions. + priors: Mapping of parameter names to prior objects. + ranges: Optional plotting or reporting ranges by parameter. + labels: Optional labels by parameter. + fixed_parameters: Mapping of parameter names to fixed values. + ellipsoid: Whether to include the AbacusSummit ellipsoid prior. + markers: Optional marker styling for plots. + sample_in_transformed_space: If True, use transformed outputs. + **kwargs: Extra arguments for subclasses. + """ self.logger = logging.getLogger(self.__class__.__name__) self.theory_model = theory_model if fixed_parameters is None: fixed_parameters = {} self.fixed_parameters = fixed_parameters - self.observation = observation self.priors = priors self.ranges = ranges self.labels = labels - self.precision_matrix = precision_matrix self.ellipsoid = ellipsoid self.markers = markers + self.sample_in_transformed_space = sample_in_transformed_space + + # Handle transformation of observations and covariance + if sample_in_transformed_space: + # Validate that the observable has an output transform + if not hasattr(theory_model.__self__.model, 'transform_output'): + raise ValueError('Cannot sample in transformed space: observable does not have a transform_output. ' + 'Either set sample_in_transformed_space=False or use an observable with transform_output.') + + # Check if transform_output is valid (not None or empty list) + transform = theory_model.__self__.model.transform_output + if transform is None: + raise ValueError('Cannot sample in transformed space: transform_output is None. ' + 'Either set sample_in_transformed_space=False or use an observable with transform_output.') + + # For combined observables, transform_output is a list + if isinstance(transform, list): + if all(t is None for t in transform): + raise ValueError('Cannot sample in transformed space: all transforms in combined observable are None. ' + 'Either set sample_in_transformed_space=False or use observables with transform_output.') + + self.logger.warning('Sampling in transformed space (skip_output_inverse_transform=True). ' + 'Ensure observations and covariance matrix are also transformed to match!') + + self.observation = observation + self.precision_matrix = precision_matrix + if self.ellipsoid: self.abacus_ellipsoid = AbacusSummitEllipsoid() + self.ndim = len(self.priors.keys()) - len(self.fixed_parameters.keys()) self.logger.info(f'Free parameters: {[key for key in priors.keys() if key not in fixed_parameters.keys()]}') self.logger.info(f'Fixed parameters: {[key for key in priors.keys() if key in fixed_parameters.keys()]}') def save_chain(self, save_fn, metadata=None): - """Save the chain to a file + """Save a chain dictionary to a NumPy file. + + Args: + save_fn: Output filename for the NumPy archive. + metadata: Optional extra metadata to include. """ data = self.get_chain(flat=True) names = [param for param in self.priors.keys() if param not in self.fixed_parameters] @@ -66,6 +118,11 @@ def save_chain(self, save_fn, metadata=None): np.save(save_fn, cout) def save_table(self, save_fn): + """Write a summary table with MAP/mean/std values. + + Args: + save_fn: Output filename for the text table. + """ chain = self.get_chain(flat=True) maxp = chain['samples'][chain['log_posterior'].argmax()] mean = chain['samples'].mean(axis=0) diff --git a/sunbird/inference/pocomc.py b/sunbird/inference/pocomc.py index 6230575..e75862a 100644 --- a/sunbird/inference/pocomc.py +++ b/sunbird/inference/pocomc.py @@ -1,3 +1,5 @@ +"""Samplers based on the `pocomc` inference engine.""" + import torch import pocomc import numpy as np @@ -5,17 +7,20 @@ class PocoMCSampler(BaseSampler): + """PoCoMC sampler wrapper with optional ellipsoid prior support.""" + def __init__(self, **kwargs): + """Initialize the PoCoMC sampler wrapper.""" super().__init__(**kwargs) def fill_params(self, theta): - """Fill the parameter vector to include fixed parameters + """Fill a parameter vector to include fixed parameters. Args: - theta (np.array): input parameters + theta: Free parameter vector. Returns: - np.array: filled parameters + Filled parameter vector with fixed values inserted. """ params = np.ones(len(self.priors.keys())) itheta = 0 @@ -28,13 +33,13 @@ def fill_params(self, theta): return params def fill_params_batch(self, thetas): - """Fill the batch of parameter vectors to include fixed parameters + """Fill a batch of parameter vectors to include fixed parameters. Args: - thetas (np.array): input parameters + thetas: Batch of free parameter vectors. Returns: - np.array: filled parameters + Filled parameter array with fixed values inserted. """ params = np.ones((len(thetas), len(self.priors.keys()))) for i, theta in enumerate(thetas): @@ -42,29 +47,27 @@ def fill_params_batch(self, thetas): return params def get_model_prediction(self, theta): - """Get model prediction + """Return the model prediction for the given parameters. Args: - theta (np.array): input parameters + theta: Parameter vector or batch of vectors. Returns: - np.array: model prediction + Model prediction as a NumPy array. """ - # pred = self.theory_model.get_prediction(x=theta) - pred = self.theory_model(x=theta) - # detach if using torch + pred = self.theory_model(x=theta, skip_output_inverse_transform=self.sample_in_transformed_space) if isinstance(pred, torch.Tensor): pred = pred.detach().numpy() return pred def log_likelihood(self, theta): - """Log likelihood function + """Compute the log likelihood for a parameter vector or batch. Args: - theta (np.array): input parameters + theta: Free parameter vector or batch of vectors. Returns: - float: log likelihood + Log likelihood value(s). """ batch = len(theta.shape) > 1 params = self.fill_params_batch(theta) if batch else self.fill_params(theta) @@ -80,14 +83,24 @@ def log_likelihood(self, theta): logl += self.abacus_ellipsoid.log_likelihood(params[:8]) return logl - def __call__(self, vectorize=True, random_state=0, precondition=True, n_total=4096, progress=True, **kwargs): - """Run the sampler + def __call__( + self, + vectorize=True, + random_state=0, + precondition=True, + n_total=4096, + progress=True, + **kwargs, + ): + """Run the PoCoMC sampler. Args: - vectorize (bool, optional): Vectorize the log likelihood call. Defaults to False. - random_state (int, optional): Random seed. Defaults to 0. - precondition (bool, optional): If False, use standard MCMC without normalizing flow. Defaults to True. - kwargs: Additional arguments for the sampler + vectorize: Vectorize the log likelihood call. + random_state: Random seed for the sampler. + precondition: If False, disable normalizing flow preconditioning. + n_total: Total number of samples to draw. + progress: Whether to display progress output. + **kwargs: Additional arguments for `pocomc.Sampler`. """ prior = pocomc.Prior([value for key, value in self.priors.items() if key not in self.fixed_parameters.keys()]) @@ -103,27 +116,20 @@ def __call__(self, vectorize=True, random_state=0, precondition=True, n_total=40 self.sampler.run(progress=progress, n_total=n_total) def get_chain(self, **kwargs): - """Get the chain from the sampler - - Returns: - np.array: chain - """ + """Return the posterior samples and derived quantities.""" samples, weights, loglike, logprior = self.sampler.posterior() logz, logz_err = self.sampler.evidence() logposterior = loglike + logprior - logz return {'samples': samples, 'weights': weights, 'log_likelihood': loglike, 'log_prior': logprior, 'log_posterior': logposterior} - def evidence(self,): - """Get the evidence from the sampler - - Returns: - tuple: logz, logz_err - """ + def evidence(self): + """Return the evidence estimate and its error.""" return self.sampler.evidence() class PocoMCPriorSampler(PocoMCSampler): + """PoCoMC sampler that returns a flat likelihood over the prior.""" def __init__( self, observation=None, @@ -131,17 +137,11 @@ def __init__( theory_model=None, **kwargs, ): + """Initialize a prior-only sampler.""" super().__init__(observation, precision_matrix, theory_model, **kwargs) def log_likelihood(self, theta): - """Log likelihood function - - Args: - theta (np.array): input parameters - - Returns: - float: log likelihood - """ + """Return a flat log likelihood with optional ellipsoid term.""" batch = len(theta.shape) > 1 params = self.fill_params_batch(theta) if batch else self.fill_params(theta) if batch: