# Path: cora/signal/lsscontainers.py
# class InterpolatedFunction(memh5.BasicCont):
# class CosmologyContainer(containers.ContainerBase):
# class FZXContainer(CosmologyContainer):
# class MatterPowerSpectrum(CosmologyContainer, InterpolatedFunction):
# class CorrelationFunction(CosmologyContainer, InterpolatedFunction):
# class InitialLSS(FZXContainer, containers.HealpixContainer):
# class BiasedLSS(FZXContainer, containers.HealpixContainer):
#     def __init__(self, *args, **kwargs):
#     def _finish_setup(self):
#     def get_function(self, name: str) -> Callable[[FloatArrayLike], FloatArrayLike]:
#     def add_function(
#         self, name: str, x: np.ndarray, f: np.ndarray, type: str = "linear", **kwargs
#     ):
#     def __init__(self, cosmology: Union[Cosmology, dict, None] = None, *args, **kwargs):
#     def _resolve_args(
#         cosmology: Union[Cosmology, dict, None] = None,
#         attrs_from: Optional[containers.ContainerBase] = None,
#         **kwargs,
#     ):
#     def cosmology(self):
#     def __init__(
#         self,
#         freq: Optional[np.ndarray] = None,
#         redshift: Optional[np.ndarray] = None,
#         *args,
#         **kwargs,
#     ):
#     def chi(self):
#     def redshift(self):
#     def freq(self):
#     def __init__(
#         self,
#         k: FloatArrayLike,
#         ps: FloatArrayLike,
#         *args,
#         ps_redshift: float = 0.0,
#         **kwargs,
#     ):
#     def powerspectrum(
#         self, k: FloatArrayLike, z: FloatArrayLike = 0.0
#     ) -> FloatArrayLike:
#     def powerspectrum_at_z(
#         self, z: FloatArrayLike
#     ) -> Callable[[FloatArrayLike], FloatArrayLike]:
#         def _ps(k):
#     def _ps_redshift(self):
#     def __init__(self, *args, **kwargs):
#     def delta(self):
#     def phi(self):
#     def __init__(
#         self,
#         *args,
#         lightcone: Optional[bool] = None,
#         fixed_redshift: Optional[float] = None,
#         **kwargs,
#     ):
#     def lightcone(self) -> bool:
#     def fixed_redshift(self) -> Optional[float]:
#     def delta(self) -> np.ndarray:
# 
# Path: cora/util/cosmology.py
# class Cosmology(object):
#     H0: float = 67.66
#         H = (
#             self.H0
#             * (
#                 self.omega_r * (1 + z) ** 4
#                 + self.omega_m * (1 + z) ** 3
#                 + self.omega_k * (1 + z) ** 2
#                 + self.omega_l
#                 * (1 + z) ** (3 * (1 + self.w_0 + self.w_a))
#                 * np.exp(-3 * self.w_a * z / (1 + z))
#             )
#             ** 0.5
#         )
#     def omega_m(self) -> float:
#     def omega_r(self) -> float:
#     def omega_k(self) -> float:
#     def from_physical(
#         cls,
#         ombh2: float = 0.02242,
#         omch2: float = 0.11933,
#         H0: float = 67.66,
#         omk: float = 0.0,
#         TCMB=2.7255,
#         nnu=3.046,
#     ) -> "Cosmology":
#     def to_dict(self) -> dict:
#     def H(self, z: FloatArrayLike = 0.0) -> FloatArrayLike:
#     def comoving_distance(self, z: FloatArrayLike) -> FloatArrayLike:
#         def f(z1):
#     def proper_distance(self, z: FloatArrayLike) -> FloatArrayLike:
#     def angular_distance(self, z: FloatArrayLike) -> FloatArrayLike:
#     def luminosity_distance(self, z: FloatArrayLike) -> FloatArrayLike:
#     def lookback_time(self, z: FloatArrayLike) -> FloatArrayLike:
#         def f(z1):
#     def _unit_distance(self) -> float:
#     def _unit_time(self) -> float:
#     def growth_factor(self, z: FloatArrayLike) -> FloatArrayLike:
#     def growth_rate(self, z: FloatArrayLike) -> FloatArrayLike:
# def _intf_0_z(f, z):
#     def _yp(y, z):
# def sound_horizon(c=None):
# def ps_nowiggle(kh, z=0.0, c=None):


import camb
from cora.signal import lsscontainers
from cora.util import cosmology
    "planck2018_z1.0_halofit-original": {
        "nonlinear": True,
        "halofit_version": "original",
    },
    "planck2018_z1.0_halofit-mead": {
        "nonlinear": True,
        "halofit_version": "mead2020",
    },
    "planck2018_z1.0_halofit-mead-feedback": {
        "nonlinear": True,
        "halofit_version": "mead2020_feedback",
    },
    "planck2018_z1.0_halofit-takahashi": {
        "nonlinear": True,
        "halofit_version": "takahashi",
    },
}

c = cosmology.Cosmology.from_physical(**background_params)

for name, p in extended_params.items():
    print(f"Generating {name}")

    full_params = dict(**background_params, **base_params, **p)

    cp = camb.set_params(**full_params)
    r = camb.get_results(cp)

    k, _, pk = r.get_matter_power_spectrum(minkh=1e-4, maxkh=1000, npoints=1024)


# Path: cora/signal/lsscontainers.py
# class InterpolatedFunction(memh5.BasicCont):
# class CosmologyContainer(containers.ContainerBase):
# class FZXContainer(CosmologyContainer):
# class MatterPowerSpectrum(CosmologyContainer, InterpolatedFunction):
# class CorrelationFunction(CosmologyContainer, InterpolatedFunction):
# class InitialLSS(FZXContainer, containers.HealpixContainer):
# class BiasedLSS(FZXContainer, containers.HealpixContainer):
#     def __init__(self, *args, **kwargs):
#     def _finish_setup(self):
#     def get_function(self, name: str) -> Callable[[FloatArrayLike], FloatArrayLike]:
#     def add_function(
#         self, name: str, x: np.ndarray, f: np.ndarray, type: str = "linear", **kwargs
#     ):
#     def __init__(self, cosmology: Union[Cosmology, dict, None] = None, *args, **kwargs):
#     def _resolve_args(
#         cosmology: Union[Cosmology, dict, None] = None,
#         attrs_from: Optional[containers.ContainerBase] = None,
#         **kwargs,
#     ):
#     def cosmology(self):
#     def __init__(
#         self,
#         freq: Optional[np.ndarray] = None,
#         redshift: Optional[np.ndarray] = None,
#         *args,
#         **kwargs,
#     ):
#     def chi(self):
#     def redshift(self):
#     def freq(self):
#     def __init__(
#         self,
#         k: FloatArrayLike,
#         ps: FloatArrayLike,
#         *args,
#         ps_redshift: float = 0.0,
#         **kwargs,
#     ):
#     def powerspectrum(
#         self, k: FloatArrayLike, z: FloatArrayLike = 0.0
#     ) -> FloatArrayLike:
#     def powerspectrum_at_z(
#         self, z: FloatArrayLike
#     ) -> Callable[[FloatArrayLike], FloatArrayLike]:
#         def _ps(k):
#     def _ps_redshift(self):
#     def __init__(self, *args, **kwargs):
#     def delta(self):
#     def phi(self):
#     def __init__(
#         self,
#         *args,
#         lightcone: Optional[bool] = None,
#         fixed_redshift: Optional[float] = None,
#         **kwargs,
#     ):
#     def lightcone(self) -> bool:
#     def fixed_redshift(self) -> Optional[float]:
#     def delta(self) -> np.ndarray:
# 
# Path: cora/util/cosmology.py
# class Cosmology(object):
#     H0: float = 67.66
#         H = (
#             self.H0
#             * (
#                 self.omega_r * (1 + z) ** 4
#                 + self.omega_m * (1 + z) ** 3
#                 + self.omega_k * (1 + z) ** 2
#                 + self.omega_l
#                 * (1 + z) ** (3 * (1 + self.w_0 + self.w_a))
#                 * np.exp(-3 * self.w_a * z / (1 + z))
#             )
#             ** 0.5
#         )
#     def omega_m(self) -> float:
#     def omega_r(self) -> float:
#     def omega_k(self) -> float:
#     def from_physical(
#         cls,
#         ombh2: float = 0.02242,
#         omch2: float = 0.11933,
#         H0: float = 67.66,
#         omk: float = 0.0,
#         TCMB=2.7255,
#         nnu=3.046,
#     ) -> "Cosmology":
#     def to_dict(self) -> dict:
#     def H(self, z: FloatArrayLike = 0.0) -> FloatArrayLike:
#     def comoving_distance(self, z: FloatArrayLike) -> FloatArrayLike:
#         def f(z1):
#     def proper_distance(self, z: FloatArrayLike) -> FloatArrayLike:
#     def angular_distance(self, z: FloatArrayLike) -> FloatArrayLike:
#     def luminosity_distance(self, z: FloatArrayLike) -> FloatArrayLike:
#     def lookback_time(self, z: FloatArrayLike) -> FloatArrayLike:
#         def f(z1):
#     def _unit_distance(self) -> float:
#     def _unit_time(self) -> float:
#     def growth_factor(self, z: FloatArrayLike) -> FloatArrayLike:
#     def growth_rate(self, z: FloatArrayLike) -> FloatArrayLike:
# def _intf_0_z(f, z):
#     def _yp(y, z):
# def sound_horizon(c=None):
# def ps_nowiggle(kh, z=0.0, c=None):


# Path: cora/signal/data/generate_powerspectra.py
import camb
from cora.signal import lsscontainers
from cora.util import cosmology

    "planck2018_z1.0_halofit-original": {
        "nonlinear": True,
        "halofit_version": "original",
    },
    "planck2018_z1.0_halofit-mead": {
        "nonlinear": True,
        "halofit_version": "mead2020",
    },
    "planck2018_z1.0_halofit-mead-feedback": {
        "nonlinear": True,
        "halofit_version": "mead2020_feedback",
    },
    "planck2018_z1.0_halofit-takahashi": {
        "nonlinear": True,
        "halofit_version": "takahashi",
    },
}

c = cosmology.Cosmology.from_physical(**background_params)

for name, p in extended_params.items():
    print(f"Generating {name}")

    full_params = dict(**background_params, **base_params, **p)

    cp = camb.set_params(**full_params)
    r = camb.get_results(cp)

    k, _, pk = r.get_matter_power_spectrum(minkh=1e-4, maxkh=1000, npoints=1024)


    pkc = lsscontainers.MatterPowerSpectrum(