At nyx (#105)

* models Nyx alphaP

* add script for plotting

* add script with plotting function

* remove .png from ignore file

* add validation figures

* fix problem with mpg test simulations

* add plot for central

* add plotting of mpg in the AT workflow

* python file with the plotting AT for mpg

* fix bug

* fix bug

* fix bug

* disable second test to debug

* enable second test to debug

* extend timeout

* debugging

* save from workflow

* save from workflow

* same figure, error in upload

* upload artifcat in different path

* naming change

* try new namign for debugging

* testing

* test with tmp diretory

* another test

* test

* another test

* fix bug

* deploy integration

* deploy integration

* deploy integration

* deploy integration

* clear debugging prints

* only CI on branch main

* bug testing simulations

* bug testing simulations-check

* add GP to CI

* indent error

* indent error

* add gpy to installation

* add gpy to installation

* add gpy to installation

* only push to main

---------

Co-authored-by: Laura Cabayol-Garcia <lcabayol@login26.chn.perlmutter.nersc.gov>
Co-authored-by: Laura Cabayol-Garcia <lcabayol@login35.chn.perlmutter.nersc.gov>
Co-authored-by: Laura Cabayol-Garcia <lcabayol@login34.chn.perlmutter.nersc.gov>

.github/workflows/python-tests.yml

@@ -3,8 +3,9 @@ name: LaCE automated testing
 
 on:
   push:
-    branches: '*'
+    branches: 'main'
   workflow_dispatch:
+
 jobs:
   build:
     runs-on: ubuntu-latest
@@ -24,11 +25,26 @@ jobs:
         pip install fdasrsf
         pip install setuptools
         pip install pytest wheel pytest-cov
-    - name: Install lace
-      run: pip install -e .
 
+    - name: Install lace
+      run: pip install -e .[explicit]
 
+    # Run the tests and specify the temporary directory as output_dir
     - name: Run tests
-      timeout-minutes: 5
+      timeout-minutes: 70
       run: |
-        pytest tests/test_lace.py  
+        pytest tests/test_lace.py
+        pytest tests/plot_mpg_gp.py
+        pytest tests/plot_mpg_nn.py
+
+    - name: List generated plots
+      run: |
+        echo "Files in the data/tmp_validation_figures directory:"
+        ls data/tmp_validation_figures/Gadget
+
+    # Archive the generated plots
+    - name: Archive generated plots
+      uses: actions/upload-artifact@v4
+      with:
+        name: validation-figures
+        path: data/tmp_validation_figures/Gadget/*.png    

.gitignore

@@ -1,6 +1,5 @@
 .ipynb_checkpoints
 *.ipynb
-*.png
 lace.egg-info/
 __pycache__
 .github/workflows/.python-tests.yml.swn

lace/utils/plotting_functions.py

@@ -0,0 +1,77 @@
+import os
+import numpy as np
+from matplotlib import pyplot as plt
+from lace.utils import poly_p1d
+import matplotlib.cm as cm
+
+def plot_p1d_vs_emulator(testing_data, emulator, save_path=None):
+    """
+    Plots the true P1D and the emulated P1D for each redshift in the testing data.
+    
+    Parameters:
+    - testing_data: List of dictionaries containing the data for different redshifts.
+    - emulator: An emulator object with a method emulate_p1d_Mpc() for generating P1D.
+    
+    Returns:
+    - None (displays the plot).
+    """
+    # Initialize arrays to store true and emulated P1D values
+    Nk = len(testing_data[0]['p1d_Mpc'][(testing_data[0]['k_Mpc']>0)&(testing_data[0]['k_Mpc']<4)])
+    z = [d['z'] for d in testing_data if d['z']<4.8 ]
+    Nz = len(z)
+
+    p1ds_true = np.zeros(shape=(Nz, Nk))
+    p1ds = np.zeros(shape=(Nz, Nk))
+
+    # Create a colormap with 11 unique colors
+    cmap = cm.get_cmap('viridis', Nz)
+
+    plt.figure()
+
+    # Loop through the entries in the testing data
+    for m in range(Nz):
+        if 'kF_Mpc' not in testing_data[m]:
+            continue
+
+        # Extract relevant data
+        p1d_true = testing_data[m]['p1d_Mpc']
+        kMpc = testing_data[m]['k_Mpc']
+
+        # Filter kMpc values between 0 and 4
+        kMpc_test = kMpc[(kMpc > 0) & (kMpc < 4)]
+        p1d_true = p1d_true[(kMpc > 0) & (kMpc < 4)]
+
+        # Fit the true P1D using a polynomial
+        fit_p1d = poly_p1d.PolyP1D(kMpc_test, p1d_true, kmin_Mpc=1e-3, kmax_Mpc=4, deg=5)
+        p1d_true = fit_p1d.P_Mpc(kMpc_test)
+
+        # Emulate P1D using the provided emulator
+        p1d = emulator.emulate_p1d_Mpc(testing_data[m], kMpc_test)
+
+        # Store the results
+        p1ds_true[m] = p1d_true
+        p1ds[m] = p1d
+
+        # Get the redshift and assign a unique color from the colormap
+        redshift = testing_data[m]['z']
+        color = cmap(m)
+
+        # Plot the true and emulated P1D
+        plt.scatter(kMpc_test, kMpc_test * p1d, label=f'$z={redshift:.1f}$', color=color, marker='^')
+        plt.plot(kMpc_test, kMpc_test * p1d_true, color=color)
+
+    # Set axis labels
+    plt.xlabel(r'$k$ [1/Mpc]')
+    plt.ylabel(r'$k$ * P1D')
+
+    # Add the legend outside the plot for better readability
+    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
+
+    # Save the plot
+    if save_path != None:
+        plt.savefig(save_path, bbox_inches='tight')
+
+    # Show the plot
+    plt.show()
+    plt.close()
+    return

notebooks/Tutorial_emulator.py

@@ -6,11 +6,11 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.1
+#       jupytext_version: 1.15.2
 #   kernelspec:
-#     display_name: Python 3 (ipykernel)
+#     display_name: emulators2
 #     language: python
-#     name: python3
+#     name: emulators2
 # ---
 
 # %% [markdown]

notebooks/plot_archive.py

@@ -36,6 +36,10 @@
     nyx_version = "Oct2023"
     archive = NyxArchive(nyx_version=nyx_version, verbose=True)
 
+for sim in archive.list_sim:
+    if sim in archive.list_sim_test:
+        print(sim)
+
 emu_params=['Delta2_p', 'n_p','mF', 'sigT_Mpc', 'gamma', 'kF_Mpc']
 training_data=archive.get_training_data(emu_params=emu_params)
 

scripts/make_plots_emulator.py

@@ -0,0 +1,133 @@
+# Import necessary modules
+## General python modules
+import numpy as np
+from matplotlib import pyplot as plt
+import matplotlib.cm as cm
+import os
+import argparse  # Used for parsing command-line arguments
+
+## LaCE specific modules
+import lace
+from lace.emulator.nn_emulator import NNEmulator
+from lace.archive import nyx_archive, gadget_archive
+from lace.utils import poly_p1d
+from lace.utils.plotting_functions import plot_p1d_vs_emulator
+
+
+def plot_emulated_p1d(archive_name='Nyx'):
+    """
+    Function to plot emulated P1D using specified archive (Nyx or Gadget).
+    
+    Parameters:
+    archive (str): Archive to use for data ('Nyx' or 'Gadget')
+    """
+
+    # Print the type of archive selected for debugging
+    print(f"Selected archive: {archive_name}")
+
+    # Get the base directory of the lace module
+    repo = os.path.dirname(lace.__path__[0]) + "/"
+    print(f"Repository base path: {repo}")  # Print path for clarity
+
+    # Choose archive and emulator parameters based on the archive type
+    if archive_name == 'Nyx':
+        # Define the parameters for the emulator specific to Nyx
+        nyx_emu_params = ['Delta2_p', 'n_p', 'alpha_p', 'mF', 'sigT_Mpc', 'gamma', 'kF_Mpc']
+        emulator_label='Nyx_alphap'
+        training_set = 'Nyx23_Oct2023'
+        model_path = f'{repo}data/NNmodels/Nyxap_Oct2023/Nyx_alphap_drop_sim'
+
+        # Initialize a NyxArchive instance for postprocessing data
+        archive = nyx_archive.NyxArchive(verbose=True)
+        print("Initialized NyxArchive.")
+
+    elif archive_name == 'Gadget':
+        # Define the parameters for the emulator specific to Gadget
+        nyx_emu_params = ['Delta2_p', 'n_p', 'mF', 'sigT_Mpc', 'gamma', 'kF_Mpc']
+        training_set='Cabayol23'
+        emulator_label='Cabayol23+'
+        model_path = f'{repo}data/NNmodels/Cabayol23+/Cabayol23+_drop_sim'
+
+        # Initialize a GadgetArchive instance for postprocessing data
+        archive = gadget_archive.GadgetArchive(postproc="Cabayol23")
+        print("Initialized GadgetArchive.")
+
+    else:
+        # Raise an error if an invalid archive is provided
+        raise ValueError("archive must be 'Nyx' or 'Gadget'")
+
+    # Iterate over each simulation in the archive
+    skip_sims = ['nyx_14', 'nyx_15', 'nyx_16', 'nyx_17', 'nyx_seed', 'nyx_wdm']
+    for ii, sim in enumerate(archive.list_sim):
+        if sim in skip_sims:
+            continue
+
+        print(f"Processing simulation: {sim}")  # Print simulation being processed
+
+        # Initialize the emulator with the given parameters and model path
+        if (sim=='nyx_central')|(sim=='mpg_central'):
+            if archive_name == 'Gadget':
+                model_path = f'{repo}data/NNmodels/Cabayol23+/Cabayol23+.pt'
+            elif archive_name == 'Nyx':
+                model_path = f'{repo}data/NNmodels/Nyxap_Oct2023/Nyx_alphap.pt'
+
+            emulator = NNEmulator(
+                training_set=training_set,
+                emulator_label=emulator_label,
+                emu_params=nyx_emu_params,
+                model_path=model_path,
+                drop_sim=None,
+                train=False,
+            )
+        else:
+            emulator = NNEmulator(
+                training_set=training_set,
+                emulator_label=emulator_label,
+                emu_params=nyx_emu_params,
+                model_path=model_path+f'_{sim}.pt',
+                drop_sim=sim,
+                train=False,
+            )
+
+        print(f"Initialized emulator for simulation {sim}")
+
+        # Get testing data for the current simulation
+        testing_data = archive.get_testing_data(sim_label=f'{sim}')
+        if sim!='nyx_central':
+            testing_data = [d for d in testing_data if d['val_scaling']==1]
+
+        print(f"Retrieved testing data for simulation {sim}")
+
+        # Plot and save the emulated P1D
+        plot_p1d_vs_emulator(
+            testing_data,
+            emulator,
+            save_path=f'{repo}data/validation_figures/{archive_name}/{sim}.png'
+        )
+        print(f"Saved plot for simulation {sim}")
+
+    print("Plotting complete!")
+    return
+
+
+if __name__ == "__main__":
+    # Parse command-line arguments
+    parser = argparse.ArgumentParser(description="Plot emulated P1D")
+
+    # Add an argument for specifying the archive type ('Nyx' or 'Gadget')
+    parser.add_argument(
+        "--archive", 
+        type=str, 
+        default="Nyx", 
+        choices=["Nyx", "Gadget"], 
+        help="Specify the archive to use, either 'Nyx' or 'Gadget'"
+    )
+
+    # Parse the arguments provided via command line
+    args = parser.parse_args()
+
+    # Print the parsed arguments for clarity
+    print(f"Command-line argument for archive: {args.archive}")
+
+    # Call the function with the parsed argument
+    plot_emulated_p1d(archive_name=args.archive)

tests/plot_mpg_gp.py

@@ -0,0 +1,72 @@
+# Import necessary modules
+## General python modules
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')  # Use a non-interactive backend for rendering plots
+from matplotlib import pyplot as plt
+import matplotlib.cm as cm
+import os
+import argparse  # Used for parsing command-line arguments
+
+## LaCE specific modules
+import lace
+from lace.emulator.gp_emulator import GPEmulator
+from lace.archive import nyx_archive, gadget_archive
+from lace.utils import poly_p1d
+from lace.utils.plotting_functions import plot_p1d_vs_emulator
+
+
+def test():
+    """
+    Function to plot emulated P1D using specified archive (Nyx or Gadget).
+    
+    Parameters:
+    archive (str): Archive to use for data ('Nyx' or 'Gadget')
+    """
+    archive_name = 'Gadget'
+    # Get the base directory of the lace module
+    repo = os.path.dirname(lace.__path__[0]) + "/"
+
+    # Define the parameters for the emulator specific to Gadget
+    emu_params = ['Delta2_p', 'n_p', 'mF', 'sigT_Mpc', 'gamma', 'kF_Mpc']
+    training_set='Pedersen21'
+    emulator_label='Pedersen23'
+
+
+    # Initialize a GadgetArchive instance for postprocessing data
+    archive = gadget_archive.GadgetArchive(postproc="Cabayol23")
+
+    # Directory for saving plots
+    save_dir = f'{repo}data/tmp_validation_figures/{archive_name}/'
+    #save_dir = '{repo}tmp/validation_figures/'
+    # Create the directory if it does not exist
+    os.makedirs(save_dir, exist_ok=True)
+
+    for ii, sim in enumerate(archive.list_sim):
+        if sim in archive.list_sim_test:
+            emulator = GPEmulator(
+                training_set=training_set,
+                emulator_label=emulator_label,
+                emu_params=emu_params,
+            )
+        else:
+            emulator = GPEmulator(
+                training_set=training_set,
+                emulator_label=emulator_label,
+                emu_params=emu_params,
+                drop_sim=sim,
+            )
+
+        # Get testing data for the current simulation
+        testing_data = archive.get_testing_data(sim_label=f'{sim}')
+        if sim != 'nyx_central':
+            testing_data = [d for d in testing_data if d['val_scaling'] == 1]
+
+        # Plot and save the emulated P1D
+        save_path = f'{save_dir}{sim}{emulator_label}.png'
+        plot_p1d_vs_emulator(testing_data, emulator, save_path=save_path)
+
+    return  
+
+# Call the function to execute the test
+test()