Merge pull request #17 from haukekoehn/train_fluxes

Updates on NN architecture, Filters, redshift and tests

.github/workflows/pylint.yml

@@ -7,7 +7,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.10", "3.11", "3.12"]
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
@@ -18,6 +18,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install pylint
+        python -m pip install .
     - name: Analysing the code with pylint
       run: |
-        pylint $(git ls-files '*.py')
+        pylint --fail-under=5 --disable=C,R,W $(git ls-files 'src/*.py')

.github/workflows/unittest.yml

@@ -0,0 +1,40 @@
+name: Unittest
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+    defaults:
+      run:
+        shell: bash
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["ubuntu-latest"]
+        python-version: ["3.10", "3.11", "3.12"]
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python with Conda ${{ matrix.python-version }}
+      uses: conda-incubator/setup-miniconda@v3
+      with:
+          python-version: ${{ matrix.python-version }}
+          architecture: ${{ 'x86_64' }}
+          miniforge-version: latest
+          use-mamba: true
+          mamba-version: "*"
+          activate-environment: fiesta_env
+
+    - name: Install Python dependencies
+      run: |
+        python -m pip install --upgrade pip setuptools wheel
+        python -m pip install pytest pytest-cov pytest-aiohttp sqlparse freezegun PyJWT joblib coveralls
+        python -m pip install .
+
+    - name: Perform tests with pytest
+      run: |
+          python -m coverage run --source fiesta -m pytest tests/*.py

.gitignore

@@ -1 +1,11 @@
 __pycache__
+.coverage
+*egg-info
+/flux_models/*/benchmarks
+/lightcurve_models/*/benchmarks
+/flux_models/*/model/*.h5
+/lightcurve_models/*/model/*.h5
+*nohup.out
+*.sbatch
+*/log
+*/err

benchmarks/GRB/benchmark_afterglowpy_tophat.py

@@ -18,7 +18,6 @@
         file_ending = "Linf"
 
     B = Benchmarker(name = name,
-                parameter_grid = parameter_grid,
                 model_dir = model_dir,
                 MODEL = AfterglowpyLightcurvemodel,
                 filters = FILTERS,

flux_models/afterglowpy_gaussian/benchmark_afterglowpy_gaussian.py

@@ -1,45 +1,35 @@
-import numpy as np 
-import matplotlib.pyplot as plt
-
-from fiesta.train.BenchmarkerFluxes import Benchmarker
-from fiesta.inference.lightcurve_model import AfterglowpyPCA
-from fiesta.utils import Filter
+from fiesta.train.Benchmarker import Benchmarker
+from fiesta.inference.lightcurve_model import AfterglowFlux
 
 
 name = "gaussian"
 model_dir = f"./model/"
 FILTERS = ["radio-3GHz", "radio-6GHz", "bessellv", "X-ray-1keV"]
 
-for metric_name in ["$\\mathcal{L}_2$", "$\\mathcal{L}_\infty$"]:   
-    if metric_name == "$\\mathcal{L}_2$":
-        file_ending = "L2"
-    else:
-        file_ending = "Linf"
-
-    B = Benchmarker(name = name,
-                model_dir = model_dir,
-                MODEL = AfterglowpyPCA,
-                filters = FILTERS,
-                metric_name = metric_name,
-                )
-
-
-    for filt in FILTERS:
-
-        fig, ax = B.plot_lightcurves_mismatch(filter =filt, parameter_labels = ["$\\iota$", "$\log_{10}(E_0)$", "$\\theta_{\\mathrm{c}}$", "$\\alpha_w$", "$\log_{10}(n_{\mathrm{ism}})$", "$p$", "$\\epsilon_E$", "$\\epsilon_B$"])
-        fig.savefig(f"./benchmarks/benchmark_{filt}_{file_ending}.pdf", dpi = 200)
+lc_model = AfterglowFlux(name,
+                         directory = model_dir, 
+                         filters = FILTERS,
+                         model_type= "MLP")
+
+for metric_name in ["L2", "Linf"]:   
+
 
-        B.print_correlations(filter = filt)    
+    benchmarker = Benchmarker(
+                    model = lc_model,
+                    data_file = "./model/afterglowpy_raw_data.h5",
+                    metric_name = metric_name
+                    )
 
-    fig, ax = B.plot_worst_lightcurves()
-    fig.savefig(f"./benchmarks/worst_lightcurves_{file_ending}.pdf", dpi = 200)
-
-
-fig, ax = B.plot_error_distribution()
-fig.savefig(f"./benchmarks/error_distribution.pdf", dpi = 200)
-
-fig, ax = B.plot_error_over_time()
-fig.savefig(f"./benchmarks/error_over_time.pdf", dpi = 200)
+    benchmarker.benchmark()
+
+    benchmarker.plot_lightcurves_mismatch(parameter_labels = ["$\\iota$", 
+                                                              "$\log_{10}(E_0)$", 
+                                                              "$\\theta_{\\mathrm{c}}$",
+                                                              "$\\alpha_{\\mathrm{w}}$",
+                                                              "$\log_{10}(n_{\mathrm{ism}})$", 
+                                                              "$p$", 
+                                                              "$\log_{10}(\\epsilon_e)$", 
+                                                              "$\log_{10}(\\epsilon_B)$"])
 
 
 

flux_models/afterglowpy_gaussian/create_data_afterglowpy_gaussian.py

@@ -1,19 +1,17 @@
 import numpy as np 
-import matplotlib.pyplot as plt
-
 from fiesta.train.AfterglowData import AfterglowpyData
 
 #############
 ### SETUP ###
 #############
 
-tmin = 0.1 # days
+tmin = 1e-4 # days
 tmax = 2000 # days
-n_times = 200
+n_times = 250
 
 
 numin = 1e9 # Hz 
-numax = 2.5e18 # Hz (10 keV)
+numax = 2.5e19 # Hz (100 keV)
 n_nu = 256
 
 
@@ -34,7 +32,7 @@
 jet_conversion = {"tophat": -1,
                   "gaussian": 0}
 
-n_training = 45_000
+n_training = 20_000
 n_val = 0
 n_test = 0
 

flux_models/afterglowpy_gaussian/create_special_data_afterglowpy_gaussian.py

@@ -1,32 +1,24 @@
+from jax.random import PRNGKey
 import numpy as np 
+
 from fiesta.train.AfterglowData import AfterglowpyData
+from fiesta.inference.prior_dict import ConstrainedPrior
+from fiesta.inference.prior import Uniform, Constraint
 
 #############
 ### SETUP ###
 #############
 
 
-parameter_distributions = {
-    'inclination_EM': (0, np.pi/2, "uniform"),
-    'log10_E0': (47, 57, "uniform"),
-    'thetaCore': (0.01, np.pi/5, "loguniform"),
-    'alphaWing': (0.2, 3.5, "uniform"),
-    'log10_n0': (-6, 2, "uniform"),
-    'p': (2.01, 3, "uniform"),
-    'log10_epsilon_e': (-4, 0, "uniform"),
-    'log10_epsilon_B': (-8, 0, "uniform")
-}
-
-
-
-name = "tophat"
+name = "gaussian"
 outdir = f"./model/"
 
 n_training = 0
 n_val = 0
 n_test = 0
 
 n_pool = 24
+
 size = 20_000
 
 
@@ -39,26 +31,28 @@
                           n_test = 0,
                           n_pool = n_pool)
 
-#import h5py
-#with h5py.File(creator.outfile, "r+") as f:
-#    unproblematic = np.unique(np.where(~np.isinf(f["special_train"]["01"]["y"]))[0])
-#
-#    X = f["special_train"]["01"]["X"][unproblematic]
-#    y = f["special_train"]["01"]["y"][unproblematic]
-#    breakpoint()
-#    creator._save_to_file(X, y, group = "special_train", label = "02", comment = "log10_E0 (54, 57)  log10_n0 (-6, -4) thetaCore (0.4, np.pi/5)")
-
-
-
-inclination = np.random.uniform(0, np.pi/2, size = size)
-log10_E0 = np.random.uniform(54, 57, size = size)
-thetaCore = np.random.uniform(0.4, np.pi/5, size= size)
-alphaWing = np.random.uniform(0.2, 3.5, size = size)
-log10_n0 = np.random.uniform(-6, -4, size = size)
-p = np.random.uniform(2, 3, size = size)
-log10_epsilon_e = np.random.uniform(-4, 0, size = size)
-log10_epsilon_B = np.random.uniform(-8, 0, size = size)
-
-X = np.array([inclination, log10_E0, thetaCore, alphaWing, log10_n0, p, log10_epsilon_e, log10_epsilon_B]).T
-
-creator.create_special_data(X, label = "01", comment = "log10_E0 (54, 57)  log10_n0 (-6, -4) thetaCore (0.4, np.pi/5)")
+def conversion_function(sample):
+    converted_sample = sample
+    converted_sample["thetaWing"] = converted_sample["thetaCore"] * converted_sample["alphaWing"]
+    converted_sample["epsilon_tot"] = 10**(converted_sample["log10_epsilon_B"]) + 10**(converted_sample["log10_epsilon_e"]) 
+    return converted_sample
+
+prior = ConstrainedPrior([
+                    Uniform(xmin=0., xmax=np.pi/2, naming=["inclination_EM"]),
+                    Uniform(xmin=54., xmax=57., naming=["log10_E0"]),
+                    Uniform(xmin=0.35, xmax=np.pi/5, naming=["thetaCore"]),
+                    Uniform(0.2, 3.5, naming=["alphaWing"]),
+                    Uniform(xmin=-6.,xmax=-4.,naming=["log10_n0"]),
+                    Uniform(xmin=2., xmax=3., naming=["p"]),
+                    Uniform(xmin=-4., xmax=0., naming=["log10_epsilon_e"]),
+                    Uniform(xmin=-8.,xmax=0., naming=["log10_epsilon_B"]),
+                    Constraint(xmin=0., xmax=1., naming=["epsilon_tot"]),
+                    Constraint(xmin=0., xmax=np.pi/2, naming=["thetaWing"])
+                    ],
+                    conversion_function)
+
+X = prior.sample(PRNGKey(2728), n_samples=size)
+X = [X[p] for p in creator.parameter_names]
+X = np.transpose(X)
+
+creator.create_special_data(X, label = "01", comment = "log10_E0 (54, 57)  log10_n0 (-6, -4) thetaCore (0.35, np.pi/5)")