Training data generation with CLASS

cosmopower/cosmopower_NN.py

@@ -118,22 +118,32 @@ def __init__(self, parameters: Optional[Sequence[str]] = None,
         self.alphas = []
         self.betas = []
         for i in range(self.n_layers):
-            self.W.append(tf.Variable(
-                tf.random.normal([
-                    self.architecture[i], self.architecture[i + 1]
-                ], 0.0, 1e-3),
-                name="W_" + str(i), trainable=trainable))
-            self.b.append(tf.Variable(tf.zeros([self.architecture[i + 1]]),
-                                      name="b_" + str(i),
-                                      trainable=trainable))
+            self.W.append(self.add_weight(
+                initializer=tf.keras.initializers.RandomNormal(stddev=1e-3),
+                shape=(self.architecture[i], self.architecture[i + 1]),
+                name="W_" + str(i),
+                trainable=trainable
+            ))
+            self.b.append(self.add_weight(
+                initializer="zeros",
+                shape=(self.architecture[i + 1],),
+                name="b_" + str(i),
+                trainable=trainable
+            ))
 
         for i in range(self.n_layers - 1):
-            self.alphas.append(tf.Variable(
-                tf.random.normal([self.architecture[i + 1]]),
-                name="alphas_" + str(i), trainable=trainable))
-            self.betas.append(tf.Variable(
-                tf.random.normal([self.architecture[i + 1]]),
-                name="betas_" + str(i), trainable=trainable))
+            self.alphas.append(self.add_weight(
+                initializer="random_normal",
+                shape=(self.architecture[i + 1],),
+                name="alphas_" + str(i),
+                trainable=trainable
+            ))
+            self.betas.append(self.add_weight(
+                initializer="random_normal",
+                shape=(self.architecture[i + 1],),
+                name="betas_" + str(i),
+                trainable=trainable
+            ))
 
         # restore weights if restoring
         if restore_filename is not None:
@@ -546,7 +556,6 @@ def ten_to_rescaled_predictions_np(self, parameters_dict: dict
         return 10.**self.rescaled_predictions_np(parameters_dict)
 
     # Infrastructure for network training
-
     @tf.function
     def compute_loss(self, training_parameters: tf.Tensor,
                      training_features: tf.Tensor) -> tf.Tensor:
@@ -595,7 +604,6 @@ def compute_loss_and_gradients(self,
         """
         # compute loss on the tape
         with tf.GradientTape() as tape:
-
             # loss
             loss = tf.sqrt(
                 tf.reduce_mean(
@@ -888,7 +896,6 @@ def train(self, training_data: Sequence,
                 for epoch in t:
                     # loop over batches
                     for theta, feats in training_data:
-
                         # training step: check whether to accumulate gradients
                         # or not (only worth doing this for very large batch
                         # sizes)

cosmopower/cosmopower_PCAplusNN.py

@@ -128,21 +128,32 @@ def __init__(self, cp_pca: Optional[object] = None,
         self.alphas = []
         self.betas = []
         for i in range(self.n_layers):
-            self.W.append(tf.Variable(
-                tf.random.normal([
-                    self.architecture[i], self.architecture[i + 1]
-                ], 0.0, np.sqrt(2.0 / self.n_parameters)),
-                name="W_" + str(i), trainable=trainable))
-            self.b.append(tf.Variable(tf.zeros([self.architecture[i + 1]]),
-                                      name="b_" + str(i),
-                                      trainable=trainable))
+            self.W.append(self.add_weight(
+                initializer=tf.keras.initializers.RandomNormal(stddev=1e-3),
+                shape=(self.architecture[i], self.architecture[i + 1]),
+                name="W_" + str(i),
+                trainable=trainable
+            ))
+            self.b.append(self.add_weight(
+                initializer="zeros",
+                shape=(self.architecture[i + 1],),
+                name="b_" + str(i),
+                trainable=trainable
+            ))
+
         for i in range(self.n_layers - 1):
-            self.alphas.append(tf.Variable(
-                tf.random.normal([self.architecture[i + 1]]),
-                name="alphas_" + str(i), trainable=trainable))
-            self.betas.append(tf.Variable(
-                tf.random.normal([self.architecture[i + 1]]),
-                name="betas_" + str(i), trainable=trainable))
+            self.alphas.append(self.add_weight(
+                initializer="random_normal",
+                shape=(self.architecture[i + 1],),
+                name="alphas_" + str(i),
+                trainable=trainable
+            ))
+            self.betas.append(self.add_weight(
+                initializer="random_normal",
+                shape=(self.architecture[i + 1],),
+                name="betas_" + str(i),
+                trainable=trainable
+            ))
 
         # restore weights if restoring
         if restore_filename is not None:
@@ -923,8 +934,8 @@ def train(self, training_data: Sequence,
                 ).shuffle(n_training).batch(batch_sizes[i])
 
             # set up training loss
-            validation_loss = [np.infty]
-            best_loss = np.infty
+            validation_loss = [np.inf]
+            best_loss = np.inf
             early_stopping_counter = 0
 
             # loop over epochs

cosmopower/dataset.py

@@ -170,7 +170,7 @@ def read_parameters(self, indices: Union[int, np.ndarray]) -> dict:
         cast_1d = (type(indices) is int)
         indices = np.atleast_1d(indices)
 
-        entries = np.where(np.in1d(self.file["data"]["indices"][:],
+        entries = np.where(np.isin(self.file["data"]["indices"][:],
                                    indices))[0]
 
         parameters = {
@@ -179,7 +179,7 @@ def read_parameters(self, indices: Union[int, np.ndarray]) -> dict:
         }
 
         if cast_1d:
-            parameters = {k: float(v) for k, v in parameters.items()}
+            parameters = {k: float(v[0]) for k, v in parameters.items()}
 
         return parameters
 
@@ -193,7 +193,7 @@ def read_spectra(self, indices: Union[int, np.ndarray]) -> np.ndarray:
         cast_1d = (type(indices) is int)
         indices = np.atleast_1d(indices)
 
-        entries = np.where(np.in1d(self.file["data"]["indices"][:],
+        entries = np.where(np.isin(self.file["data"]["indices"][:],
                                    indices))[0]
         spec = self.file["data"]["spectra"][entries, :]
         if cast_1d:

cosmopower/examples/1_create_training_data.py

@@ -2,7 +2,7 @@
 
 # Author: A. Spurio Mancini, H. T. Jense
 
-import tqdm
+from tqdm import tqdm
 import numpy as np
 from cosmopower.parser import YAMLParser
 from cosmopower.spectra import init_boltzmann_code, get_boltzmann_spectra
@@ -51,7 +51,8 @@
 samples, validation_samples, testing_samples = \
     parser.get_parameter_samples()
 
-print(samples)
+for k in samples:
+    print(k, samples[k][:10])
 
 parser.save_samples_to_file(samples, validation_samples)
 
@@ -128,7 +129,7 @@
             # This function returns False if the sample is invalid, so it's a
             # simple check to discard invalid datapoints.
             network_params = np.array([
-                params[k]
+                samples[k][n]
                 for k in parser.network_input_parameters("Pk/lin")
             ])
             dataset.write_data(n, network_params,
@@ -147,7 +148,7 @@
             # This function returns False if the sample is invalid, so it's a
             # simple check to discard invalid datapoints.
             network_params = np.array([
-                params[k]
+                samples[k][n]
                 for k in parser.network_input_parameters("Pk/lin")
             ])
             dataset.write_data(n, network_params,
@@ -163,7 +164,7 @@
             # This function returns False if the sample is invalid, so it's a
             # simple check to discard invalid datapoints.
             network_params = np.array([
-                params[k]
+                samples[k][n]
                 for k in parser.network_input_parameters("Pk/lin")
             ])
             dataset.write_data(n, network_params,

cosmopower/examples/2b_train_PCA_emulator.py

@@ -42,7 +42,7 @@
 network = cosmopower_PCAplusNN(cp_pca=pca, verbose=True, trainable=True,
                                **settings.get("n_traits", {}))
 
-with tf.device("/device:CPU:0"):
+with tf.device(None):
     network.train(training_data=datasets,
                   filename_saved_model=output_file,
                   validation=validation,

cosmopower/examples/example.yaml

@@ -4,7 +4,7 @@ path: example
 emulated_code:
   name: camb
   # which parameters are passed to camb
-  inputs: [ombh2, omch2, H0, ns, As]
+  inputs: [ombh2, omch2, H0, ns, As, z]
   extra_args:
     # Some extra accuracy and speed parameters
     kmax: 20.0

cosmopower/examples/planck_cmb.yaml

@@ -40,10 +40,6 @@ networks:
   - quantity: "derived"
     type: NN
     inputs: [ombh2, omch2, cosmomc_theta, ns, logA, tau]
-    log: True
-    modes:
-      label: l
-      range: [2, 2508]
     n_traits:
       n_hidden: [512,512,512,512]
     training:

cosmopower/spectra.py

@@ -59,7 +59,7 @@ def init_boltzmann_code(parser: YAMLParser) -> dict:
     res = {}
 
     code = parser.boltzmann_code
-    extra_args = parser.boltzmann_extra_args
+    extra_args = parser.boltzmann_extra_args or {}
 
     # Try to import the correct python module for spectrum generation
     try:
@@ -256,7 +256,7 @@ def generate_spectra(args: list = None) -> None:
           f"files {first_file}--{last_file}.")
 
     state = init_boltzmann_code(parser)
-    extra_args = parser.boltzmann_extra_args
+    extra_args = parser.boltzmann_extra_args or {}
 
     accepted = 0
     tbar = tqdm.tqdm(np.arange(first, last + 1))
@@ -265,7 +265,7 @@ def generate_spectra(args: list = None) -> None:
 
     for n in tbar:
         tbar.set_description(("" if MPI is None else f"[{rank}] ")
-                             + f"{accepted/max(n,1):.1%} success rate")
+                             + f"{accepted/max(n, 1):.1%} success rate")
 
         boltzmann_params = {k: training[k][n] for k in parser.boltzmann_inputs}
 

cosmopower/spectra_camb.py

@@ -48,13 +48,13 @@ def initialize(parser: YAMLParser, extra_args: dict = {}) -> dict:
 
         # Check that the maximum values for mode ranges matches
         # the requested outputs.
-        if parser.modes_label(quantity) == "l" and \
-           extra_args["lmax"] < parser.modes(quantity).max():
-            extra_args["lmax"] = parser.modes(quantity).max()
+        if parser.modes_label(quantity) == "l":
+            extra_args["lmax"] = max(extra_args["lmax"],
+                                     parser.modes(quantity).max())
 
-        if parser.modes_label(quantity) == "k" and \
-           extra_args["kmax"] < parser.modes(quantity).max():
-            extra_args["kmax"] = parser.modes(quantity).max()
+        if parser.modes_label(quantity) == "k":
+            extra_args["kmax"] = max(extra_args["kmax"],
+                                     parser.modes(quantity).max())
 
         if parser.modes_label(quantity) == "z":
             z1 = extra_args["redshifts"]
@@ -223,12 +223,12 @@ def get_spectra(parser: YAMLParser, state: dict, args: dict = {},
     except (camb.CAMBError, camb.CAMBFortranError):
         return False
 
+    state["derived"] = get_camb_derived(state["params"], state["results"],
+                                        list(state["derived"].keys()))
+
     for quantity in quantities:
         qpath = quantity.split("/")
 
-        state["derived"] = get_camb_derived(state["params"], state["results"],
-                                            list(state["derived"].keys()))
-
         if qpath[0] == "derived":
             continue
         elif qpath[0] == "Cl":
@@ -275,14 +275,14 @@ def get_spectra(parser: YAMLParser, state: dict, args: dict = {},
                 state["params"], state["results"], parser.modes(qpath[0])
             )
         elif qpath[0] == "DA" or qpath[0] == "angular_diameter_distance":
-            # angular diameter distance
-            state[quantity] = get_camb_sigma8(
+            # angular diameter distance DA(z)
+            state[quantity] = get_camb_angular_diameter_distance(
                 state["params"], state["results"], parser.modes(qpath[0])
             )
         elif qpath[0] in ["Omega_b", "Omega_cdm"]:
             # background evolution of densities
             var = {"Omega_b": "baryon", "Omega_cdm": "cdm"}.get(qpath[0])
-            state[quantity] = get_camb_sigma8(
+            state[quantity] = get_camb_Omega(
                 state["params"], state["results"], parser.modes(qpath[0]), var
             )
         else: