diff --git a/.github/workflows/delete-branch-cache.yml b/.github/workflows/delete-branch-cache.yml index 55498c922..93ab43a63 100644 --- a/.github/workflows/delete-branch-cache.yml +++ b/.github/workflows/delete-branch-cache.yml @@ -10,18 +10,18 @@ jobs: steps: - name: Check out code uses: actions/checkout@v3 - + - name: Cleanup run: | gh extension install actions/gh-actions-cache - + REPO=${{ github.repository }} BRANCH="refs/pull/${{ github.event.pull_request.number }}/merge" echo "Fetching list of cache key" cacheKeysForPR=$(gh actions-cache list -R $REPO -B $BRANCH -L 100 | cut -f 1 ) - ## Setting this to not fail the workflow while deleting cache keys. + ## Setting this to not fail the workflow while deleting cache keys. set +e echo "Deleting caches..." for cacheKey in $cacheKeysForPR @@ -30,4 +30,4 @@ jobs: done echo "Done" env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} \ No newline at end of file + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/Dockerfile b/Dockerfile index 8e1c8f04c..1346d440a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -37,7 +37,7 @@ RUN apt-get install -y \ cmake \ liblapacke-dev \ liblapack-dev \ - libblas-dev + libblas-dev RUN cd build && cmake .. && make diff --git a/doc/conf.py b/doc/conf.py index 3bc355045..1120206eb 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -5,7 +5,7 @@ import nmma -extensions = ["myst_parser", "sphinx_copybutton","sphinx_github_changelog"] +extensions = ["myst_parser", "sphinx_copybutton", "sphinx_github_changelog"] myst_enable_extensions = [ "amsmath", "colon_fence", @@ -58,7 +58,6 @@ author = "The NMMA Team" - version = nmma.__version__ release = version diff --git a/doc/index.rst b/doc/index.rst index 6b319b7a3..4d7909190 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -138,7 +138,7 @@ Install parallel-bilby: .. code:: - conda install -c conda-forge parallel-bilby + pip install parallel-bilby .. note:: @@ -161,7 +161,7 @@ specifc instructions below) .. code:: - conda install -c conda-forge pymultinest + conda install conda-forge::pymultinest .. warning:: @@ -184,14 +184,14 @@ requirements.txt file which are necessary for NMMA: .. code:: - pip install . + pip install -e . .. note:: There is an issue pip installing ``pyfftw`` on arm64 Mac systems; see the dedicated section below for a solution. If any package appeared to have an issue installing, you can first check by attempting to install it again using pip: .. code:: - + pip install keras, tensorflow pip install importlib_resources pip install extinction pip install dill @@ -201,6 +201,8 @@ requirements.txt file which are necessary for NMMA: pip install sncosmo pip install scikit-learn pip install joblib + pip install keras + pip install tensorflow conda install -c conda-forge p-tqdm .. note:: diff --git a/doc/joint_inference.md b/doc/joint_inference.md index e07f23e16..fdee900b8 100644 --- a/doc/joint_inference.md +++ b/doc/joint_inference.md @@ -1,6 +1,6 @@ ## Multi-messenger inference -A joint inference on gravitational-wave and electromagnetic signals requires NMMA to run on a supercomputer cluster because large memory space are required and need to be shared across many CPU cores. Here, we consider a full joint inference on the binary neutron star merger observed on 17th August 2017. +A joint inference on gravitational-wave and electromagnetic signals requires NMMA to run on a supercomputer cluster because large memory space are required and need to be shared across many CPU cores. Here, we consider a full joint inference on the binary neutron star merger observed on 17th August 2017. In order to run a multi-messenger inference, we need to follow to main steps: @@ -27,11 +27,11 @@ Moreover, a prior on all observed messengers is required and needs to be tailore **Electroamagnetic data and models** -In order to not only sample on gravitational-wave data, we provide further electromagnetic signal related flags. The flag `with-grb=True` will turn on the sampling on a GRB data. As NMMA currently only includes one GRB model, this model does not need to be further specified. If `with-grb=False`, a joint inference of GW+KN data is possible, excluding the GRB part. With regard to the kilonova model, we need to provide a specific model under `kilonova-model`, its respective reduced model grid (if applicable) under `kilonova-model-svd` and a `kilonova-interpolation-type` which can be either `sklearn_gp` or `tensorflow`. The `light-curve-data` flag should include both GRB and kilonova data if a joint inference on GW-GRB-KN is desired (meaning use: `with-grb=True`) or should just include the kilonova data if a GW-KN inference is targeted (meaning use: `with-grb=False`). The kilonova start/end time and time steps apply to both the GRB and kilonova model which will generate light curves during the inference to match the observed data provided. +In order to not only sample on gravitational-wave data, we provide further electromagnetic signal related flags. The flag `with-grb=True` will turn on the sampling on a GRB data. As NMMA currently only includes one GRB model, this model does not need to be further specified. If `with-grb=False`, a joint inference of GW+KN data is possible, excluding the GRB part. With regard to the kilonova model, we need to provide a specific model under `kilonova-model`, its respective reduced model grid (if applicable) under `kilonova-model-svd` and a `kilonova-interpolation-type` which can be either `sklearn_gp` or `tensorflow`. The `light-curve-data` flag should include both GRB and kilonova data if a joint inference on GW-GRB-KN is desired (meaning use: `with-grb=True`) or should just include the kilonova data if a GW-KN inference is targeted (meaning use: `with-grb=False`). The kilonova start/end time and time steps apply to both the GRB and kilonova model which will generate light curves during the inference to match the observed data provided. **Including EOS information** -NMMA enables to include nuclear information by using equations-of-state (EOS) and sample over the EOS during the inference. In order to include a set of EOSs, each EOS.dat file needs to include information on Mass, Radius and Tidal deformability. For the example shown in the `config.ini` file below, we see that `Neos = 5000` meaning that we include 5000 EOS.dat files each containing information on mass, radius and tidal deformability. We also see that a constraint from NICER measurements has been folded in and thus the `eos-weight` reflects this in a weighting. The EOS set should be sorted according to this weighting in order to reduce runtime for the sampling on the EOSs. +NMMA enables to include nuclear information by using equations-of-state (EOS) and sample over the EOS during the inference. In order to include a set of EOSs, each EOS.dat file needs to include information on Mass, Radius and Tidal deformability. For the example shown in the `config.ini` file below, we see that `Neos = 5000` meaning that we include 5000 EOS.dat files each containing information on mass, radius and tidal deformability. We also see that a constraint from NICER measurements has been folded in and thus the `eos-weight` reflects this in a weighting. The EOS set should be sorted according to this weighting in order to reduce runtime for the sampling on the EOSs. **Running the config.ini generation** @@ -40,9 +40,9 @@ In order to prepare the joint inference, a `config.ini` file is required which s ################################################################################ ## Data generation arguments ################################################################################ - + trigger_time = 1187008882.43 - + ################################################################################ ## Detector arguments ################################################################################ @@ -52,15 +52,15 @@ In order to prepare the joint inference, a `config.ini` file is required which s channel_dict = {H1=LOSC-STRAIN, L1=LOSC-STRAIN, V1=LOSC-STRAIN} data_dict = {H1=data/GW170817/H-H1_LOSC_CLN_16_V1-1187007040-2048.gwf, L1=data/GW170817/L-L1_LOSC_CLN_16_V1-1187007040-2048.gwf, V1=data/GW170817/V-V1_LOSC_CLN_16_V1-1187007040-2048.gwf} duration = 128 - + ################################################################################ ## Calibration arguments ################################################################################ - + calibration-model = CubicSpline spline-calibration-nodes = 10 spline-calibration-envelope-dict = {H1:data/GW170817/Feb-20-2018_O2_LHO_GPSTime_1187008882_C02_RelativeResponseUncertainty_FinalResults.txt, L1:data/GW170817/Feb-20-2018_O2_LLO_GPSTime_1187008882_C02_RelativeResponseUncertainty_FinalResults.txt, V1:data/GW170817/V_calibrationUncertaintyEnvelope_magnitude5p1percent_phase40mraddeg20microsecond.txt} - + ################################################################################ ## Job submission arguments ################################################################################ @@ -75,24 +75,24 @@ In order to prepare the joint inference, a `config.ini` file is required which s distance-marginalization=False phase-marginalization=False time-marginalization=False - + ################################################################################ ## Prior arguments ################################################################################ - + prior-file = GW170817_AT2017gfo_GRB170817A.prior - + ################################################################################ ## Waveform arguments ################################################################################ - + frequency-domain-source-model = lal_binary_neutron_star waveform_approximant = IMRPhenomPv2_NRTidalv2 - + ################################################################################ ## EM arguments ################################################################################ - + binary-type=BNS light-curve-data=data/AT2017gfo-GRB170817A/AT2017gfo_GRB170817A.dat kilonova-model=Bu2019lm @@ -107,22 +107,22 @@ In order to prepare the joint inference, a `config.ini` file is required which s kilonova-interpolation-type=sklearn_gp grb-resolution=12 with-grb=True - + ################################################################################ ## EOS arguments ################################################################################ - - with-eos=True + + with-eos=True eos-data=eos/with_NICER_J0740/EOS_024_uniform_5k_sorted Neos=5000 eos-weight=eos/with_NICER_J0740/EOS_sorted_weight.dat The joint inference generation can be performed by running: - + nmma-generation config.ini -This will generate a `GW170817-AT2017gfo-GRB170817A_data_dump.pickle` file under `outdir/data/` which need to be provided for the joint inference function `nmma-analysis`. +This will generate a `GW170817-AT2017gfo-GRB170817A_data_dump.pickle` file under `outdir/data/` which need to be provided for the joint inference function `nmma-analysis`. **Running the analysis** @@ -136,7 +136,7 @@ As detailed above, running the analysis with the command `nmma-analysis --data-d #PBS -o ./outdir/log_data_analysis/out.txt #PBS -m abe #PBS -M - + module load python module load mpt module load mpi4py @@ -144,25 +144,24 @@ As detailed above, running the analysis with the command `nmma-analysis --data-d export MPI_UNBUFFERED_STDIO=true export MPI_LAUNCH_TIMEOUT=240 - + cd $PBS_O_WORKDIR mpirun -np 512 omplace -c 0-127:st=4 nmma-analysis --data-dump /outdir/data/GW170817-AT2017gfo-GRB170817A_data_dump.pickle --nlive 1024 --nact 10 --maxmcmc 10000 --sampling-seed 20210213 --no-plot --outdir -Note that settings might differ from cluster to cluster and also the installation of NMMA might be changed (conda vs. python installation). +Note that settings might differ from cluster to cluster and also the installation of NMMA might be changed (conda vs. python installation). **Maximum mass constraint from a joint analysis** -From a joint posterior of GW and lightcurve data from a BNS, one can derive an upper limit on the TOV mass, if one assumes that the remnant collapsed to a black hole. The idea is to determine the posterior distribution on the remnant's mass from the posterior distribution of the individual neutron star masses $m_1$, $m_2$ and the ejecta and compare this to the TOV mass of EOSs. +From a joint posterior of GW and lightcurve data from a BNS, one can derive an upper limit on the TOV mass, if one assumes that the remnant collapsed to a black hole. The idea is to determine the posterior distribution on the remnant's mass from the posterior distribution of the individual neutron star masses $m_1$, $m_2$ and the ejecta and compare this to the TOV mass of EOSs. + +This can be done via the command -This can be done via the command - maximum-mass-constraint --outdir --joint-posterior --prior --eos-path-macro --eos-path-micro [--use-M-Kepler] -The last flag determines whether the remnant mass is compared against the TOV mass or the maximum mass limit for a rotating NS (Kepler limit). The latter is less conservative. The joint posterior should contain the parameters chirp mass, eta_star, log10_mdisk, log10_mej_dyn as named columns. Here, eta_star is $η* = \ln(0.25-η)$ from the symmetric mass ratio $η$. The macroscopic EOS curves must have the central pressure p0 in MeV/fm³ of each NS mass as last column. +The last flag determines whether the remnant mass is compared against the TOV mass or the maximum mass limit for a rotating NS (Kepler limit). The latter is less conservative. The joint posterior should contain the parameters chirp mass, eta_star, log10_mdisk, log10_mej_dyn as named columns. Here, eta_star is $η* = \ln(0.25-η)$ from the symmetric mass ratio $η$. The macroscopic EOS curves must have the central pressure p0 in MeV/fm³ of each NS mass as last column. -If --use-M-Kepler is set, the prior file needs to contain two additional fiducial paramters for the quasi-universal relations: +If --use-M-Kepler is set, the prior file needs to contain two additional fiducial paramters for the quasi-universal relations: ratio_R = Gaussian(name = "R", mu = 1.255, sigma = 0.024) delta = Uniform(name="delta", minimum = -0.0125, maximum = 0.0125) - diff --git a/doc/lfi_analysis.md b/doc/lfi_analysis.md index ec378a4c8..cb2f0e96c 100644 --- a/doc/lfi_analysis.md +++ b/doc/lfi_analysis.md @@ -1,6 +1,6 @@ # Perform Parameter Estimation Using Liklihood Free Inference (LFI) -NMMA is adding machine learning functionality to its currently offered analysis methods. In this initial incorporation, a neural network approach will perform parameter estimation on light curves from BNS events. We will address the limitations first, and then provide an example run. +NMMA is adding machine learning functionality to its currently offered analysis methods. In this initial incorporation, a neural network approach will perform parameter estimation on light curves from BNS events. We will address the limitations first, and then provide an example run. ## Limitations @@ -21,7 +21,7 @@ In addition to installing the standard requirements.txt file, you must also run ### Generate a simulation set -First, we will create an injection file that describes our light curves. Running the following command line will generate a json file (injection.json). For Ka2017, this will include the parameters: luminosity_distance, timeshift, log10_mej, log10_vej, log10_Xlan, and geocent_time. +First, we will create an injection file that describes our light curves. Running the following command line will generate a json file (injection.json). For Ka2017, this will include the parameters: luminosity_distance, timeshift, log10_mej, log10_vej, log10_Xlan, and geocent_time. nmma_create_injection --prior-file ./priors/Ka2017 --eos-file ./example_files/eos/ALF2.dat --binary-type BNS --filename ./output/injection --n-injection 10 --original-parameters --extension json @@ -29,10 +29,8 @@ We can generate light curves using this injection file with the following comman lightcurve-generation --model Ka2017 --outdir outdir --outfile-type json --label test --tmin -2 --tmax 20 --dt 0.25 --filters ztfg,ztfr,ztfi --injection ./outdir/injection.json --injection-detection-limit 22.0,22.0,22.0 --ztf-uncertainties -There are two options when running the analysis. One can run it using the injection file directly, which will cause the analysis to generate a light curve on the fly. Or, the command will accept a pre-generated light curve. To run using LFI, the --sampler argument must be neuralnet. In the first example, we call the injection file and pass the first injection to the analysis. In the second, we take the first light curve generated by the injection file. Both of these commands will output a posterior plot to the provided outdir. +There are two options when running the analysis. One can run it using the injection file directly, which will cause the analysis to generate a light curve on the fly. Or, the command will accept a pre-generated light curve. To run using LFI, the --sampler argument must be neuralnet. In the first example, we call the injection file and pass the first injection to the analysis. In the second, we take the first light curve generated by the injection file. Both of these commands will output a posterior plot to the provided outdir. lightcurve-analysis --sampler neuralnet --model Ka2017 --outdir inferences --label with_inj --prior priors/Ka2017.prior --injection outdir/injection.json --injection-num 0 --dt 0.25 - lightcurve-analysis --sampler neuralnet --model Ka2017 --outdir inferences --label with_data --data outdir/test_0.json --prior priors/Ka2017.prior - - + lightcurve-analysis --sampler neuralnet --model Ka2017 --outdir inferences --label with_data --data outdir/test_0.json --prior priors/Ka2017.prior diff --git a/doc/observing-scenarios-light-curves.md b/doc/observing-scenarios-light-curves.md index 3121685d1..88fd5ca78 100644 --- a/doc/observing-scenarios-light-curves.md +++ b/doc/observing-scenarios-light-curves.md @@ -7,7 +7,7 @@ This document provides a detailed description of how to use the **injections.dat Running the following command will generate a JSON file (`injection_Bu2019lm.json`) with the BILBY processing of compact binary merger events. Here, we consider binaries of type BNS (Binary Neutron Stars), although NSBH (Neutron Star-Black Hole) is also an option. This injection contains a simulation set of parameters: `luminosity_distance`, `log10_mej_wind`, `KNphi`, `inclination_EM`, `timeshift`, `geocent_time` for the Bu2019lm model. -### Command to Run [nmma-create-injection]: +### Command to Run [nmma-create-injection]: Run this command to create the JSON injection file: @@ -70,4 +70,4 @@ Install `gwpy` in your NMMA environment before splitting the BNS and NSBH event [paper]: https://doi.org/10.3847/1538-4357/acfcb1 [nmma-create-injection]: https://github.com/nuclear-multimessenger-astronomy/nmma/blob/main/nmma/eos/create_injection.py -[lightcurve-analysis]: https://github.com/nuclear-multimessenger-astronomy/nmma/blob/main/nmma/em/analysis.py \ No newline at end of file +[lightcurve-analysis]: https://github.com/nuclear-multimessenger-astronomy/nmma/blob/main/nmma/em/analysis.py diff --git a/doc/systematics.rst b/doc/systematics.rst index 249bbe9b7..e2765291d 100644 --- a/doc/systematics.rst +++ b/doc/systematics.rst @@ -78,7 +78,7 @@ In this configuration, the ``sdssu`` and ``ztfr`` filters are sampled together f type: Uniform minimum: 0 maximum: 2 - + Distribution types ================== diff --git a/example_files/lightcurves/AT2017gfo_reduced.dat b/example_files/lightcurves/AT2017gfo_reduced.dat index 9fda7c3c0..e3310ce14 100644 --- a/example_files/lightcurves/AT2017gfo_reduced.dat +++ b/example_files/lightcurves/AT2017gfo_reduced.dat @@ -105,4 +105,4 @@ 2017-08-27T00:00:00.000 ps1::y 20.78000 0.11000 2017-08-27T00:00:00.000 2massj 20.23000 0.10000 2017-08-27T00:00:00.000 2massh 19.66000 0.14000 -2017-08-27T00:00:00.000 2massks 18.50000 0.20000 \ No newline at end of file +2017-08-27T00:00:00.000 2massks 18.50000 0.20000 diff --git a/example_files/multi_config_analysis/config.yaml b/example_files/multi_config_analysis/config.yaml index 6e2879fdf..79af8b1b2 100644 --- a/example_files/multi_config_analysis/config.yaml +++ b/example_files/multi_config_analysis/config.yaml @@ -72,4 +72,4 @@ Me2017_4: generation-seed: 42 filters: ztfg,ztfr,ztfi remove-nondetections: True - #process-per-config: 7 \ No newline at end of file + #process-per-config: 7 diff --git a/example_files/multi_config_analysis/injection.json b/example_files/multi_config_analysis/injection.json index ac74590d4..fd61af7b6 100644 --- a/example_files/multi_config_analysis/injection.json +++ b/example_files/multi_config_analysis/injection.json @@ -28,4 +28,4 @@ ] } } -} \ No newline at end of file +} diff --git a/example_files/tools/gwem_resampling/gwsamples_generation.py b/example_files/tools/gwem_resampling/gwsamples_generation.py index 1e46951d2..37e2dc32f 100644 --- a/example_files/tools/gwem_resampling/gwsamples_generation.py +++ b/example_files/tools/gwem_resampling/gwsamples_generation.py @@ -3,15 +3,15 @@ import bilby # load posterior file -eos_post = np.loadtxt('/example_files/tools/gwem_resampling/posterior_probability.txt') +eos_post = np.loadtxt("/example_files/tools/gwem_resampling/posterior_probability.txt") -npts = 150000 +npts = 150000 Neos = 5000 nparams = 3 ############# [mass1, mass2, DL] -params_low = [1.001398, 1.001398, 1] -params_high = [2.2, 2.2, 75] +params_low = [1.001398, 1.001398, 1] +params_high = [2.2, 2.2, 75] # 1) create dummy EOS samples with eos_post from nature paper EOS_raw = np.arange(0, Neos) # the gwem_resampling will add one to this @@ -20,13 +20,22 @@ # 2) generate samples for masses and distance mass_1 = np.random.uniform(params_low[0], params_high[0], size=npts) mass_2 = np.random.uniform(params_low[1], params_high[1], size=npts) -mass_1, mass_2 = np.maximum(mass_1, mass_2), np.minimum(mass_1, mass_2) +mass_1, mass_2 = np.maximum(mass_1, mass_2), np.minimum(mass_1, mass_2) mass_ratio = mass_2 / mass_1 # mass ratio q < 1 convention is used chirp_mass = bilby.gw.conversion.component_masses_to_chirp_mass(mass_1, mass_2) lum_distance = np.random.uniform(params_low[2], params_high[2], size=npts) # 3) create pandas dataframe -dataset = pd.DataFrame({'mass_1': mass_1, 'mass_2': mass_2, 'chirp_mass': chirp_mass, 'mass_ratio': mass_ratio, 'luminosity_distance': lum_distance, 'EOS': EOS_samples}) +dataset = pd.DataFrame( + { + "mass_1": mass_1, + "mass_2": mass_2, + "chirp_mass": chirp_mass, + "mass_ratio": mass_ratio, + "luminosity_distance": lum_distance, + "EOS": EOS_samples, + } +) # 4) save GWsamples.dat file -dataset.to_csv('GWsamples_new.dat', index=False, sep=' ') +dataset.to_csv("GWsamples_new.dat", index=False, sep=" ") diff --git a/nmma/em/analysis.py b/nmma/em/analysis.py index 192834d6a..0c31fad62 100644 --- a/nmma/em/analysis.py +++ b/nmma/em/analysis.py @@ -14,15 +14,15 @@ import pandas as pd from astropy import time from bilby.core.likelihood import ZeroLikelihood -import matplotlib.pyplot as plt from ..utils.models import refresh_models_list from .injection import create_light_curve_data from .likelihood import OpticalLightCurve from .model import create_light_curve_model_from_args, model_parameters_dict from .prior import create_prior_from_args -from .utils import getFilteredMag, dataProcess -from .io import loadEvent, detection_limit_from_m4opt_fits_file +from .utils import getFilteredMag, dataProcess, set_mission_name + +from .io import loadEvent matplotlib.use("agg") @@ -90,7 +90,10 @@ def get_parser(**kwargs): "--dt", type=float, default=0.1, help="Time step in day (default: 0.1)" ) parser.add_argument( - "--dt-inj", type=float, default=1, help="Time step in day for injection (default: 1.0)" + "--dt-inj", + type=float, + default=1, + help="Time step in day for injection (default: 1.0)", ) parser.add_argument( "--log-space-time", @@ -280,7 +283,7 @@ def get_parser(**kwargs): "--rubin-ToO-type", help="Type of ToO observation. Won't work w/o --rubin-ToO", type=str, - choices=["platinum","gold","gold_z","silver","silver_z"], + choices=["platinum", "gold", "gold_z", "silver", "silver_z"], ) parser.add_argument( "--xlim", @@ -394,11 +397,16 @@ def get_parser(**kwargs): parser.add_argument( "--fits-file", - help="Fits file output from Bayestar, to be used for constructing dL-iota prior" + help="Fits file output from Bayestar, to be used for constructing dL-iota prior", ) + # FIXME: to be remove latter parser.add_argument( "--detection-limit-fits-file", - help="Fits file output from m4opt which contain the detection limit of a given sky location" + help="Fits file output from m4opt which contain the detection limit of a given sky location", + ) + parser.add_argument( + "--mission-name", + help="The Telescope name. When running with the detection limit from M4OPT, this will be caught in utils.py to read the appropriate bandpass (UVEX or ULTRASAT)", ) parser.add_argument( "--cosiota-node-num", @@ -415,17 +423,17 @@ def get_parser(**kwargs): parser.add_argument( "--ra", type=float, - help="Right ascension of the sky location; to be used together with fits file" + help="Right ascension of the sky location; to be used together with fits file", ) parser.add_argument( "--dec", type=float, - help="Declination of the sky location; to be used together with fits file" + help="Declination of the sky location; to be used together with fits file", ) parser.add_argument( "--dL", type=float, - help="Distance of the location; to be used together with fits file" + help="Distance of the location; to be used together with fits file", ) parser.add_argument( "--fetch-Ebv-from-dustmap", @@ -445,6 +453,14 @@ def get_parser(**kwargs): def analysis(args): + # Global mission name storage for M4OPT mission of specific filters + # When a mission name is provided (ULTRASAT or UVEX), this global variable + # stores it so that get_default_filts_lambdas() can dynamically add the + # appropriate bandpass filters (NUV for ULTRASAT, FUV+NUV for UVEX). + # Set once at the start of analysis() via set_mission_name(), + # then accessed throughout the code via get_mission_name(). + set_mission_name(args.mission_name) + if args.sampler == "pymultinest": if len(args.outdir) > 64: print( @@ -470,20 +486,19 @@ def analysis(args): filters = filters.split(",") if len(filters) == 0: raise ValueError("Need at least one valid filter.") - elif args.rubin_ToO_type == 'platinum': - filters = ["ps1__g","ps1__r","ps1__i","ps1__z","ps1__y"] - elif args.rubin_ToO_type == 'gold': - filters = ["ps1__g","ps1__r","ps1__i"] - elif args.rubin_ToO_type == 'gold_z': - filters = ["ps1__g","ps1__r","ps1__z"] - elif args.rubin_ToO_type == 'silver': - filters = ["ps1__g""ps1__i"] - elif args.rubin_ToO_type == 'silver_z': - filters = ["ps1__g","ps1__z"] + elif args.rubin_ToO_type == "platinum": + filters = ["ps1__g", "ps1__r", "ps1__i", "ps1__z", "ps1__y"] + elif args.rubin_ToO_type == "gold": + filters = ["ps1__g", "ps1__r", "ps1__i"] + elif args.rubin_ToO_type == "gold_z": + filters = ["ps1__g", "ps1__r", "ps1__z"] + elif args.rubin_ToO_type == "silver": + filters = ["ps1__g" "ps1__i"] + elif args.rubin_ToO_type == "silver_z": + filters = ["ps1__g", "ps1__z"] else: filters = None - # initialize light curve model timeshift = 0 if args.log_space_time: @@ -496,7 +511,7 @@ def analysis(args): ) else: sample_times = np.arange(args.tmin, args.tmax + args.dt, args.dt) - + print("Creating light curve model for inference") if args.filters: @@ -514,7 +529,7 @@ def analysis(args): f, object_hook=bilby.core.utils.decode_bilby_json ) injection_df = injection_dict["injections"] - row = injection_df.loc[injection_df['simulation_id'] == args.injection_num] + row = injection_df.loc[injection_df["simulation_id"] == args.injection_num] injection_parameters = row.squeeze().to_dict() if "geocent_time" in injection_parameters: @@ -524,10 +539,10 @@ def analysis(args): else: print("Need either geocent_time or geocent_time_x") exit(1) - - timeshift = 0 + + timeshift = 0 trigger_time = tc_gps.mjd + timeshift - + if args.ignore_timeshift: if "timeshift" in injection_parameters: timeshift = injection_parameters["timeshift"] @@ -535,9 +550,9 @@ def analysis(args): timeshift = 0 trigger_time = tc_gps.mjd + timeshift - - #print("the trigger time from the injection file is: ", trigger_time) - + + # print("the trigger time from the injection file is: ", trigger_time) + # initialize light curve model if args.log_space_time: if args.n_tstep: @@ -548,10 +563,12 @@ def analysis(args): np.log10(args.tmin), np.log10(args.tmax + args.dt), n_step ) else: - sample_times = np.arange(args.tmin + timeshift, args.tmax + timeshift + args.dt, args.dt) - + sample_times = np.arange( + args.tmin + timeshift, args.tmax + timeshift + args.dt, args.dt + ) + print("Creating light curve model for inference") - + injection_parameters["kilonova_trigger_time"] = trigger_time if args.prompt_collapse: injection_parameters["log10_mej_wind"] = -3.0 @@ -592,41 +609,49 @@ def analysis(args): ) print(f"Injection generated with parameters {injection_parameters}") - #checking produced data for magnitudes dimmer than the detection limit + # checking produced data for magnitudes dimmer than the detection limit if filters is not None: if args.detection_limit is None: if args.rubin_ToO_type: - detection_limit = {'ps1__g':25.8,'ps1__r':25.5,'ps1__i':24.8,'ps1__z':24.1,'ps1__y':22.9} - #elif args.ztf_sampling: - # detection_limit = {} - elif args.detection_limit_fits_file is not None: - limit_given_radec = detection_limit_from_m4opt_fits_file( - args.detection_limit_fits_file, args.ra, args.dec - ) - print(f"Detection limit from {args.detection_limit_fits_file} is used") - print(f"Given ra:{args.ra} and dec:{args.dec}, the limiting mag is {limit_given_radec}") detection_limit = { - x: float(limit_given_radec) - for x in filters + "ps1__g": 25.8, + "ps1__r": 25.5, + "ps1__i": 24.8, + "ps1__z": 24.1, + "ps1__y": 22.9, } + # elif args.ztf_sampling: + # detection_limit = {} + + # FIXME : to be remove we could just provide the detection limit from M4OPT instead of recalculat them + # elif args.detection_limit_fits_file is not None: + # limit_given_radec = detection_limit_from_m4opt_fits_file( + # args.detection_limit_fits_file, args.ra, args.dec + # ) + # print(f"Detection limit from {args.detection_limit_fits_file} is used") + # print(f"Given ra:{args.ra} and dec:{args.dec}, the limiting mag is {limit_given_radec}") + # detection_limit = { + # x: float(limit_given_radec) + # for x in filters + # } else: detection_limit = {x: np.inf for x in filters} else: detection_limit = literal_eval(args.detection_limit) - #print("the detection limits for this run are: ", detection_limit) + # print("the detection limits for this run are: ", detection_limit) for filt in filters: - i=0 + i = 0 for row in data[filt]: - #print('the old data is {data}'.format(data=data[filt])) + # print('the old data is {data}'.format(data=data[filt])) mjd, mag, mag_unc = row - #print("the data for {f} is: ".format(f=filt), row) + # print("the data for {f} is: ".format(f=filt), row) if mag > detection_limit[filt]: - data[filt][i,:] = [mjd, detection_limit[filt], -np.inf] - - #print("the new data is: ", data[filt]) - i+=1 + data[filt][i, :] = [mjd, detection_limit[filt], -np.inf] + + # print("the new data is: ", data[filt]) + i += 1 if args.injection_outfile is not None: if filters is not None: @@ -724,7 +749,7 @@ def analysis(args): detection = False notallnan = False - #print('data before checking for detections: ', data[filt]) + # print('data before checking for detections: ', data[filt]) for filt in data.keys(): idx = np.where(np.isfinite(data[filt][:, 2]))[0] @@ -786,7 +811,7 @@ def analysis(args): # setup the prior priors = create_prior_from_args(model_names, args) - #print('the data passed to likelihood is: ', data) + # print('the data passed to likelihood is: ', data) # setup the likelihood if args.detection_limit: @@ -796,12 +821,12 @@ def analysis(args): filters=filters_to_analyze, light_curve_data=data, trigger_time=trigger_time, - tmin=args.tmin+timeshift, - tmax=args.tmax+timeshift, + tmin=args.tmin + timeshift, + tmax=args.tmax + timeshift, error_budget=error_budget, verbose=args.verbose, detection_limit=args.detection_limit, - systematics_file=args.systematics_file + systematics_file=args.systematics_file, ) likelihood = OpticalLightCurve(**likelihood_kwargs) @@ -832,7 +857,7 @@ def analysis(args): elif args.sampler == "dynesty": sampler_kwargs["maxiter"] = 1 - #print("passing arguments to bilby") + # print("passing arguments to bilby") result = bilby.run_sampler( likelihood, @@ -938,7 +963,11 @@ def analysis(args): # calculate the chi2 # ###################### processed_data = dataProcess( - data, filters_to_analyze, trigger_time, args.tmin+timeshift, args.tmax+timeshift + data, + filters_to_analyze, + trigger_time, + args.tmin + timeshift, + args.tmax + timeshift, ) chi2 = 0.0 dof = 0.0 @@ -954,12 +983,18 @@ def analysis(args): print("the time values before adding timeshift are: ", t) # shift t values by timeshift if "timeshift" in bestfit_params: - print("timeshift found in bestfit_params is: ",bestfit_params["timeshift"]) + print( + "timeshift found in bestfit_params is: ", + bestfit_params["timeshift"], + ) t += bestfit_params["timeshift"] # only the detection data are needed finite_idx = np.where(np.isfinite(sigma_y))[0] print("the {f} data being analyzed is: ".format(f=filt), samples) - print("for {f} the length of the detections array is: ".format(f=filt), len(finite_idx)) + print( + "for {f} the length of the detections array is: ".format(f=filt), + len(finite_idx), + ) if len(finite_idx) > 0: # fetch the erorr_budget if "em_syserr" in bestfit_params: @@ -1044,25 +1079,29 @@ def analysis(args): bspace = 0.7 trspace = 0.2 hpanel = 2.25 - wpanel = 3. + wpanel = 3.0 ncol = 2 nrow = int(np.ceil(len(filters_plot) / ncol)) fig, axes = plt.subplots(nrow, ncol) - figsize = (1.5 * (lspace + wpanel * ncol + wspace * (ncol - 1) + trspace), - 1.5 * (bspace + hpanel * nrow + hspace * (nrow - 1) + trspace)) + figsize = ( + 1.5 * (lspace + wpanel * ncol + wspace * (ncol - 1) + trspace), + 1.5 * (bspace + hpanel * nrow + hspace * (nrow - 1) + trspace), + ) # Create the figure and axes. fig, axes = plt.subplots(nrow, ncol, figsize=figsize, squeeze=False) - fig.subplots_adjust(left=lspace / figsize[0], - bottom=bspace / figsize[1], - right=1. - trspace / figsize[0], - top=1. - trspace / figsize[1], - wspace=wspace / wpanel, - hspace=hspace / hpanel) + fig.subplots_adjust( + left=lspace / figsize[0], + bottom=bspace / figsize[1], + right=1.0 - trspace / figsize[0], + top=1.0 - trspace / figsize[1], + wspace=wspace / wpanel, + hspace=hspace / hpanel, + ) if len(filters_plot) % 2: - axes[-1, -1].axis('off') + axes[-1, -1].axis("off") cnt = 0 for filt, color in zip(filters_plot, colors): @@ -1074,9 +1113,7 @@ def analysis(args): ax_sum = axes[row, col] # adding the ax for the Delta divider = make_axes_locatable(ax_sum) - ax_delta = divider.append_axes('bottom', - size='30%', - sharex=ax_sum) + ax_delta = divider.append_axes("bottom", size="30%", sharex=ax_sum) # configuring ax_sum ax_sum.set_ylabel("AB magnitude", rotation=90) @@ -1110,16 +1147,13 @@ def analysis(args): marker="v", color=color, ) - + mag_plot = getFilteredMag(mag, filt) # calculating the chi2 - mag_per_data = np.interp( - t[det_idx], - mag["bestfit_sample_times"], - mag_plot) + mag_per_data = np.interp(t[det_idx], mag["bestfit_sample_times"], mag_plot) diff_per_data = mag_per_data - y[det_idx] - sigma_per_data = np.sqrt((sigma_y[det_idx]**2 + error_budget[filt]**2)) + sigma_per_data = np.sqrt((sigma_y[det_idx] ** 2 + error_budget[filt] ** 2)) chi2_per_data = diff_per_data**2 chi2_per_data /= sigma_per_data**2 chi2_total = np.sum(chi2_per_data) @@ -1127,12 +1161,12 @@ def analysis(args): # plot the mismatch between the model and the data ax_delta.scatter(t[det_idx], diff_per_data / sigma_per_data, color=color) - ax_delta.axhline(0, linestyle='--', color='k') + ax_delta.axhline(0, linestyle="--", color="k") ax_sum.plot( mag["bestfit_sample_times"], mag_plot, - color='coral', + color="coral", linewidth=3, linestyle="--", ) @@ -1142,7 +1176,7 @@ def analysis(args): mag["bestfit_sample_times"], mag_plot + error_budget[filt], mag_plot - error_budget[filt], - facecolor='coral', + facecolor="coral", alpha=0.2, label="combined", ) @@ -1151,7 +1185,7 @@ def analysis(args): mag["bestfit_sample_times"], mag_plot + error_budget[filt], mag_plot - error_budget[filt], - facecolor='coral', + facecolor="coral", alpha=0.2, ) @@ -1161,7 +1195,7 @@ def analysis(args): ax_sum.plot( mag["bestfit_sample_times"], mag_plot, - color='coral', + color="coral", linewidth=3, linestyle="--", ) @@ -1174,36 +1208,55 @@ def analysis(args): label=models[ii].model, ) - ax_sum.set_title(f'{filt}: ' + fr'$\chi^2 / d.o.f. = {round(chi2_total / N_data, 2)}$') + ax_sum.set_title( + f"{filt}: " + rf"$\chi^2 / d.o.f. = {round(chi2_total / N_data, 2)}$" + ) ax_sum.set_xlim([float(x) for x in args.xlim.split(",")]) ax_sum.set_ylim([float(x) for x in args.ylim.split(",")]) ax_delta.set_xlim([float(x) for x in args.xlim.split(",")]) - plt.savefig(plotName, bbox_inches='tight') + plt.savefig(plotName, bbox_inches="tight") plt.close() + def nnanalysis(args): + set_mission_name(args.mission_name) + # import functions - from ..mlmodel.dataprocessing import gen_prepend_filler, gen_append_filler, pad_the_data - from ..mlmodel.resnet import ResNet + from ..mlmodel.dataprocessing import ( + # gen_prepend_filler, # noqa: F401 + # gen_append_filler, # noqa: F401 + pad_the_data, + ) + + # from ..mlmodel.resnet import ResNet # noqa: F401 from ..mlmodel.embedding import SimilarityEmbedding from ..mlmodel.normalizingflows import normflow_params from ..mlmodel.inference import cast_as_bilby_result # need to add these packages: import torch - import torch.nn as nn - from torch.utils.data import Dataset, DataLoader, TensorDataset, random_split - import torch.nn.functional as F - from nflows.nn.nets.resnet import ResidualNet - from nflows import transforms, distributions, flows - from nflows.distributions import StandardNormal + + # import torch.nn as nn # noqa: F401 + # from torch.utils.data import ( + # Dataset, # noqa: F401 + # DataLoader, # noqa: F401 + # TensorDataset, # noqa: F401 + # random_split, # noqa: F401 + # ) + # import torch.nn.functional as F # noqa: F401 + # from nflows.nn.nets.resnet import ResidualNet # noqa: F401 + # from nflows import transforms, distributions, flows # noqa: F401 + # from nflows.distributions import StandardNormal # noqa: F401 from nflows.flows import Flow - from nflows.transforms.autoregressive import MaskedAffineAutoregressiveTransform - from nflows.transforms import CompositeTransform, RandomPermutation - import nflows.utils as torchutils + + # from nflows.transforms.autoregressive import ( + # MaskedAffineAutoregressiveTransform, # noqa: F401 + # ) + # from nflows.transforms import CompositeTransform, RandomPermutation # noqa: F401 + # import nflows.utils as torchutils # noqa: F401 # only continue if the Kasen model is selected if args.model != "Ka2017": @@ -1211,22 +1264,24 @@ def nnanalysis(args): "WARNING: model selected is not currently compatible with this inference method" ) exit() - else: + else: pass - print('Starting LFI') + print("Starting LFI") # only can use ztfr, ztfg, and ztfi filters in the light curve data - if args.filters: + if args.filters: filters = args.filters.replace(" ", "") # remove all whitespace filters = filters.split(",") - if ('ztfr' in filters) and ('ztfi' in filters) and ('ztfg' in filters): + if ("ztfr" in filters) and ("ztfi" in filters) and ("ztfg" in filters): pass else: raise ValueError("Need the ztfr, ztfi, and ztfg filters.") - else: - print('Currently filters are hardcoded to ztfr, ztfi, and ztfg. Continuing with these filters.') - filters = 'ztfg,ztfi,ztfr' + else: + print( + "Currently filters are hardcoded to ztfr, ztfi, and ztfg. Continuing with these filters." + ) + filters = "ztfg,ztfi,ztfr" filters = filters.replace(" ", "") # remove all whitespace filters = filters.split(",") @@ -1244,7 +1299,7 @@ def nnanalysis(args): bilby.core.utils.setup_logger(outdir=args.outdir, label=args.label) bilby.core.utils.check_directory_exists_and_if_not_mkdir(args.outdir) - print('Setting up logger and storage directory') + print("Setting up logger and storage directory") if args.log_space_time: if args.n_tstep: @@ -1293,18 +1348,20 @@ def nnanalysis(args): if args.dt: time_step = args.dt if args.dt != 0.25: - raise ValueError("Need dt to be 0.25 until interpolation feature is incorporated.") + raise ValueError( + "Need dt to be 0.25 until interpolation feature is incorporated." + ) # currently no linear interpolation function - do_lin_interpolation = True + do_lin_interpolation = True # noqa: F841 else: - do_lin_interpolation = False + do_lin_interpolation = False # noqa: F841 args.kilonova_tmin = args.tmin args.kilonova_tmax = args.tmax args.kilonova_tstep = args.dt_inj args.kilonova_error = args.photometric_error_budget - current_points = int(round(args.tmax - args.tmin))/args.dt + 1 + # current_points = int(round(args.tmax - args.tmin)) / args.dt + 1 if not args.injection_model: args.kilonova_injection_model = args.model @@ -1338,16 +1395,17 @@ def nnanalysis(args): args.injection_detection_limit.split(","), ) } - elif args.detection_limit_fits_file is not None: - limit_given_radec = detection_limit_from_m4opt_fits_file( - args.detection_limit_fits_file, args.ra, args.dec - ) - print(f"Detection limit from {args.detection_limit_fits_file} is used") - print(f"Given ra:{args.ra} and dec:{args.dec}, the limiting mag is {limit_given_radec}") - detection_limit = { - x: float(limit_given_radec) - for x in filters - } + # FIXME : to be remove + # elif args.detection_limit_fits_file is not None: + # limit_given_radec = detection_limit_from_m4opt_fits_file( + # args.detection_limit_fits_file, args.ra, args.dec + # ) + # print(f"Detection limit from {args.detection_limit_fits_file} is used") + # print(f"Given ra:{args.ra} and dec:{args.dec}, the limiting mag is {limit_given_radec}") + # detection_limit = { + # x: float(limit_given_radec) + # for x in filters + # } else: detection_limit = {x: np.inf for x in filters} else: @@ -1409,7 +1467,7 @@ def nnanalysis(args): # load the lightcurve data data = loadEvent(args.data) res = next(iter(data)) - current_points = len(data[res]) + # current_points = len(data[res]) if args.trigger_time is None: # load the minimum time as trigger time @@ -1493,17 +1551,21 @@ def nnanalysis(args): # setup the prior priors = create_prior_from_args(model_names, args) - + # now that we have the kilonova light curve, we need to pad it with non-detections # this part is currently hard coded in terms of the times !!!! likely will need the most work # (so that the 'fixed' and 'shifted' are properly represented) num_points = 121 num_channels = 3 - bands = ['ztfg', 'ztfr', 'ztfi'] # will need to edit to not be hardcoded - t_zero = 44242.00021937881 - t_min = 44240.00050450478 - t_max = 44269.99958898723 - days = int(round(t_max - t_min)) + # bands = [ + # "ztfg", + # "ztfr", + # "ztfi", + # ] # noqa: F841 # will need to edit to not be hardcoded + # t_zero = 44242.00021937881 # noqa: F841 + # t_min = 44240.00050450478 + # t_max = 44269.99958898723 + # days = int(round(t_max - t_min)) # noqa: F841 time_step = 0.25 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -1511,14 +1573,14 @@ def nnanalysis(args): if args.detection_limit: detection_limit = args.detection_limit - else: + else: detection_limit = 22.0 data_df = pd.DataFrame() t_list = [] for i in range(len(data[res])): t_list.append(data[res][i][0]) - data_df['t'] = t_list + data_df["t"] = t_list for key in data: mag_list = [] for i, val in enumerate(data[key]): @@ -1527,52 +1589,78 @@ def nnanalysis(args): data_df[key] = mag_list column_list = data_df.columns.to_list() - # pad the data + # pad the data padded_data_df = pad_the_data( - data_df, + data_df, column_list, - desired_count=num_points, - filler_time_step=time_step, - filler_data=detection_limit + desired_count=num_points, + filler_time_step=time_step, + filler_data=detection_limit, ) # change the data into pytorch tensors - data_tensor = torch.tensor(padded_data_df.iloc[:, 1:4].values.reshape(1, num_points, num_channels), dtype=torch.float32).transpose(1, 2) - - # set up the embedding - similarity_embedding = SimilarityEmbedding(num_dim=7, num_hidden_layers_f=1, num_hidden_layers_h=1, num_blocks=4, kernel_size=5, num_dim_final=5).to(device) + data_tensor = torch.tensor( + padded_data_df.iloc[:, 1:4].values.reshape(1, num_points, num_channels), + dtype=torch.float32, + ).transpose(1, 2) + + # set up the embedding + similarity_embedding = SimilarityEmbedding( + num_dim=7, + num_hidden_layers_f=1, + num_hidden_layers_h=1, + num_blocks=4, + kernel_size=5, + num_dim_final=5, + ).to(device) num_dim = 7 - SAVEPATH = os.getcwd() + '/nmma/mlmodel/similarity_embedding_weights.pth' + SAVEPATH = os.getcwd() + "/nmma/mlmodel/similarity_embedding_weights.pth" similarity_embedding.load_state_dict(torch.load(SAVEPATH, map_location=device)) for name, param in similarity_embedding.named_parameters(): param.requires_grad = False # set up the normalizing flows - transform, base_dist, embedding_net = normflow_params(similarity_embedding, 9, 5, 90, context_features=num_dim, num_dim=num_dim) + transform, base_dist, embedding_net = normflow_params( + similarity_embedding, 9, 5, 90, context_features=num_dim, num_dim=num_dim + ) flow = Flow(transform, base_dist, embedding_net).to(device=device) - PATH_nflow = os.getcwd() + '/nmma/mlmodel/frozen-flow-weights.pth' + PATH_nflow = os.getcwd() + "/nmma/mlmodel/frozen-flow-weights.pth" flow.load_state_dict(torch.load(PATH_nflow, map_location=device)) nsamples = 20000 with torch.no_grad(): samples = flow.sample(nsamples, context=data_tensor) - samples = samples.cpu().reshape(nsamples,3) - + samples = samples.cpu().reshape(nsamples, 3) + if args.injection: avail_parameters = injection_parameters.keys() - if ('log10_mej' in avail_parameters) and ('log10_vej' in avail_parameters) and ('log10_Xlan' in avail_parameters): - param_tensor = torch.tensor([injection_parameters['log10_mej'], injection_parameters['log10_vej'], injection_parameters['log10_Xlan']], dtype=torch.float32) + if ( + ("log10_mej" in avail_parameters) + and ("log10_vej" in avail_parameters) + and ("log10_Xlan" in avail_parameters) + ): + param_tensor = torch.tensor( + [ + injection_parameters["log10_mej"], + injection_parameters["log10_vej"], + injection_parameters["log10_Xlan"], + ], + dtype=torch.float32, + ) with torch.no_grad(): - truth = param_tensor + truth = param_tensor flow_result = cast_as_bilby_result(samples, truth, priors=priors) - fig = flow_result.plot_corner(save=True, label = args.label, outdir=args.outdir) - print('saved posterior plot') + flow_result.plot_corner(save=True, label=args.label, outdir=args.outdir) + print("saved posterior plot") else: - raise ValueError('The injection parameters provided do not match the parameters the flow has been trained on') + raise ValueError( + "The injection parameters provided do not match the parameters the flow has been trained on" + ) else: flow_result = cast_as_bilby_result(samples, truth=None, priors=priors) - fig = flow_result.plot_corner(save=True, label = args.label, outdir=args.outdir) - print('saved posterior plot') + flow_result.plot_corner(save=True, label=args.label, outdir=args.outdir) + print("saved posterior plot") + def main(args=None): if args is None: @@ -1592,4 +1680,3 @@ def main(args=None): nnanalysis(args) else: analysis(args) - diff --git a/nmma/em/analysis_lbol.py b/nmma/em/analysis_lbol.py index fee93265a..c8d341aaa 100644 --- a/nmma/em/analysis_lbol.py +++ b/nmma/em/analysis_lbol.py @@ -21,7 +21,7 @@ from .utils import running_in_ci matplotlib.use("agg") -matplotlib.rcParams['text.usetex'] = not running_in_ci() +matplotlib.rcParams["text.usetex"] = not running_in_ci() def get_parser(**kwargs): @@ -37,9 +37,7 @@ def get_parser(**kwargs): help="Name of the configuration file containing parameter values.", ) parser.add_argument( - "--model", - type=str, - help="Name of the kilonova model to be used" + "--model", type=str, help="Name of the kilonova model to be used" ) parser.add_argument( "--interpolation-type", @@ -73,11 +71,7 @@ def get_parser(**kwargs): type=str, help="Path to the data file in [time(isot) filter magnitude error] format", ) - parser.add_argument( - "--prior", - type=str, - help="Path to the prior file" - ) + parser.add_argument("--prior", type=str, help="Path to the prior file") parser.add_argument( "--tmin", type=float, @@ -91,10 +85,7 @@ def get_parser(**kwargs): help="Days to stop analysing from the trigger time (default: 14)", ) parser.add_argument( - "--dt", - type=float, - default=0.1, - help="Time step in day (default: 0.1)" + "--dt", type=float, default=0.1, help="Time step in day (default: 0.1)" ) parser.add_argument( "--log-space-time", @@ -170,10 +161,7 @@ def get_parser(**kwargs): help="Number of cores to be used, only needed for dynesty (default: 1)", ) parser.add_argument( - "--nlive", - type=int, - default=2048, - help="Number of live points (default: 2048)" + "--nlive", type=int, default=2048, help="Number of live points (default: 2048)" ) parser.add_argument( "--reactive-sampling", @@ -189,16 +177,10 @@ def get_parser(**kwargs): help="Sampling seed (default: 42)", ) parser.add_argument( - "--injection", - metavar="PATH", - type=str, - help="Path to the injection json file" + "--injection", metavar="PATH", type=str, help="Path to the injection json file" ) parser.add_argument( - "--plot", - action="store_true", - default=False, - help="add best fit plot" + "--plot", action="store_true", default=False, help="add best fit plot" ) parser.add_argument( "--bestfit", @@ -224,7 +206,7 @@ def get_parser(**kwargs): action="store_true", default=False, ) - parser.add_argument( #no use in this script + parser.add_argument( # no use in this script "--systematics-file", metavar="PATH", help="Path to systematics configuration file", @@ -269,12 +251,11 @@ def analysis(args): error_budget = args.error_budget light_curve_model = SimpleBolometricLightCurveModel( - model=args.model, - sample_times=sample_times + model=args.model, sample_times=sample_times ) # setup the prior - priors = create_prior_from_args(args.model.split(','), args) + priors = create_prior_from_args(args.model.split(","), args) # setup the likelihood likelihood_kwargs = dict( @@ -338,6 +319,7 @@ def analysis(args): if args.bestfit or args.plot: import matplotlib.pyplot as plt + posterior_file = os.path.join( args.outdir, f"{args.label}_posterior_samples.dat" ) @@ -367,10 +349,7 @@ def analysis(args): lbol_dict["bestfit_sample_times"] + bestfit_params["timeshift"] ) - matplotlib.rcParams.update( - {'font.size': 12, - 'font.family': 'Times New Roman'} - ) + matplotlib.rcParams.update({"font.size": 12, "font.family": "Times New Roman"}) plt.figure(1) plotName = os.path.join(args.outdir, f"{args.label}_lightcurves.png") @@ -385,20 +364,25 @@ def analysis(args): idx = np.where(np.isfinite(sigma_y))[0] plt.errorbar( - t[idx], y[idx], sigma_y[idx], fmt="o", color="k", markersize=12, + t[idx], + y[idx], + sigma_y[idx], + fmt="o", + color="k", + markersize=12, ) idx = np.where(~np.isfinite(sigma_y))[0] - plt.errorbar( - t[idx], y[idx], sigma_y[idx], fmt="v", color="k", markersize=12 + plt.errorbar(t[idx], y[idx], sigma_y[idx], fmt="v", color="k", markersize=12) + + plt.plot( + lbol_dict["bestfit_sample_times"], + lbol_dict["lbol"], + color=color, + linewidth=3, + linestyle="--", ) - plt.plot(lbol_dict['bestfit_sample_times'], lbol_dict['lbol'], - color=color, - linewidth=3, - linestyle="--", - ) - plt.ylabel("L [erg / s]") plt.xlabel("Time [days]") plt.tight_layout() diff --git a/nmma/em/combine_EOS.py b/nmma/em/combine_EOS.py index 3a778e7e1..713e9141d 100644 --- a/nmma/em/combine_EOS.py +++ b/nmma/em/combine_EOS.py @@ -13,7 +13,7 @@ from .utils import running_in_ci matplotlib.use("agg") -matplotlib.rcParams['text.usetex'] = not running_in_ci() +matplotlib.rcParams["text.usetex"] = not running_in_ci() c = seaborn.color_palette("colorblind") diff --git a/nmma/em/create_lightcurves.py b/nmma/em/create_lightcurves.py index 5c6539355..54432fde9 100644 --- a/nmma/em/create_lightcurves.py +++ b/nmma/em/create_lightcurves.py @@ -392,7 +392,7 @@ def main(args=None): import matplotlib matplotlib.use("agg") - matplotlib.rcParams['text.usetex'] = not running_in_ci() + matplotlib.rcParams["text.usetex"] = not running_in_ci() params = { "backend": "pdf", "axes.labelsize": 42, diff --git a/nmma/em/detect_lightcurves.py b/nmma/em/detect_lightcurves.py index 5c1e77b11..e37431c3d 100644 --- a/nmma/em/detect_lightcurves.py +++ b/nmma/em/detect_lightcurves.py @@ -45,33 +45,16 @@ def main(): required=True, help="lightcurve file directory with lightcurves", ) + parser.add_argument("-i", "--indices-file", type=str) + parser.add_argument("-d", "--detections-file", type=str) parser.add_argument( - "-i", - "--indices-file", - type=str + "--binary-type", type=str, required=True, help="Either BNS or NSBH" ) parser.add_argument( - "-d", - "--detections-file", - type=str + "--outdir", type=str, default="outdir", help="Path to the output directory" ) parser.add_argument( - "--binary-type", - type=str, - required=True, - help="Either BNS or NSBH" - ) - parser.add_argument( - "--outdir", - type=str, - default="outdir", - help="Path to the output directory" - ) - parser.add_argument( - "--telescope", - type=str, - default="ZTF", - help="telescope to recover" + "--telescope", type=str, default="ZTF", help="telescope to recover" ) parser.add_argument( "--tmin", @@ -86,10 +69,7 @@ def main(): help="Days to be stoped analysing from the trigger time (default: 14)", ) parser.add_argument( - "--dt", - type=float, - default=0.1, - help="Time step in day (default: 0.1)" + "--dt", type=float, default=0.1, help="Time step in day (default: 0.1)" ) parser.add_argument( "--filters", @@ -109,24 +89,22 @@ def main(): type=int, required=True ) - parser.add_argument( - "--parallel", - action="store_true", - default=False, - help="parallel the runs" + "--Tobs", + type=str, + default="0.0,1.0" ) parser.add_argument( - "--number-of-cores", - type=int, - default=1, - help="Number of cores" + "--parallel", action="store_true", default=False, help="parallel the runs" ) parser.add_argument( - "--label", - type=str, - required=True, - help="Label for the run" + "--number-of-cores", type=int, default=1, help="Number of cores" + ) + parser.add_argument("--label", type=str, required=True, help="Label for the run") + parser.add_argument( + "--plots", + help="comma delimited list of plot types (skymap,tiles,coverage,schedule,efficiency)", + default="skymap,tiles,coverage,schedule,efficiency", ) args = parser.parse_args() @@ -135,11 +113,11 @@ def main(): if args.injection_file.endswith(".json"): with open(args.injection_file, "rb") as f: injection_data = json.load(f) - datadict = injection_data["injections"]["content"] + datadict = injection_data["injections"]["content"] dataframe_from_inj = pd.DataFrame.from_dict(datadict) - - # get the injection index from the oserving scenarios simulations - simulation_id = dataframe_from_inj['simulation_id'] + + # get the injection index from the oserving scenarios simulations + simulation_id = dataframe_from_inj["simulation_id"] else: print("Only json supported.") exit(1) @@ -147,7 +125,7 @@ def main(): if len(dataframe_from_inj) > 0: args.n_injection = len(dataframe_from_inj) - #indices = np.loadtxt(args.indices_file) + # indices = np.loadtxt(args.indices_file) indices = simulation_id commands = [] @@ -156,18 +134,22 @@ def main(): outdir = os.path.join(args.outdir, str(index)) if not os.path.isdir(outdir): os.makedirs(outdir) - + skymap_file = os.path.join(args.skymap_dir, "%d.fits" % indices[index]) - lc_file = os.path.join(args.lightcurve_dir, "%s.dat" % f"{args.label}_{indices[index]}") - + lc_file = os.path.join( + args.lightcurve_dir, "%s.dat" % f"{args.label}_{indices[index]}" + ) + # fixed scheduling time as observation plan - #number_shot = int(1 + (args.tmax - args.tmin)/args.dt) - - #lcs[index] = np.loadtxt(lc_file)[0:number_shot,] - + # number_shot = int(1 + (args.tmax - args.tmin)/args.dt) + + # lcs[index] = np.loadtxt(lc_file)[0:number_shot,] + lcs[index] = np.loadtxt(lc_file) - efffile = os.path.join(outdir, f"efficiency_true_{args.label}_{indices[index]}.txt") + efffile = os.path.join( + outdir, f"efficiency_true_{args.label}_{indices[index]}.txt" + ) if os.path.isfile(efffile): continue @@ -184,15 +166,15 @@ def main(): exposuretime = ",".join( [str(args.exposuretime) for i in args.filters.split(",")] ) - + if args.telescope.lower() == 'ultrasat': - command = f"gwemopt-run --telescopes {args.telescope} --geometry 3d --doTiles --doSchedule --scheduleType greedy_slew --doTrueLocation --true_location --true_ra {ra} --true_dec {dec} --true_distance {dist} --powerlaw_cl 0.9 --doObservability --doObservabilityExit --timeallocationType powerlaw --gpstime {gpstime} --event {skymap_file} --filters {args.filters} --exposuretimes {exposuretime} --doSingleExposure --doAlternatingFilters --doEfficiency --lightcurveFiles {lc_file} -o {outdir} --modelType file --doPlots --doUsePrimary --ignore_observability" + command = f"gwemopt-run --telescopes {args.telescope} --geometry 3d --doTiles --doSchedule --scheduleType greedy_slew --doTrueLocation --true_location --true_ra {ra} --true_dec {dec} --true_distance {dist} --powerlaw_cl 0.9 --doObservability --doObservabilityExit --timeallocationType powerlaw --gpstime {gpstime} --event {skymap_file} --filters {args.filters} --exposuretimes {exposuretime} --doSingleExposure --doAlternatingFilters --doEfficiency --lightcurveFiles {lc_file} -o {outdir} --modelType file --doPlots --doUsePrimary --ignore_observability --Tobs {args.Tobs}" else: - command = f"gwemopt-run --telescopes {args.telescope} --geometry 3d --doTiles --doSchedule --scheduleType greedy_slew --doTrueLocation --true_location --true_ra {ra} --true_dec {dec} --true_distance {dist} --powerlaw_cl 0.9 --doObservability --doObservabilityExit --timeallocationType powerlaw --gpstime {gpstime} --event {skymap_file} --filters {args.filters} --exposuretimes {exposuretime} --doSingleExposure --doAlternatingFilters --doEfficiency --lightcurveFiles {lc_file} -o {outdir} --modelType file --doPlots --doUsePrimary" + command = f"gwemopt-run --telescopes {args.telescope} --geometry 3d --doTiles --doSchedule --scheduleType greedy_slew --doTrueLocation --true_location --true_ra {ra} --true_dec {dec} --true_distance {dist} --powerlaw_cl 0.9 --doObservability --doObservabilityExit --timeallocationType powerlaw --gpstime {gpstime} --event {skymap_file} --filters {args.filters} --exposuretimes {exposuretime} --doSingleExposure --doAlternatingFilters --doEfficiency --lightcurveFiles {lc_file} -o {outdir} --modelType file --doPlots --doUsePrimary --Tobs {args.Tobs}" commands.append(command) - + print("Number of jobs remaining... %d." % len(commands)) if args.parallel: @@ -207,27 +189,29 @@ def main(): absmag, effs, probs = [], [], [] fid = open(args.detections_file, "w") - #fid = open('lc_skymap_detection', "w") + # fid = open('lc_skymap_detection', "w") for index, row in dataframe_from_inj.iterrows(): outdir = os.path.join(args.outdir, str(index)) - - efffile = os.path.join(outdir, f"efficiency_true_{args.label}_{indices[index]}.txt") - - ## Choose band for the best proxy - # for ZTF r-band give the best proxy - if args.telescope == 'ZTF': + + efffile = os.path.join( + outdir, f"efficiency_true_{args.label}_{indices[index]}.txt" + ) + + # Choose band for the best proxy + # for ZTF r-band give the best proxy + if args.telescope == "ZTF": absmag.append(np.min(lcs[index][:, 2])) - - # For Rubin i-band or z-band are best - elif args.telescope == 'LSST': + + # For Rubin i-band or z-band are best + elif args.telescope == "LSST": absmag.append(np.min(lcs[index][:, 3])) - - elif args.telescope.lower() == 'ultrasat': + + elif args.telescope.lower() == "ultrasat": absmag.append(np.min(lcs[index][:, 1])) - + else: absmag.append(np.min(lcs[index][:, 3])) - + if not os.path.isfile(efffile): fid.write("0\n") effs.append(0.0) @@ -243,7 +227,7 @@ def main(): with open(efffile, "r") as file: data_out = file.read() effs.append(float(data_out.split("\n")[1].split("\t")[4])) - + efffile = os.path.join(outdir, f"efficiency_{args.label}_{indices[index]}.txt") if not os.path.isfile(efffile): probs.append(0.0) @@ -382,7 +366,7 @@ def main(): plt.axes(ax4) plt.axis("off") - cbar = fig.colorbar(sm, ax=ax4, orientation='horizontal') + cbar = fig.colorbar(sm, ax=ax4, orientation="horizontal") cbar.set_label(r"Detection Efficiency") plt.axes(ax3) @@ -460,5 +444,5 @@ def main(): ax2.set_xticklabels([]) plotName = os.path.join(plotdir, f"eff_{args.binary_type}_{args.telescope}.pdf") - plt.savefig(plotName,bbox_inches='tight', dpi=2500) - plt.close() \ No newline at end of file + plt.savefig(plotName, bbox_inches="tight", dpi=2500) + plt.close() diff --git a/nmma/em/gwem_Hubble_estimate.py b/nmma/em/gwem_Hubble_estimate.py index b01361599..676058645 100644 --- a/nmma/em/gwem_Hubble_estimate.py +++ b/nmma/em/gwem_Hubble_estimate.py @@ -10,8 +10,10 @@ def reweight_to_flat_mass_prior(df): - total_mass = bilby.gw.conversion.chirp_mass_and_mass_ratio_to_total_mass(df.chirp_mass, df.mass_ratio) - m1 = total_mass / (1. + df.mass_ratio) + total_mass = bilby.gw.conversion.chirp_mass_and_mass_ratio_to_total_mass( + df.chirp_mass, df.mass_ratio + ) + m1 = total_mass / (1.0 + df.mass_ratio) jacobian = m1 * m1 / df.chirp_mass df_new = df.sample(frac=0.3, weights=jacobian) return df_new @@ -19,24 +21,68 @@ def reweight_to_flat_mass_prior(df): def main(): - parser = argparse.ArgumentParser(description="Calculate the combination and seperate trend of estimated Hubble constant with GW and EM input") + parser = argparse.ArgumentParser( + description="Calculate the combination and seperate trend of estimated Hubble constant with GW and EM input" + ) parser.add_argument("--output-label", metavar="NAME", type=str, required=True) - parser.add_argument("--GWsamples", metavar='PATH', type=str, required=True, help="Path to the GW posterior samples directory, the samples files are expected to be in the format of posterior_samples_{i}.dat") - parser.add_argument("--EMsamples", metavar='PATH', type=str, required=True, help="Same as the GW samples but for EM samples") - parser.add_argument("--injection", metavar='PATH', type=str, required=True) - parser.add_argument("--inject-Hubble", metavar='H0', type=float, required=True) + parser.add_argument( + "--GWsamples", + metavar="PATH", + type=str, + required=True, + help="Path to the GW posterior samples directory, the samples files are expected to be in the format of posterior_samples_{i}.dat", + ) + parser.add_argument( + "--EMsamples", + metavar="PATH", + type=str, + required=True, + help="Same as the GW samples but for EM samples", + ) + parser.add_argument("--injection", metavar="PATH", type=str, required=True) + parser.add_argument("--inject-Hubble", metavar="H0", type=float, required=True) parser.add_argument("--detectable", type=str, required=False) parser.add_argument("--Nevent", type=int, required=False) - parser.add_argument("--N-prior-samples", type=int, required=False, default=10000, help="Number of prior samples to be used for resampling (default: 10000)") - parser.add_argument("--N-posterior-samples", type=int, required=False, default=6000, help="Number of posterior samples to be drawn during the resampling (default: 6000)") - parser.add_argument("--N-reordering", type=int, default=100, help="Number of reodering realisation to be comsidered (default: 100)") - parser.add_argument("--cred-interval", type=float, default=0.95, help="Credible interval to be calculated (default: 0.95)") - parser.add_argument("--seed", type=int, required=False, default=42, help="Random seed (default: 42)") - parser.add_argument("--p-value-threshold", type=float, required=False, help="p-value threshold used to remove badly recoved injections") + parser.add_argument( + "--N-prior-samples", + type=int, + required=False, + default=10000, + help="Number of prior samples to be used for resampling (default: 10000)", + ) + parser.add_argument( + "--N-posterior-samples", + type=int, + required=False, + default=6000, + help="Number of posterior samples to be drawn during the resampling (default: 6000)", + ) + parser.add_argument( + "--N-reordering", + type=int, + default=100, + help="Number of reodering realisation to be comsidered (default: 100)", + ) + parser.add_argument( + "--cred-interval", + type=float, + default=0.95, + help="Credible interval to be calculated (default: 0.95)", + ) + parser.add_argument( + "--seed", type=int, required=False, default=42, help="Random seed (default: 42)" + ) + parser.add_argument( + "--p-value-threshold", + type=float, + required=False, + help="p-value threshold used to remove badly recoved injections", + ) args = parser.parse_args() - assert args.detectable or args.Nevent, \ - "Input for detectable or number of injection is required" + assert ( + args.detectable or args.Nevent + ), "Input for detectable or number of injection is required" # load the injected Hubble constant H0_true = args.inject_Hubble @@ -46,10 +92,8 @@ def main(): detectable = np.loadtxt(args.detectable) # read the injection table - with open(args.injection, 'r') as f: - injection_dict = json.load( - f, object_hook=bilby.core.utils.decode_bilby_json - ) + with open(args.injection, "r") as f: + injection_dict = json.load(f, object_hook=bilby.core.utils.decode_bilby_json) injection_df = injection_dict["injections"] if args.detectable: @@ -66,22 +110,36 @@ def main(): p_values = [] for i in event_to_loop: try: - samplesEM = pd.read_csv('{0}/posterior_samples_{1}.dat'.format(args.EMsamples, i), header=0, delimiter=' ') - samplesGW = pd.read_csv('{0}/posterior_samples_{1}.dat'.format(args.GWsamples, i), header=0, delimiter=' ') + samplesEM = pd.read_csv( + "{0}/posterior_samples_{1}.dat".format(args.EMsamples, i), + header=0, + delimiter=" ", + ) + samplesGW = pd.read_csv( + "{0}/posterior_samples_{1}.dat".format(args.GWsamples, i), + header=0, + delimiter=" ", + ) except IOError: # this is designed for running on incomplete dataset continue distanceEM = samplesEM.luminosity_distance.to_numpy() # resample GW to flat in component mass - distanceGW = reweight_to_flat_mass_prior(samplesGW).luminosity_distance.to_numpy() + distanceGW = reweight_to_flat_mass_prior( + samplesGW + ).luminosity_distance.to_numpy() # Have the true redshift centered be the center of the posterior distance_injected = injection_df.luminosity_distance[i] - z_true = H0_true * distance_injected * (1e3 / scipy.constants.c) # the factor 100 is to convert speed of light to km s^-1 + z_true = ( + H0_true * distance_injected * (1e3 / scipy.constants.c) + ) # the factor 100 is to convert speed of light to km s^-1 if args.p_value_threshold: # check if the GW posterior is good - p_value = scipy.stats.percentileofscore(distanceGW, distance_injected) / 100 # as it function is returning percentage, we divide it by 100 + p_value = ( + scipy.stats.percentileofscore(distanceGW, distance_injected) / 100 + ) # as it function is returning percentage, we divide it by 100 p_value = 2 * min(p_value, 1 - p_value) p_values.append(p_value) # remove bad injections @@ -91,14 +149,20 @@ def main(): # generate redshift samples and Hubble samples redshift_samples = np.random.normal(float(z_true), 1e-3, size=len(distanceEM)) - velocity_samples = redshift_samples * scipy.constants.c / 1e3 # velocity in km s^-1 - H0_samples_EM = velocity_samples / distanceEM # Hubble samples in km s^-1 Mpc^-1 + velocity_samples = ( + redshift_samples * scipy.constants.c / 1e3 + ) # velocity in km s^-1 + H0_samples_EM = ( + velocity_samples / distanceEM + ) # Hubble samples in km s^-1 Mpc^-1 # generate the posterior on H0 with a gaussian_kde with weighting back to the uniform in volume # the reason for reweighting back to uniform in volume (p(d) ~ d^2) is because the selection effect # for such prior is known (N(H0) ~ H0^-3), which makes the further analysis easier # if the EM analysis is already using a uniform-in-volume prior for distance, this line is to be commented out - probs_EM[i] = scipy.stats.gaussian_kde(H0_samples_EM, weights=distanceEM * distanceEM) + probs_EM[i] = scipy.stats.gaussian_kde( + H0_samples_EM, weights=distanceEM * distanceEM + ) # now do it for GW redshift_samples = np.random.normal(float(z_true), 1e-3, size=len(distanceGW)) velocity_samples = redshift_samples * scipy.constants.c / 1e3 @@ -134,9 +198,9 @@ def main(): for idx, i in enumerate(index): if idx == 0: # the calculation is done in the log-space to prevent both overflowing and underflowing - logprob_combined_GW = 0. - logprob_combined_EM = 0. - logprob_combined_total = 0. + logprob_combined_GW = 0.0 + logprob_combined_EM = 0.0 + logprob_combined_total = 0.0 logprob_combined_GW = probs_GW[i].logpdf(H0_prior_samples) logprob_combined_GW -= scipy.special.logsumexp(logprob_combined_GW) @@ -152,24 +216,38 @@ def main(): # Although the selection correction is not needed for one event # with GW+EM it is effectively two events - logprob_combined_total = logprob_combined_GW + logprob_combined_EM + 3 * np.log(H0_prior_samples) - logprob_combined_total -= scipy.special.logsumexp(logprob_combined_total) + logprob_combined_total = ( + logprob_combined_GW + + logprob_combined_EM + + 3 * np.log(H0_prior_samples) + ) + logprob_combined_total -= scipy.special.logsumexp( + logprob_combined_total + ) prob_cumprod_total.append(np.exp(logprob_combined_total)) else: # when combing events, we are multiplying the posterior distribution and divided by the selection correction # in general we also need to divide the prior on H0, but since we are using a flat prior on it (line 147), it can be ignored - logprob_combined_GW = logprob_combined_GW + probs_GW[i].logpdf(H0_prior_samples) + logprob_combined_GW = logprob_combined_GW + probs_GW[i].logpdf( + H0_prior_samples + ) logprob_combined_GW = logprob_combined_GW + 3 * np.log(H0_prior_samples) logprob_combined_GW -= scipy.special.logsumexp(logprob_combined_GW) prob_cumprod_GW.append(np.exp(logprob_combined_GW)) - logprob_combined_EM = logprob_combined_EM + probs_EM[i].logpdf(H0_prior_samples) + logprob_combined_EM = logprob_combined_EM + probs_EM[i].logpdf( + H0_prior_samples + ) logprob_combined_EM = logprob_combined_EM + 3 * np.log(H0_prior_samples) logprob_combined_EM -= scipy.special.logsumexp(logprob_combined_EM) prob_cumprod_EM.append(np.exp(logprob_combined_EM)) - logprob_combined_total = logprob_combined_total + logprob_combined_GW + logprob_combined_EM - logprob_combined_total -= scipy.special.logsumexp(logprob_combined_total) + logprob_combined_total = ( + logprob_combined_total + logprob_combined_GW + logprob_combined_EM + ) + logprob_combined_total -= scipy.special.logsumexp( + logprob_combined_total + ) prob_cumprod_total.append(np.exp(logprob_combined_total)) # cast the list back to array for easier element-wise calculation @@ -193,7 +271,9 @@ def main(): for i in range(len(prob_cumprod_GW)): # calculate the posterior distribution using the prior samples # and the weighting that we previously calculated - GW_kde = scipy.stats.gaussian_kde(H0_prior_samples, weights=prob_cumprod_GW[i]) + GW_kde = scipy.stats.gaussian_kde( + H0_prior_samples, weights=prob_cumprod_GW[i] + ) # generate sample based on this distribution samples_GW = GW_kde.resample(size=args.N_posterior_samples)[0] # calculate the median and append it to the list @@ -204,14 +284,18 @@ def main(): H0_uplim_per_event_GW.append(cred_range[1]) H0_lowlim_per_event_GW.append(cred_range[0]) - EM_kde = scipy.stats.gaussian_kde(H0_prior_samples, weights=prob_cumprod_EM[i]) + EM_kde = scipy.stats.gaussian_kde( + H0_prior_samples, weights=prob_cumprod_EM[i] + ) samples_EM = EM_kde.resample(size=args.N_posterior_samples)[0] H0_med_per_event_EM.append(np.median(samples_EM)) cred_range = arviz.hdi(samples_EM, hdi_prob=args.cred_interval) H0_uplim_per_event_EM.append(cred_range[1]) H0_lowlim_per_event_EM.append(cred_range[0]) - total_kde = scipy.stats.gaussian_kde(H0_prior_samples, weights=prob_cumprod_total[i]) + total_kde = scipy.stats.gaussian_kde( + H0_prior_samples, weights=prob_cumprod_total[i] + ) samples_total = total_kde.resample(size=args.N_posterior_samples)[0] H0_med_per_event_total.append(np.median(samples_total)) cred_range = arviz.hdi(samples_total, hdi_prob=args.cred_interval) @@ -258,14 +342,16 @@ def main(): H0_lowlim_total_final = np.median(H0_lowlim_total, axis=0) # output the result using pandas - df_dict = dict(GW_med=H0_med_GW_final, - GW_uperr=H0_uplim_GW_final - H0_med_GW_final, - GW_lowerr=H0_med_GW_final - H0_lowlim_GW_final, - EM_med=H0_med_EM_final, - EM_uperr=H0_uplim_EM_final - H0_med_EM_final, - EM_lowerr=H0_med_EM_final - H0_lowlim_EM_final, - total_med=H0_med_total_final, - total_uperr=H0_uplim_total_final - H0_med_total_final, - total_lowerr=H0_med_total_final - H0_lowlim_total_final) + df_dict = dict( + GW_med=H0_med_GW_final, + GW_uperr=H0_uplim_GW_final - H0_med_GW_final, + GW_lowerr=H0_med_GW_final - H0_lowlim_GW_final, + EM_med=H0_med_EM_final, + EM_uperr=H0_uplim_EM_final - H0_med_EM_final, + EM_lowerr=H0_med_EM_final - H0_lowlim_EM_final, + total_med=H0_med_total_final, + total_uperr=H0_uplim_total_final - H0_med_total_final, + total_lowerr=H0_med_total_final - H0_lowlim_total_final, + ) df = pd.DataFrame.from_dict(df_dict) - df.to_csv('GW_EM_H0_trend_{0}.dat'.format(args.output_label), index=False, sep=' ') + df.to_csv("GW_EM_H0_trend_{0}.dat".format(args.output_label), index=False, sep=" ") diff --git a/nmma/em/gwem_resampling.py b/nmma/em/gwem_resampling.py index f49c0cab0..dbb9c8a46 100644 --- a/nmma/em/gwem_resampling.py +++ b/nmma/em/gwem_resampling.py @@ -13,38 +13,35 @@ def main(): ) parser.add_argument("--outdir", metavar="PATH", type=str, required=True) parser.add_argument( - "--GWsamples", - metavar="PATH", - type=str, - required=True, - help="If no posterior files are available, use gwsamples_creation.py to generate dummy GWsamples." + "--GWsamples", + metavar="PATH", + type=str, + required=True, + help="If no posterior files are available, use gwsamples_creation.py to generate dummy GWsamples.", ) parser.add_argument( - "--EMsamples", - metavar="PATH", - type=str, - required=True, - help="posterior samples file from a previous Bayesian inference run on EM signals (e.g. Kilonova inference or Kilonova+GRB inference.") + "--EMsamples", + metavar="PATH", + type=str, + required=True, + help="posterior samples file from a previous Bayesian inference run on EM signals (e.g. Kilonova inference or Kilonova+GRB inference.", + ) parser.add_argument( - "--EOSpath", - metavar="PATH", - type=str, - required=True, - help="Path of EOS folder, e.g. 15nsat_cse_uniform_R14 (located: https://zenodo.org/record/6106130#.YoysIHVBwUG)" + "--EOSpath", + metavar="PATH", + type=str, + required=True, + help="Path of EOS folder, e.g. 15nsat_cse_uniform_R14 (located: https://zenodo.org/record/6106130#.YoysIHVBwUG)", ) parser.add_argument( - "--Neos", - metavar="Neos", - type=int, - required=True, - help="Number of EOS files used for the inference." + "--Neos", + metavar="Neos", + type=int, + required=True, + help="Number of EOS files used for the inference.", ) parser.add_argument( - "--nlive", - metavar="nlive", - type=int, - required=False, - default=1024 + "--nlive", metavar="nlive", type=int, required=False, default=1024 ) parser.add_argument( "--GWprior", @@ -66,7 +63,7 @@ def main(): help="To run with total ejecta mass, if not activated, the two ejecta are consider seperately", ) parser.add_argument( - "--withNSBH", + "--withNSBH", action="store_true", help="Compute GW-EM-resampling for NSBH source, else: for BNS source.", ) @@ -87,7 +84,7 @@ def main(): # read the prior files GWprior = bilby.gw.prior.PriorDict(args.GWprior) EMprior = bilby.gw.prior.PriorDict(args.EMprior) - + if args.withNSBH: try: os.makedirs(args.outdir + "/pm/") @@ -140,9 +137,9 @@ def main(): args.withNSBH, **pymulti_kwargs ) - + if args.withNSBH: - + samples = solution.samples.T posterior_samples = dict() posterior_samples["chirp_mass"] = samples[0] @@ -154,13 +151,17 @@ def main(): posterior_samples["chi_2"] = samples[6] posterior_samples = pd.DataFrame.from_dict(posterior_samples) - posterior_samples.to_csv("{0}/posterior_samples.dat".format(args.outdir), sep=" ", index=False) + posterior_samples.to_csv( + "{0}/posterior_samples.dat".format(args.outdir), sep=" ", index=False + ) - sampler_functions.corner_plot(posterior_samples, solution, args.outdir,args.withNSBH) + sampler_functions.corner_plot( + posterior_samples, solution, args.outdir, args.withNSBH + ) else: samples = solution.samples.T - + posterior_samples = dict() posterior_samples["chirp_mass"] = samples[0] posterior_samples["mass_ratio"] = samples[1] @@ -169,9 +170,13 @@ def main(): posterior_samples["zeta"] = samples[4] posterior_samples = pd.DataFrame.from_dict(posterior_samples) - posterior_samples.to_csv("{0}/posterior_samples.dat".format(args.outdir), sep=" ", index=False) + posterior_samples.to_csv( + "{0}/posterior_samples.dat".format(args.outdir), sep=" ", index=False + ) - sampler_functions.corner_plot(posterior_samples, solution, args.outdir, args.withNSBH) + sampler_functions.corner_plot( + posterior_samples, solution, args.outdir, args.withNSBH + ) if __name__ == "__main__": diff --git a/nmma/em/gwem_resampling_utils.py b/nmma/em/gwem_resampling_utils.py index 44df218a4..90ebfcdd2 100644 --- a/nmma/em/gwem_resampling_utils.py +++ b/nmma/em/gwem_resampling_utils.py @@ -11,25 +11,33 @@ from pymultinest.solve import Solver - def luminosity_distance_to_redshift(distance): - from astropy import units - from astropy import cosmology as cosmo - cosmology = cosmo.Planck15 - - if hasattr(distance, '__len__') and len(distance)>50: #luminosity_distance_to_redshift gets really slow if too many distances are put in at once - zmin = cosmo.z_at_value(cosmology.luminosity_distance, distance.min() * units.Mpc) - zmax = cosmo.z_at_value(cosmology.luminosity_distance, distance.max() * units.Mpc) - zgrid = np.geomspace(zmin, zmax, 50) - distance_grid = cosmology.luminosity_distance(zgrid).value - return np.interp(distance, distance_grid, zgrid).value - else: - return cosmo.z_at_value(cosmology.luminosity_distance, distance *units.Mpc) + from astropy import units + from astropy import cosmology as cosmo + + cosmology = cosmo.Planck15 + + if ( + hasattr(distance, "__len__") and len(distance) > 50 + ): # luminosity_distance_to_redshift gets really slow if too many distances are put in at once + zmin = cosmo.z_at_value( + cosmology.luminosity_distance, distance.min() * units.Mpc + ) + zmax = cosmo.z_at_value( + cosmology.luminosity_distance, distance.max() * units.Mpc + ) + zgrid = np.geomspace(zmin, zmax, 50) + distance_grid = cosmology.luminosity_distance(zgrid).value + return np.interp(distance, distance_grid, zgrid).value + else: + return cosmo.z_at_value(cosmology.luminosity_distance, distance * units.Mpc) def construct_EM_KDE_seperate(EMsamples): - kde = scipy.stats.gaussian_kde((EMsamples.log10_mej_dyn.to_numpy(), EMsamples.log10_mej_wind.to_numpy())) + kde = scipy.stats.gaussian_kde( + (EMsamples.log10_mej_dyn.to_numpy(), EMsamples.log10_mej_wind.to_numpy()) + ) return kde @@ -37,79 +45,125 @@ def construct_EM_KDE_seperate(EMsamples): def construct_EM_KDE(EMsamples): if "log10_mej" in EMsamples.columns: - kde = scipy.stats.gaussian_kde(10**EMsamples.log10_mej.to_numpy()) + kde = scipy.stats.gaussian_kde(10 ** EMsamples.log10_mej.to_numpy()) elif "log10_mej_dyn" in EMsamples.columns and "log10_mej_wind" in EMsamples.columns: - total_eject_mass = 10**EMsamples.log10_mej_dyn.to_numpy() + 10**EMsamples.log10_mej_wind.to_numpy() + total_eject_mass = ( + 10 ** EMsamples.log10_mej_dyn.to_numpy() + + 10 ** EMsamples.log10_mej_wind.to_numpy() + ) kde = scipy.stats.gaussian_kde(total_eject_mass) else: - raise ValueError("EM samples must either contain total ejecta mass as 'log10_mej' or seperate ejecta masses as 'log10_mej_dyn' and 'log10_mej_wind'.") + raise ValueError( + "EM samples must either contain total ejecta mass as 'log10_mej' or seperate ejecta masses as 'log10_mej_dyn' and 'log10_mej_wind'." + ) return kde def mceta2m1m2(mc, eta): - M = mc / np.power(eta, 3. / 5.) - q = (1 - np.sqrt(1. - 4. * eta) - 2 * eta) / (2. * eta) + M = mc / np.power(eta, 3.0 / 5.0) + q = (1 - np.sqrt(1.0 - 4.0 * eta) - 2 * eta) / (2.0 * eta) - m1 = M / (1. + q) - m2 = M * q / (1. + q) + m1 = M / (1.0 + q) + m2 = M * q / (1.0 + q) return m1, m2 def lambdas2lambdaTs(lambda1, lambda2, q): - eta = q / np.power(1. + q, 2.) + eta = q / np.power(1.0 + q, 2.0) eta2 = eta * eta eta3 = eta2 * eta - root14eta = np.sqrt(1. - 4 * eta) - - lambdaT = (8. / 13.) * ((1. + 7 * eta - 31 * eta2) * (lambda1 + lambda2) + root14eta * (1. + 9 * eta - 11. * eta2) * (lambda1 - lambda2)) - dlambdaT = 0.5 * (root14eta * (1. - 13272. * eta / 1319. + 8944. * eta2 / 1319.) * (lambda1 + lambda2) + (1. - 15910. * eta / 1319. + 32850. * eta2 / 1319. + 3380. * eta3 / 1319.) * (lambda1 - lambda2)) + root14eta = np.sqrt(1.0 - 4 * eta) + + lambdaT = (8.0 / 13.0) * ( + (1.0 + 7 * eta - 31 * eta2) * (lambda1 + lambda2) + + root14eta * (1.0 + 9 * eta - 11.0 * eta2) * (lambda1 - lambda2) + ) + dlambdaT = 0.5 * ( + root14eta + * (1.0 - 13272.0 * eta / 1319.0 + 8944.0 * eta2 / 1319.0) + * (lambda1 + lambda2) + + ( + 1.0 + - 15910.0 * eta / 1319.0 + + 32850.0 * eta2 / 1319.0 + + 3380.0 * eta3 / 1319.0 + ) + * (lambda1 - lambda2) + ) return lambdaT, dlambdaT + def corner_plot(posterior_samples, solution, outdir, withNSBH): import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt import corner - matplotlib.rcParams['text.usetex'] = not running_in_ci() - matplotlib.rcParams.update({'font.size': 16, 'font.family': 'Times New Roman'}) - kwargs = dict(bins=50, smooth=1.3, label_kwargs=dict(fontsize=16), show_titles=True, - title_kwargs=dict(fontsize=16), color='#0072C1', - truth_color='tab:orange', quantiles=[0.05, 0.95], - levels=(0.10, 0.32, 0.68, 0.95), median_line=True, - plot_density=False, plot_datapoints=False, fill_contours=True, - max_n_ticks=4, hist_kwargs={'density': True}) - - - mc = posterior_samples['chirp_mass'].to_numpy() - invq = posterior_samples['mass_ratio'].to_numpy() - alpha = posterior_samples['alpha'].to_numpy() - zeta = posterior_samples['zeta'].to_numpy() - - eta = invq / np.power(1. + invq, 2) + + matplotlib.rcParams["text.usetex"] = not running_in_ci() + matplotlib.rcParams.update({"font.size": 16, "font.family": "Times New Roman"}) + kwargs = dict( + bins=50, + smooth=1.3, + label_kwargs=dict(fontsize=16), + show_titles=True, + title_kwargs=dict(fontsize=16), + color="#0072C1", + truth_color="tab:orange", + quantiles=[0.05, 0.95], + levels=(0.10, 0.32, 0.68, 0.95), + median_line=True, + plot_density=False, + plot_datapoints=False, + fill_contours=True, + max_n_ticks=4, + hist_kwargs={"density": True}, + ) + + mc = posterior_samples["chirp_mass"].to_numpy() + invq = posterior_samples["mass_ratio"].to_numpy() + alpha = posterior_samples["alpha"].to_numpy() + zeta = posterior_samples["zeta"].to_numpy() + + eta = invq / np.power(1.0 + invq, 2) m1, m2 = mceta2m1m2(mc, eta) q = m1 / m2 # inverted - - - if withNSBH:#NSBH resampling result corner plot - labels = [r'$\mathcal{M}_c[M_{\odot}]$', r'$q$', r'$\alpha[M_{\odot}]$', r'$\zeta$'] + if withNSBH: # NSBH resampling result corner plot + labels = [ + r"$\mathcal{M}_c[M_{\odot}]$", + r"$q$", + r"$\alpha[M_{\odot}]$", + r"$\zeta$", + ] plotSamples = np.vstack((mc, q, alpha, zeta)) - limits = ((np.amin(mc), np.amax(mc)), (np.amin(q), 3), (np.amin(alpha), np.amax(alpha)), (np.amin(zeta), np.amax(zeta))) + limits = ( + (np.amin(mc), np.amax(mc)), + (np.amin(q), 3), + (np.amin(alpha), np.amax(alpha)), + (np.amin(zeta), np.amax(zeta)), + ) plt.figure(1) corner.corner(plotSamples.T, labels=labels, range=limits, **kwargs) - plt.savefig('{0}/corner.pdf'.format(outdir), bbox_inches='tight') - - - else: #BNS resampling result corner plot - labels = [r'$\mathcal{M}_c[M_{\odot}]$', r'$q$', r'$\tilde{\Lambda}$', r'$\alpha[M_{\odot}]$', r'$\zeta$', r'$M_{\rm{max}}[M_{\odot}]$'] - - EOS = posterior_samples['EOS'].astype(int) + 1 + plt.savefig("{0}/corner.pdf".format(outdir), bbox_inches="tight") + + else: # BNS resampling result corner plot + labels = [ + r"$\mathcal{M}_c[M_{\odot}]$", + r"$q$", + r"$\tilde{\Lambda}$", + r"$\alpha[M_{\odot}]$", + r"$\zeta$", + r"$M_{\rm{max}}[M_{\odot}]$", + ] + + EOS = posterior_samples["EOS"].astype(int) + 1 lambda1 = [] lambda2 = [] for i in range(len(EOS)): @@ -126,38 +180,50 @@ def corner_plot(posterior_samples, solution, outdir, withNSBH): MTOV = np.array(MTOV) plotSamples = np.vstack((mc, q, lambdaT, alpha, zeta, MTOV)) - limits = ((np.amin(mc), np.amax(mc)), (np.amin(q), 3), (np.amin(lambdaT), np.amax(lambdaT)), (np.amin(alpha), np.amax(alpha)), (np.amin(zeta), np.amax(zeta)), (np.amin(MTOV), 2.7)) + limits = ( + (np.amin(mc), np.amax(mc)), + (np.amin(q), 3), + (np.amin(lambdaT), np.amax(lambdaT)), + (np.amin(alpha), np.amax(alpha)), + (np.amin(zeta), np.amax(zeta)), + (np.amin(MTOV), 2.7), + ) plt.figure(1) corner.corner(plotSamples.T, labels=labels, range=limits, **kwargs) - plt.savefig('{0}/corner.pdf'.format(outdir), bbox_inches='tight') - print("The 90% upper bound for lambdaT is {0}".format(np.quantile(lambdaT, 0.9))) + plt.savefig("{0}/corner.pdf".format(outdir), bbox_inches="tight") + print( + "The 90% upper bound for lambdaT is {0}".format(np.quantile(lambdaT, 0.9)) + ) class TotalEjectaMassInference(Solver): - - def __init__(self, GWsamples, EMsamples, GWprior, EMprior, Neos, EOSpath, withNSBH, **kwargs): + def __init__( + self, GWsamples, EMsamples, GWprior, EMprior, Neos, EOSpath, withNSBH, **kwargs + ): self.GWsamples = GWsamples self.EMsamples = EMsamples self.withNSBH = withNSBH - - if (withNSBH): - self.priors = {'chirp_mass': GWprior['chirp_mass'], - 'mass_ratio': GWprior['mass_ratio'], - 'EOS': Uniform(minimum=0, maximum=Neos + 1), - 'alpha': EMprior['alpha'], - 'zeta': EMprior['zeta'], - 'chi_1': GWprior['chi_1'], - 'chi_2': GWprior['chi_2'], - } - + + if withNSBH: + self.priors = { + "chirp_mass": GWprior["chirp_mass"], + "mass_ratio": GWprior["mass_ratio"], + "EOS": Uniform(minimum=0, maximum=Neos + 1), + "alpha": EMprior["alpha"], + "zeta": EMprior["zeta"], + "chi_1": GWprior["chi_1"], + "chi_2": GWprior["chi_2"], + } + else: - self.priors = {'chirp_mass': GWprior['chirp_mass'], - 'mass_ratio': GWprior['mass_ratio'], - 'EOS': Uniform(minimum=0, maximum=Neos + 1), - 'alpha': EMprior['alpha'], - 'zeta': EMprior['zeta']} - - + self.priors = { + "chirp_mass": GWprior["chirp_mass"], + "mass_ratio": GWprior["mass_ratio"], + "EOS": Uniform(minimum=0, maximum=Neos + 1), + "alpha": EMprior["alpha"], + "zeta": EMprior["zeta"], + } + self.priors = PriorDict(self.priors) self._search_parameter_keys = self.priors.keys() self.Neos = Neos @@ -165,10 +231,12 @@ def __init__(self, GWsamples, EMsamples, GWprior, EMprior, Neos, EOSpath, withNS EOS_radius_interp_dict = {} EOS_lambda_interp_dict = {} for i in range(1, Neos + 1): - r, m, Lambda = np.loadtxt('{0}/{1}.dat'.format(EOSpath, i), usecols=[0, 1, 2], unpack=True) + r, m, Lambda = np.loadtxt( + "{0}/{1}.dat".format(EOSpath, i), usecols=[0, 1, 2], unpack=True + ) - interp_r_of_m = scipy.interpolate.interp1d(m, r, kind='linear') - interp_lambda_of_m = scipy.interpolate.interp1d(m, Lambda, kind='linear') + interp_r_of_m = scipy.interpolate.interp1d(m, r, kind="linear") + interp_lambda_of_m = scipy.interpolate.interp1d(m, Lambda, kind="linear") EOS_radius_interp_dict[i] = interp_r_of_m EOS_lambda_interp_dict[i] = interp_lambda_of_m @@ -176,16 +244,18 @@ def __init__(self, GWsamples, EMsamples, GWprior, EMprior, Neos, EOSpath, withNS self.EOS_radius_interp_dict = EOS_radius_interp_dict self.EOS_lambda_interp_dict = EOS_lambda_interp_dict - z = luminosity_distance_to_redshift(self.GWsamples.luminosity_distance.to_numpy()) + z = luminosity_distance_to_redshift( + self.GWsamples.luminosity_distance.to_numpy() + ) mc = self.GWsamples.chirp_mass.to_numpy() / (1 + z) q = self.GWsamples.mass_ratio.to_numpy() - eta = q/(1+q)**2 + eta = q / (1 + q) ** 2 EOS = self.GWsamples.EOS.to_numpy() - if (withNSBH): + if withNSBH: chi_1 = self.GWsamples.chi_1.to_numpy() chi_2 = self.GWsamples.chi_2.to_numpy() - chi_eff = (chi_1 + q*chi_2)/(1+q) + chi_eff = (chi_1 + q * chi_2) / (1 + q) self.priorKDE((chi_1, chi_2, mc, eta)) else: @@ -193,8 +263,8 @@ def __init__(self, GWsamples, EMsamples, GWprior, EMprior, Neos, EOSpath, withNS self.EOSsamples = EOS.astype(int) + 1 self.EMKDE = construct_EM_KDE(self.EMsamples) - - if (self.withNSBH): + + if self.withNSBH: self.NSBHEjectaFitting = NSBHEjectaFitting() else: self.BNSEjectaFitting = BNSEjectaFitting() @@ -207,35 +277,41 @@ def Prior(self, x): def LogLikelihood(self, x): if self.withNSBH: mc, q, EOS, alpha, zeta, chi_1, chi_2 = x - chi_eff = (chi_1 + q*chi_2)/(1+q) + chi_eff = (chi_1 + q * chi_2) / (1 + q) else: mc, q, EOS, alpha, zeta = x EOS = int(EOS) + 1 - eta = q / np.power(1. + q, 2.) + eta = q / np.power(1.0 + q, 2.0) m1, m2 = mceta2m1m2(mc, eta) total_mass = m1 + m2 mass_ratio = m2 / m1 if len(np.where(self.EOSsamples == EOS)[0]) == 0: return np.nan_to_num(-np.inf) - + if self.withNSBH: - logprior = self.priorKDE.logpdf(chi_1, chi_2, mc, eta) + np.log(len(np.where(self.EOSsamples == EOS)[0])) + logprior = self.priorKDE.logpdf(chi_1, chi_2, mc, eta) + np.log( + len(np.where(self.EOSsamples == EOS)[0]) + ) try: r2 = self.EOS_radius_interp_dict[EOS](m2) except ValueError: return np.nan_to_num(-np.inf) - + C2 = m2 / (r2 * 1e3 / lal.MRSUN_SI) mdyn = self.NSBHEjectaFitting.dynamic_mass_fitting(m1, m2, C2, chi_eff) - mdisk = self.NSBHEjectaFitting.remnant_disk_mass_fitting(m1, m2, C2, chi_eff) + mdisk = self.NSBHEjectaFitting.remnant_disk_mass_fitting( + m1, m2, C2, chi_eff + ) log10_mwind = np.log10(zeta) + np.log10(mdisk) total_ejecta_mass = mdyn + 10**log10_mwind + alpha loglikelihood = self.EMKDE.logpdf(total_ejecta_mass) else: - logprior = self.priorKDE.logpdf(mc, eta) + np.log(len(np.where(self.EOSsamples == EOS)[0])) + logprior = self.priorKDE.logpdf(mc, eta) + np.log( + len(np.where(self.EOSsamples == EOS)[0]) + ) try: r1 = self.EOS_radius_interp_dict[EOS](m1) @@ -248,9 +324,11 @@ def LogLikelihood(self, x): R16 = self.EOS_radius_interp_dict[EOS](1.6) * 1e3 / lal.MRSUN_SI MTOV = self.EOS_radius_interp_dict[EOS].x[-1] - + mdyn = self.BNSEjectaFitting.dynamic_mass_fitting_KrFo(m1, m2, C1, C2) - log10_mdisk = self.BNSEjectaFitting.log10_disk_mass_fitting(total_mass, mass_ratio, MTOV, R16) + log10_mdisk = self.BNSEjectaFitting.log10_disk_mass_fitting( + total_mass, mass_ratio, MTOV, R16 + ) log10_mwind = np.log10(zeta) + log10_mdisk total_ejecta_mass = mdyn + 10**log10_mwind + alpha loglikelihood = self.EMKDE.logpdf(total_ejecta_mass) @@ -259,41 +337,45 @@ def LogLikelihood(self, x): class EjectaMassInference(Solver): - - def __init__(self, GWsamples, EMsamples, GWprior, EMprior, Neos, EOSpath, withNSBH, **kwargs): + def __init__( + self, GWsamples, EMsamples, GWprior, EMprior, Neos, EOSpath, withNSBH, **kwargs + ): self.GWsamples = GWsamples self.EMsamples = EMsamples self.withNSBH = withNSBH - if (withNSBH): - self.priors = {'chirp_mass': GWprior['chirp_mass'], - 'mass_ratio': GWprior['mass_ratio'], - 'EOS': Uniform(minimum=0, maximum=Neos + 1), - 'alpha': EMprior['alpha'], - 'zeta': EMprior['zeta'], - 'chi_1': GWprior['chi_1'], - 'chi_2': GWprior['chi_2'], - } + if withNSBH: + self.priors = { + "chirp_mass": GWprior["chirp_mass"], + "mass_ratio": GWprior["mass_ratio"], + "EOS": Uniform(minimum=0, maximum=Neos + 1), + "alpha": EMprior["alpha"], + "zeta": EMprior["zeta"], + "chi_1": GWprior["chi_1"], + "chi_2": GWprior["chi_2"], + } else: - self.priors = {'chirp_mass': GWprior['chirp_mass'], - 'mass_ratio': GWprior['mass_ratio'], - 'EOS': Uniform(minimum=0, maximum=Neos + 1), - 'alpha': EMprior['alpha'], - 'zeta': EMprior['zeta'], - } + self.priors = { + "chirp_mass": GWprior["chirp_mass"], + "mass_ratio": GWprior["mass_ratio"], + "EOS": Uniform(minimum=0, maximum=Neos + 1), + "alpha": EMprior["alpha"], + "zeta": EMprior["zeta"], + } self.priors = PriorDict(self.priors) self._search_parameter_keys = self.priors.keys() self.Neos = Neos - EOS_radius_interp_dict = {} EOS_lambda_interp_dict = {} for i in range(1, Neos + 1): - r, m, Lambda = np.loadtxt('{0}/{1}.dat'.format(EOSpath, i), usecols=[0, 1, 2], unpack=True) + r, m, Lambda = np.loadtxt( + "{0}/{1}.dat".format(EOSpath, i), usecols=[0, 1, 2], unpack=True + ) - interp_r_of_m = scipy.interpolate.interp1d(m, r, kind='linear') - interp_lambda_of_m = scipy.interpolate.interp1d(m, Lambda, kind='linear') + interp_r_of_m = scipy.interpolate.interp1d(m, r, kind="linear") + interp_lambda_of_m = scipy.interpolate.interp1d(m, Lambda, kind="linear") EOS_radius_interp_dict[i] = interp_r_of_m EOS_lambda_interp_dict[i] = interp_lambda_of_m @@ -301,24 +383,26 @@ def __init__(self, GWsamples, EMsamples, GWprior, EMprior, Neos, EOSpath, withNS self.EOS_radius_interp_dict = EOS_radius_interp_dict self.EOS_lambda_interp_dict = EOS_lambda_interp_dict - z = luminosity_distance_to_redshift(self.GWsamples.luminosity_distance.to_numpy()) + z = luminosity_distance_to_redshift( + self.GWsamples.luminosity_distance.to_numpy() + ) mc = self.GWsamples.chirp_mass.to_numpy() / (1 + z) q = self.GWsamples.mass_ratio.to_numpy() EOS = self.GWsamples.EOS.to_numpy() - - if (withNSBH): + + if withNSBH: chi_1 = self.GWsamples.chi_1.to_numpy() chi_2 = self.GWsamples.chi_2.to_numpy() - chi_eff = (chi_1 + q*chi_2)/(1+q) + chi_eff = (chi_1 + q * chi_2) / (1 + q) self.chi_1KDE = scipy.stats.gaussian_kde(chi_1) self.chi_2KDE = scipy.stats.gaussian_kde(chi_2) - + self.mcKDE = scipy.stats.gaussian_kde(mc) - self.invqKDE = scipy.stats.gaussian_kde(1. / q) + self.invqKDE = scipy.stats.gaussian_kde(1.0 / q) self.EOSsamples = EOS.astype(int) + 1 self.EMKDE = construct_EM_KDE_seperate(self.EMsamples) - if (withNSBH): + if withNSBH: self.NSBHEjectaFitting = NSBHEjectaFitting() else: self.BNSEjectaFitting = BNSEjectaFitting() @@ -331,42 +415,52 @@ def Prior(self, x): def LogLikelihood(self, x): if self.withNSBH: - mc, q, EOS, alpha, zeta, chi_1, chi_2 = x - chi_eff = (chi_1 + q*chi_2)/(1+q) + mc, q, EOS, alpha, zeta, chi_1, chi_2 = x + chi_eff = (chi_1 + q * chi_2) / (1 + q) else: mc, q, EOS, alpha, zeta = x - + EOS = int(EOS) + 1 - eta = q / np.power(1. + q, 2.) + eta = q / np.power(1.0 + q, 2.0) m1, m2 = mceta2m1m2(mc, eta) total_mass = m1 + m2 mass_ratio = m2 / m1 if len(np.where(self.EOSsamples == EOS)[0]) == 0: return np.nan_to_num(-np.inf) - - + if self.withNSBH: - logprior = self.chi_1KDE.logpdf(chi_1) + self.chi_2KDE.logpdf(chi_2) + self.mcKDE.logpdf(mc) + self.invqKDE.logpdf(m1 / m2) + np.log(len(np.where(self.EOSsamples == EOS)[0])) + logprior = ( + self.chi_1KDE.logpdf(chi_1) + + self.chi_2KDE.logpdf(chi_2) + + self.mcKDE.logpdf(mc) + + self.invqKDE.logpdf(m1 / m2) + + np.log(len(np.where(self.EOSsamples == EOS)[0])) + ) try: r2 = self.EOS_radius_interp_dict[EOS](m2) except ValueError: return np.nan_to_num(-np.inf) C2 = m2 / (r2 * 1e3 / lal.MRSUN_SI) - - #mdyn = self.NSBHEjectaFitting.dynamic_mass_fitting(m1, m2, C2, chi_eff) + + # mdyn = self.NSBHEjectaFitting.dynamic_mass_fitting(m1, m2, C2, chi_eff) mdyn = self.NSBHEjectaFitting.dyn_mass2(m1, m2, C2, chi_eff) - if mdyn < 0.: + if mdyn < 0.0: return -np.inf - mdisk = self.NSBHEjectaFitting.remnant_disk_mass_fitting(m1, m2, C2, chi_eff) + mdisk = self.NSBHEjectaFitting.remnant_disk_mass_fitting( + m1, m2, C2, chi_eff + ) log10_mwind = np.log10(zeta) + np.log10(mdisk) loglikelihood = self.EMKDE.logpdf((np.log10(mdyn), log10_mwind)) - else: - logprior = self.mcKDE.logpdf(mc) + self.invqKDE.logpdf(m1 / m2) + np.log(len(np.where(self.EOSsamples == EOS)[0])) + logprior = ( + self.mcKDE.logpdf(mc) + + self.invqKDE.logpdf(m1 / m2) + + np.log(len(np.where(self.EOSsamples == EOS)[0])) + ) try: r1 = self.EOS_radius_interp_dict[EOS](m1) r2 = self.EOS_radius_interp_dict[EOS](m2) @@ -379,12 +473,16 @@ def LogLikelihood(self, x): R16 = self.EOS_radius_interp_dict[EOS](1.6) * 1e3 / lal.MRSUN_SI MTOV = self.EOS_radius_interp_dict[EOS].x[-1] - mdyn = self.BNSEjectaFitting.dynamic_mass_fitting_KrFo(m1, m2, C1, C2) + alpha - - if mdyn < 0.: + mdyn = ( + self.BNSEjectaFitting.dynamic_mass_fitting_KrFo(m1, m2, C1, C2) + alpha + ) + + if mdyn < 0.0: return -np.inf - log10_mdisk = self.BNSEjectaFitting.log10_disk_mass_fitting(total_mass, mass_ratio, MTOV, R16) + log10_mdisk = self.BNSEjectaFitting.log10_disk_mass_fitting( + total_mass, mass_ratio, MTOV, R16 + ) log10_mwind = np.log10(zeta) + log10_mdisk loglikelihood = self.EMKDE.logpdf((np.log10(mdyn), log10_mwind)) diff --git a/nmma/em/gwsamples_creation.py b/nmma/em/gwsamples_creation.py index 2266e9e0c..dcb67df23 100644 --- a/nmma/em/gwsamples_creation.py +++ b/nmma/em/gwsamples_creation.py @@ -2,16 +2,18 @@ import pandas as pd import bilby -# load posterior file, e.g. from https://zenodo.org/record/6106130#.YoysIHVBwUG -eos_post = np.loadtxt('./posterior_probability_files/Astro/15nsat_cse_uniform_R14/posterior_probability.txt') +# load posterior file, e.g. from https://zenodo.org/record/6106130#.YoysIHVBwUG +eos_post = np.loadtxt( + "./posterior_probability_files/Astro/15nsat_cse_uniform_R14/posterior_probability.txt" +) npts = 150000 Neos = 5000 nparams = 3 ############# [mass1, mass2, DL] -params_low = [1.001398, 1.001398, 1] -params_high = [2.2, 2.2, 75] +params_low = [1.001398, 1.001398, 1] +params_high = [2.2, 2.2, 75] # 1) create dummy EOS samples with eos_post from nature paper EOS_raw = np.arange(0, Neos) # the gwem_resampling will add one to this @@ -28,7 +30,16 @@ lum_distance = np.random.uniform(params_low[2], params_high[2], size=npts) # 3) create pandas dataframe -dataset = pd.DataFrame({'mass_1': mass_1, 'mass_2': mass_2, 'chirp_mass': chirp_mass, 'mass_ratio': mass_ratio, 'luminosity_distance': lum_distance, 'EOS': EOS_samples}) +dataset = pd.DataFrame( + { + "mass_1": mass_1, + "mass_2": mass_2, + "chirp_mass": chirp_mass, + "mass_ratio": mass_ratio, + "luminosity_distance": lum_distance, + "EOS": EOS_samples, + } +) # 4) save GWsamples.dat file -dataset.to_csv('GWsamples.dat', index=False, sep=' ') +dataset.to_csv("GWsamples.dat", index=False, sep=" ") diff --git a/nmma/em/injection.py b/nmma/em/injection.py index e6c61fc5b..b9609f59d 100644 --- a/nmma/em/injection.py +++ b/nmma/em/injection.py @@ -386,7 +386,9 @@ def create_light_curve_data( [3.0, ["ps1__g", "ps1__z"]], ] else: - raise ValueError("args.rubin_ToO_type should be either platinum, gold, or silver") + raise ValueError( + "args.rubin_ToO_type should be either platinum, gold, or silver" + ) # took type names from Rubin 2024 Workshop write-up mjds, passbands = [], [] @@ -414,7 +416,7 @@ def create_light_curve_data( assume_sorted=True, ) times = group["mjd"].tolist() - #print("The times of observation are: ", times) + # print("The times of observation are: ", times) data_per_filt = np.vstack([times, lc(times), lcerr(times)]).T data[filt] = data_per_filt passbands_to_keep.append(filt) diff --git a/nmma/em/injection_summary.py b/nmma/em/injection_summary.py index 9b3ab5794..ee27890c3 100644 --- a/nmma/em/injection_summary.py +++ b/nmma/em/injection_summary.py @@ -134,7 +134,7 @@ def loglike_H0(cube, ndim, nparams): c = 299792458.0 * 1e-3 vr_mean, vr_std = z * c, zerr * c # noqa: F821 - pvr = (1 / np.sqrt(2 * np.pi * vr_std ** 2)) * np.exp( + pvr = (1 / np.sqrt(2 * np.pi * vr_std**2)) * np.exp( (-1 / 2.0) * ((vr_mean - H0 * d) / vr_std) ** 2 ) prob_dist = kde_eval_single(kdedir_dist, [d])[0] # noqa: F821 @@ -555,7 +555,7 @@ def main(): m_TOV, R14_EM = np.zeros(len(data)), np.zeros(len(data)) lambda_tilde_EM = np.zeros(len(data)) for ii, (q, eos, alpha, zeta, logl) in enumerate(data): - zeta = 10 ** zeta + zeta = 10**zeta eos = np.floor(eos) params = copy.deepcopy(default_parameters) @@ -1146,7 +1146,7 @@ def main(): for ii, key in enumerate(list(data_out.keys())): vals = data_out[key][params_pair[1]] if params_pair[1] == "ratio_zeta": - vals = 10 ** vals + vals = 10**vals true = data_out[key][params_pair[0]] parts = plt.violinplot(vals, [true], widths=0.01) @@ -1176,7 +1176,7 @@ def main(): for ii, key in enumerate(list(data_out.keys())): vals = data_out[key][params_pair[1]] if params_pair[1] == "ratio_zeta": - vals = 10 ** vals + vals = 10**vals true = data_out[key][params_pair[0]] dhist, edges = np.histogram(vals, bins=bin_edges) idx = np.argmin(np.abs(bins - true)) diff --git a/nmma/em/io.py b/nmma/em/io.py index 388f19ce8..d3934d188 100644 --- a/nmma/em/io.py +++ b/nmma/em/io.py @@ -14,10 +14,10 @@ def loadEvent(filename): """ Reads in lightcurve data from a file and returns data in a dictionary format. - + Args: - filename (str): Path to lightcurve file - + Returns: - data (dict): Dictionary containing the lightcurve data from the file. The keys are generally 't' and each of the filters in the file as well as their accompanying error values. """ @@ -109,30 +109,37 @@ def read_spectroscopy_files( return data -def read_photometry_files(files: list, filters: list = None, tt: np.array = np.linspace(0, 14, 100), datatype:str ="bulla") -> dict: +def read_photometry_files( + files: list, + filters: list = None, + tt: np.array = np.linspace(0, 14, 100), + datatype: str = "bulla", +) -> dict: """ Read in a list of photometry files with given filenames and process them in a dictionary - + Args: files (list): List of filenames with the photometry files. filters (list): List of photometry filters to be extracted. tt (np.array): Array containing the time grid at which photometry values are given. datatype (str): Which model we are considering. Currently supports - + Returns: - data: Dictionary with keys being the given filenames and values being dictionaries themselves, with keys + data: Dictionary with keys being the given filenames and values being dictionaries themselves, with keys being t (time) and specified filters and values being the time grid, and values the time grid and lightcurves. """ - + # First, check whether given datatype is supported in this function supported_datatypes = ["ztf", "bulla", "standard", "hdf5"] if datatype not in supported_datatypes: space = " " - raise ValueError(f"datatype {datatype} unknown. Currently supported datatypes are: {space.join(supported_datatypes)}") + raise ValueError( + f"datatype {datatype} unknown. Currently supported datatypes are: {space.join(supported_datatypes)}" + ) - # Return value + # Return value data = {} - + # Iterate over all the given files and extract the lightcurve data from it for filename in files: name = ( @@ -186,7 +193,7 @@ def read_photometry_files(files: list, filters: list = None, tt: np.array = np.l assume_sorted=True, ) data[name][filt] = lc(tt) - + # Bulla datatype elif datatype == "bulla": with open(filename, "r") as f: @@ -283,7 +290,7 @@ def read_photometry_files(files: list, filters: list = None, tt: np.array = np.l def read_lightcurve_file(filename: str) -> dict: """ - Function to read in lightcurve file and create a dictionary containing the time (in days) at which the lightcurves + Function to read in lightcurve file and create a dictionary containing the time (in days) at which the lightcurves are evaluated and the corresponding values for different filters. """ @@ -308,17 +315,17 @@ def detection_limit_from_m4opt_fits_file(fits_file, ra, dec): # Get the BinTableHDU containing the HEALPix data bintable = hdul[1] # Assuming it's the first extension # Extract the LIMMAG data and flatten it - limmag_data = bintable.data['LIMMAG'] + limmag_data = bintable.data["LIMMAG"] limmag_map = limmag_data.flatten() - + # Get NSIDE from the header - nside = bintable.header['NSIDE'] # Should be 128 based on your header - + nside = bintable.header["NSIDE"] # Should be 128 based on your header + # Function to get limiting magnitude at a specific sky location def get_limmag_at_position(ra_deg, dec_deg, limmag_map, nside): """ Get the limiting magnitude at a specific RA, Dec position - + Parameters: ----------- ra_deg : float @@ -329,7 +336,7 @@ def get_limmag_at_position(ra_deg, dec_deg, limmag_map, nside): The HEALPix map containing limiting magnitude values nside : int The HEALPix NSIDE parameter - + Returns: -------- float @@ -339,12 +346,12 @@ def get_limmag_at_position(ra_deg, dec_deg, limmag_map, nside): # theta: co-latitude (0 at north pole, pi at south pole) # phi: longitude (0 to 2*pi) theta = np.radians(90.0 - dec_deg) # Convert Dec to co-latitude - phi = np.radians(ra_deg) # RA is already longitude-like - + phi = np.radians(ra_deg) # RA is already longitude-like + # Get the corresponding HEALPix pixel index # nest=True because your map uses NESTED ordering pixel_idx = hp.ang2pix(nside, theta, phi, nest=True) - + # Return the value at that pixel return limmag_map[pixel_idx] diff --git a/nmma/em/likelihood.py b/nmma/em/likelihood.py index dd7e20728..283b7cdb9 100644 --- a/nmma/em/likelihood.py +++ b/nmma/em/likelihood.py @@ -11,6 +11,7 @@ from bilby.core.likelihood import Likelihood from . import utils, systematics + def truncated_gaussian(m_det, m_err, m_est, lim): a, b = (-np.inf - m_est) / m_err, (lim - m_est) / m_err @@ -170,7 +171,9 @@ def log_likelihood(self): for i in range(1, n + 1): value = self.parameters.get(f"sys_err_{filt}{i}") - globals()[f"sys_err_{filt}_array"] = np.append(globals()[f"sys_err_{filt}_array"], value) + globals()[f"sys_err_{filt}_array"] = np.append( + globals()[f"sys_err_{filt}_array"], value + ) for filter_group in yaml_dict["config"]["withTime"]["filters"]: if isinstance(filter_group, list): @@ -179,8 +182,12 @@ def log_likelihood(self): filt = "all" else: filt = filter_group - globals()[f"sys_err_{filt}_interped"] = interp1d(time_array, globals()[f"sys_err_{filt}_array"], fill_value="extrapolate", bounds_error=False) - + globals()[f"sys_err_{filt}_interped"] = interp1d( + time_array, + globals()[f"sys_err_{filt}_array"], + fill_value="extrapolate", + bounds_error=False, + ) # compare the estimated light curve and the measured data minus_chisquare_total = 0.0 @@ -208,16 +215,22 @@ def log_likelihood(self): break if not filter_group_finite_idx_match: filters_to_use = "all" - data_sigma = np.sqrt(data_sigma**2 + (globals()[f"sys_err_{filters_to_use}_interped"](data_time)) ** 2) - + data_sigma = np.sqrt( + data_sigma**2 + + (globals()[f"sys_err_{filters_to_use}_interped"](data_time)) + ** 2 + ) + elif yaml_dict["config"]["withoutTime"]["value"]: - data_sigma = np.sqrt(data_sigma**2 + self.parameters["sys_err"] ** 2) + data_sigma = np.sqrt( + data_sigma**2 + self.parameters["sys_err"] ** 2 + ) # include the error budget into calculation - elif 'sys_err' in self.parameters: - data_sigma = np.sqrt(data_sigma**2 + self.parameters['sys_err']**2) + elif "sys_err" in self.parameters: + data_sigma = np.sqrt(data_sigma**2 + self.parameters["sys_err"] ** 2) else: - data_sigma = np.sqrt(data_sigma**2 + self.error_budget[filt]**2) + data_sigma = np.sqrt(data_sigma**2 + self.error_budget[filt] ** 2) # evaluate the light curve magnitude at the data points mag_est = mag_app_interp[filt](data_time) @@ -256,26 +269,40 @@ def log_likelihood(self): for yaml_filt in yaml_filters: if yaml_filt is not None and filt in yaml_filt: if isinstance(yaml_filt, str): - upperlim_sigma = globals()[f"sys_err_{yaml_filt}_interped"](data_time)[infIdx] + upperlim_sigma = globals()[ + f"sys_err_{yaml_filt}_interped" + ](data_time)[infIdx] elif isinstance(yaml_filt, list): filters_to_use = "___".join(yaml_filt) - upperlim_sigma = globals()[f"sys_err_{filters_to_use}_interped"](data_time)[infIdx] + upperlim_sigma = globals()[ + f"sys_err_{filters_to_use}_interped" + ](data_time)[infIdx] filter_group_infinite_idx_match = True break if not filter_group_infinite_idx_match: filters_to_use = "all" - upperlim_sigma = globals()[f"sys_err_{filters_to_use}_interped"](data_time)[infIdx] + upperlim_sigma = globals()[ + f"sys_err_{filters_to_use}_interped" + ](data_time)[infIdx] - gausslogsf = scipy.stats.norm.logsf(data_mag[infIdx], mag_est[infIdx], upperlim_sigma) + gausslogsf = scipy.stats.norm.logsf( + data_mag[infIdx], mag_est[infIdx], upperlim_sigma + ) elif yaml_dict["config"]["withoutTime"]["value"]: upperlim_sigma = self.parameters["sys_err"] - gausslogsf = scipy.stats.norm.logsf(data_mag[infIdx], mag_est[infIdx], upperlim_sigma) - elif 'sys_err' in self.parameters: - upperlim_sigma = self.parameters['sys_err'] - gausslogsf = scipy.stats.norm.logsf(data_mag[infIdx], mag_est[infIdx], upperlim_sigma) + gausslogsf = scipy.stats.norm.logsf( + data_mag[infIdx], mag_est[infIdx], upperlim_sigma + ) + elif "sys_err" in self.parameters: + upperlim_sigma = self.parameters["sys_err"] + gausslogsf = scipy.stats.norm.logsf( + data_mag[infIdx], mag_est[infIdx], upperlim_sigma + ) else: - gausslogsf = scipy.stats.norm.logsf(data_mag[infIdx], mag_est[infIdx], self.error_budget[filt]) + gausslogsf = scipy.stats.norm.logsf( + data_mag[infIdx], mag_est[infIdx], self.error_budget[filt] + ) gaussprob_total += np.sum(gausslogsf) log_prob = minus_chisquare_total + gaussprob_total @@ -357,21 +384,21 @@ def log_likelihood(self): lbol_interp = interp1d( self.sample_times[idx] + t0, lbol[idx], - fill_value='extrapolate', - bounds_error=False + fill_value="extrapolate", + bounds_error=False, ) # compare the estimated light curve and the measured data minus_chisquare_total = 0.0 gaussprob_total = 0.0 - data_time = self.light_curve_data['phase'].to_numpy() - data_lbol = self.light_curve_data['Lbb'].to_numpy() - data_sigma = self.light_curve_data['Lbb_unc'].to_numpy() + data_time = self.light_curve_data["phase"].to_numpy() + data_lbol = self.light_curve_data["Lbb"].to_numpy() + data_sigma = self.light_curve_data["Lbb_unc"].to_numpy() # include the error budget into calculation - if 'em_syserr' in self.parameters: - data_sigma = np.sqrt(data_sigma**2 + self.parameters['em_syserr']**2) + if "em_syserr" in self.parameters: + data_sigma = np.sqrt(data_sigma**2 + self.parameters["em_syserr"] ** 2) else: data_sigma = np.sqrt(data_sigma**2 + self.error_budget**2) @@ -402,8 +429,8 @@ def log_likelihood(self): # evaluate the data with infinite error if len(infIdx) > 0: - if 'em_syserr' in self.parameters: - upperlim_sigma = self.parameters['em_syserr'] + if "em_syserr" in self.parameters: + upperlim_sigma = self.parameters["em_syserr"] gausslogsf = scipy.stats.norm.logsf( data_lbol[infIdx], lbol_est[infIdx], upperlim_sigma ) diff --git a/nmma/em/model.py b/nmma/em/model.py index 7e3b5cdef..b45a4dea7 100644 --- a/nmma/em/model.py +++ b/nmma/em/model.py @@ -40,7 +40,6 @@ ], "Piro2021": ["log10_Menv", "log10_Renv", "log10_Ee"], "Me2017": ["log10_mej", "log10_vej", "beta", "log10_kappa_r"], - "HoNa2020": ["log10_mej", "vej_max", "vej_min", "vej_frac", "log10_kappa_low_vej", "log10_kappa_high_vej"], "Bu2022mv": ["log10_mej_dyn", "vej_dyn", "log10_mej_wind", "vej_wind", "KNtheta"], "PL_BB_fixedT": ["bb_luminosity", "temperature", "beta", "powerlaw_mag"], "blackbody_fixedT": ["bb_luminosity", "temperature"], @@ -99,7 +98,7 @@ "KNtheta", ], "HoNa2020": [ - "log10_Mej", + "log10_mej", "vej_max", "vej_min", "vej_frac", @@ -150,7 +149,10 @@ def citation(self): "TrPi2018": ["https://arxiv.org/abs/1909.11691"], "Piro2021": ["https://arxiv.org/abs/2007.08543"], "Me2017": ["https://arxiv.org/abs/1910.01617"], - "HoNa2020": ["https://arxiv.org/abs/1909.02581", "https://arxiv.org/abs/1206.2379"], + "HoNa2020": [ + "https://arxiv.org/abs/1909.02581", + "https://arxiv.org/abs/1206.2379", + ], "Sr2023": [None], # TODO: add citation, "nugent-hyper": [ "https://sncosmo.readthedocs.io/en/stable/source-list.html" @@ -266,9 +268,9 @@ def __init__( local_only=False, ): if model_parameters is None: - assert model in model_parameters_dict.keys(), ( - "Unknown model,please update model_parameters_dict at em/model.py" - ) + assert ( + model in model_parameters_dict.keys() + ), "Unknown model,please update model_parameters_dict at em/model.py" self.model_parameters = model_parameters_dict[model] else: self.model_parameters = model_parameters @@ -515,9 +517,9 @@ def __init__( give a set of parameters """ - assert model in model_parameters_dict.keys(), ( - "Unknown model,please update model_parameters_dict at em/model.py" - ) + assert ( + model in model_parameters_dict.keys() + ), "Unknown model,please update model_parameters_dict at em/model.py" self.model = model self.model_parameters = model_parameters_dict[model] self.sample_times = sample_times @@ -925,9 +927,9 @@ def __init__( give a set of parameters """ - assert model in model_parameters_dict.keys(), ( - "Unknown model,please update model_parameters_dict at em/model.py" - ) + assert ( + model in model_parameters_dict.keys() + ), "Unknown model,please update model_parameters_dict at em/model.py" self.model = model self.model_parameters = model_parameters_dict[model] self.sample_times = sample_times @@ -1048,9 +1050,9 @@ def __init__( give a set of parameters """ - assert model in model_parameters_dict.keys(), ( - "Unknown model,please update model_parameters_dict at em/model.py" - ) + assert ( + model in model_parameters_dict.keys() + ), "Unknown model,please update model_parameters_dict at em/model.py" self.model = model self.model_parameters = model_parameters_dict[model] self.sample_times = sample_times @@ -1090,19 +1092,23 @@ def generate_lightcurve(self, sample_times, parameters): vej_max = param_dict["vej_max"] vej_min = param_dict["vej_min"] vej_range = vej_max - vej_min - vej = param_dict["vej_frac"] * vej_range + vej_min + vej = param_dict["vej_frac"] * vej_range + vej_min # calculate the temperature and luminosity to feed into the blackbody radiation calculation L, T, _ = utils.lightcurve_HoNa( sample_times, - 10**param_dict["log10_mej"], + 10 ** param_dict["log10_mej"], [param_dict["vej_min"], vej, param_dict["vej_max"]], - [10**param_dict["log10_kappa_low_vej"], - 10**param_dict["log10_kappa_high_vej"]], - param_dict["n"] + [ + 10 ** param_dict["log10_kappa_low_vej"], + 10 ** param_dict["log10_kappa_high_vej"], + ], + param_dict["n"], ) param_dict["bb_luminosity"] = L.cgs.value param_dict["temperature"] = T.si.value - _, lbol, mag = utils.blackbody_constant_temperature(sample_times, param_dict, filters=self.filters) + _, lbol, mag = utils.blackbody_constant_temperature( + sample_times, param_dict, filters=self.filters + ) elif self.model == "PL_BB_fixedT": _, lbol, mag = utils.powerlaw_blackbody_constant_temperature_lc( sample_times, param_dict, filters=self.filters diff --git a/nmma/em/model_lbol.py b/nmma/em/model_lbol.py index 69762e8ef..dfd9e49a1 100644 --- a/nmma/em/model_lbol.py +++ b/nmma/em/model_lbol.py @@ -14,15 +14,12 @@ model_parameters_dict = { "Arnett": ["tau_m", "log10_mni"], - "Arnett_modified": ["tau_m", "log10_mni", "t_0"] + "Arnett_modified": ["tau_m", "log10_mni", "t_0"], } class SimpleBolometricLightCurveModel(object): - def __init__( - self, sample_times, - parameter_conversion=None, model="Arnett" - ): + def __init__(self, sample_times, parameter_conversion=None, model="Arnett"): """A light curve model object An object to evaluted the kilonova (with Me2017) light curve across filters @@ -60,12 +57,8 @@ def generate_lightcurve(self, sample_times, parameters): new_parameters = parameters.copy() if self.model == "Arnett": - lbol = utils_lbol.arnett_lc( - sample_times, new_parameters - ) + lbol = utils_lbol.arnett_lc(sample_times, new_parameters) elif self.model == "Arnett_modified": - lbol = utils_lbol.arnett_modified_lc( - sample_times, new_parameters - ) + lbol = utils_lbol.arnett_modified_lc(sample_times, new_parameters) return lbol diff --git a/nmma/em/multi_config_analysis.py b/nmma/em/multi_config_analysis.py index 345ae6ec7..6f8398ffa 100644 --- a/nmma/em/multi_config_analysis.py +++ b/nmma/em/multi_config_analysis.py @@ -9,8 +9,10 @@ def get_parser_here(): - parser = argparse.ArgumentParser(description="Multi config analysis script for NMMA.") - + parser = argparse.ArgumentParser( + description="Multi config analysis script for NMMA." + ) + parser.add_argument( "--config", type=str, @@ -96,7 +98,9 @@ def main(args=None): if not args.parallel: print(f"{'#'*100}") - print(f"Running analysis set: {analysis_set} with {processes} processes") + print( + f"Running analysis set: {analysis_set} with {processes} processes" + ) run_cmd_in_subprocess(cmd) print(f"{'#'*100}") else: diff --git a/nmma/em/prior.py b/nmma/em/prior.py index d344da61b..266a769d8 100644 --- a/nmma/em/prior.py +++ b/nmma/em/prior.py @@ -26,8 +26,10 @@ def from_list(self, systematics): prior[key] = val self.from_dictionary(prior) + setattr(_PriorDict, "from_list", from_list) + class ConditionalGaussianIotaGivenThetaCore(ConditionalTruncatedGaussian): """ This prior characterizes the conditional prior on the viewing angle @@ -89,6 +91,7 @@ def inclination_prior_from_fits(priors, args): import healpy as hp import numpy as np import matplotlib + matplotlib.use("agg") import matplotlib.pyplot as plt @@ -99,29 +102,31 @@ def inclination_prior_from_fits(priors, args): # check if the sky location is input # if not, the maximum posterior point is taken - if ('ra' in priors and 'dec' in priors) or (args.ra and args.dec): + if ("ra" in priors and "dec" in priors) or (args.ra and args.dec): if args.ra and args.dec: - print("Using command line input for sky location, ignoring the prior file input") + print( + "Using command line input for sky location, ignoring the prior file input" + ) ra = args.ra dec = args.dec else: print("Using prior file input for sky location") - ra = np.rad2deg(priors['ra'].peak) - dec = np.rad2deg(priors['dec'].peak) + ra = np.rad2deg(priors["ra"].peak) + dec = np.rad2deg(priors["dec"].peak) print(f"Using the input sky location ra={ra}, dec={dec}") # convert them back to theta and phi phi = np.deg2rad(ra) theta = 0.5 * np.pi - np.deg2rad(dec) # make use of the maP nside - maP_idx = np.argmax(skymap['PROBDENSITY']) - order, _ = moc.uniq2nest(skymap[maP_idx]['UNIQ']) + maP_idx = np.argmax(skymap["PROBDENSITY"]) + order, _ = moc.uniq2nest(skymap[maP_idx]["UNIQ"]) nside = hp.order2nside(order) # get the nested idx for the given sky location nest_idx = hp.ang2pix(nside, theta, phi, nest=True) # find the row with the closest nested index nest_idxs = [] for row in skymap: - order_per_row, nest_idx_per_row = moc.uniq2nest(row['UNIQ']) + order_per_row, nest_idx_per_row = moc.uniq2nest(row["UNIQ"]) if order_per_row == order: nest_idxs.append(nest_idx_per_row) else: @@ -130,8 +135,8 @@ def inclination_prior_from_fits(priors, args): row = skymap[np.argmin(np.absolute(nest_idxs - nest_idx))] else: print("Using the maP point from the fits file input") - maP_idx = np.argmax(skymap['PROBDENSITY']) - uniq_idx = skymap[maP_idx]['UNIQ'] + maP_idx = np.argmax(skymap["PROBDENSITY"]) + uniq_idx = skymap[maP_idx]["UNIQ"] # convert to nested indexing and find the location of that index order, nest_idx = moc.uniq2nest(uniq_idx) nside = hp.order2nside(order) @@ -146,13 +151,15 @@ def inclination_prior_from_fits(priors, args): # construct the iota prior cosiota_nodes_num = args.cosiota_node_num cosiota_nodes = np.cos(np.linspace(0, np.pi, cosiota_nodes_num)) - colnames = ['PROBDENSITY', 'DISTMU', 'DISTSIGMA', 'DISTNORM'] + colnames = ["PROBDENSITY", "DISTMU", "DISTSIGMA", "DISTNORM"] # do an all-in-one interpolation prob_density, dist_mu, dist_sigma, dist_norm = ( PchipInterpolator( - cosiota_nodes[::-1], row['{}_SAMPLES'.format(colname)][::-1], + cosiota_nodes[::-1], + row["{}_SAMPLES".format(colname)][::-1], ) - for colname in colnames) + for colname in colnames + ) # now have the joint distribution evaluated u = np.linspace(-1, 1, 1000) # this is cosiota # fetch the fixed distance @@ -161,8 +168,13 @@ def inclination_prior_from_fits(priors, args): dL = args.dL else: print("Using prior file input for distance") - dL = priors['luminosity_distance'].peak - prob_u = prob_density(u) * dist_norm(u) * np.square(dL) * norm(dist_mu(u), dist_sigma(u)).pdf(dL) + dL = priors["luminosity_distance"].peak + prob_u = ( + prob_density(u) + * dist_norm(u) + * np.square(dL) + * norm(dist_mu(u), dist_sigma(u)).pdf(dL) + ) iota = np.arccos(u) prob_iota = prob_u * np.absolute(np.sin(iota)) @@ -177,17 +189,15 @@ def inclination_prior_from_fits(priors, args): # normalize prob_iota /= np.trapz(iota_EM, prob_iota_EM) - priors['inclination_EM'] = Interped( - xx=iota_EM, - yy=prob_iota_EM, - minimum=0, - maximum=np.pi / 2) + priors["inclination_EM"] = Interped( + xx=iota_EM, yy=prob_iota_EM, minimum=0, maximum=np.pi / 2 + ) if args.plot: figIdx = np.random.randint(1000) plt.figure(figIdx) - plt.xlabel('Inclination') - plt.ylabel('PDF') + plt.xlabel("Inclination") + plt.ylabel("PDF") plt.plot(iota_EM, prob_iota_EM) plt.savefig(f"{args.outdir}/Fits_motivated_inclination_prior.png") @@ -196,8 +206,9 @@ def inclination_prior_from_fits(priors, args): def create_prior_from_args(model_names, args): - AnBa2022_intersect = list(set(['AnBa2022_linear', 'AnBa2022_log']) & - set(model_names)) + AnBa2022_intersect = list( + set(["AnBa2022_linear", "AnBa2022_log"]) & set(model_names) + ) if len(AnBa2022_intersect) > 0: @@ -205,7 +216,7 @@ def convert_mtot_mni(parameters): for param in ["mni", "mtot", "mrp"]: if param not in parameters: - parameters[param] = 10**parameters[f"log10_{param}"] + parameters[param] = 10 ** parameters[f"log10_{param}"] parameters["mni_c"] = parameters["mni"] / parameters["mtot"] parameters["mrp_c"] = ( @@ -221,7 +232,7 @@ def convert_mtot_mni(parameters): priors = bilby.gw.prior.PriorDict(args.prior) # setup for Ebv - if 'Ebv' not in priors: + if "Ebv" not in priors: if args.Ebv_max > 0.0 and args.use_Ebv: Ebv_c = 1.0 / (0.5 * args.Ebv_max) priors["Ebv"] = bilby.core.prior.Interped( @@ -259,9 +270,12 @@ def convert_mtot_mni(parameters): priors = inclination_prior_from_fits(priors, args) if args.fetch_Ebv_from_dustmap: - assert args.ra and args.dec, \ - "sky location is needed for fetching Ebv from dustmap" - print("Fetching value of Ebv from dustmap, overwriting the original prior on Ebv") + assert ( + args.ra and args.dec + ), "sky location is needed for fetching Ebv from dustmap" + print( + "Fetching value of Ebv from dustmap, overwriting the original prior on Ebv" + ) try: import dustmaps.sfd from dustmaps.config import config @@ -269,21 +283,23 @@ def convert_mtot_mni(parameters): except ImportError: print("Package dustmap is needed") import sys + sys.exit(1) # check if the dust fits are downloaded import os + default_dir = os.path.join(os.path.dirname(__file__)) # config dustmaps to fetch and place fits from default_dir - config['data_dir'] = default_dir - required_files = ['sfd/SFD_dust_4096_ngp.fits', 'sfd/SFD_dust_4096_sgp.fits'] - if any([ - not os.path.isfile(os.path.join(default_dir, f)) for f in required_files] + config["data_dir"] = default_dir + required_files = ["sfd/SFD_dust_4096_ngp.fits", "sfd/SFD_dust_4096_sgp.fits"] + if any( + [not os.path.isfile(os.path.join(default_dir, f)) for f in required_files] ): dustmaps.sfd.fetch() # fetching the Ebv value - coord = coordinates.SkyCoord(args.ra, args.dec, unit='rad') - priors['Ebv'] = float(dustmaps.sfd.SFDQuery()(coord)) + coord = coordinates.SkyCoord(args.ra, args.dec, unit="deg") + priors["Ebv"] = float(dustmaps.sfd.SFDQuery()(coord)) print(f"The prior on Ebv is set to fixed value of {priors['Ebv']}") diff --git a/nmma/em/svdmodel_benchmark.py b/nmma/em/svdmodel_benchmark.py index b7e392371..512f54bc3 100644 --- a/nmma/em/svdmodel_benchmark.py +++ b/nmma/em/svdmodel_benchmark.py @@ -15,10 +15,8 @@ import matplotlib.pyplot as plt matplotlib.use("Agg") -matplotlib.rcParams['text.usetex'] = not running_in_ci() -matplotlib.rcParams.update( - {"font.size": 16, "font.family": "Times New Roman"} -) +matplotlib.rcParams["text.usetex"] = not running_in_ci() +matplotlib.rcParams.update({"font.size": 16, "font.family": "Times New Roman"}) def get_parser(): diff --git a/nmma/em/systematics.py b/nmma/em/systematics.py index 2e0ec00fd..d41c2bc6d 100644 --- a/nmma/em/systematics.py +++ b/nmma/em/systematics.py @@ -58,7 +58,9 @@ def get_positional_args(cls): return params -DISTRIBUTION_PARAMETERS = {k: get_positional_args(v) for k, v in ALLOWED_DISTRIBUTIONS.items()} +DISTRIBUTION_PARAMETERS = { + k: get_positional_args(v) for k, v in ALLOWED_DISTRIBUTIONS.items() +} def load_yaml(file_path): @@ -71,9 +73,13 @@ def validate_only_one_true(yaml_dict): raise ValidationError(key, "'value' key must be present and be a boolean") true_count = sum(value["value"] for value in yaml_dict["config"].values()) if true_count > 1: - raise ValidationError("config", "Only one configuration key can be set to True at a time") + raise ValidationError( + "config", "Only one configuration key can be set to True at a time" + ) elif true_count == 0: - raise ValidationError("config", "At least one configuration key must be set to True") + raise ValidationError( + "config", "At least one configuration key must be set to True" + ) def validate_filters(filter_groups): @@ -127,7 +133,8 @@ def validate_distribution(distribution): missing_params = set(required_params) - set(distribution.keys()) if missing_params: raise ValidationError( - "distribution", f"Missing required parameters for {dist_type} distribution: {', '.join(missing_params)}" + "distribution", + f"Missing required parameters for {dist_type} distribution: {', '.join(missing_params)}", ) @@ -142,7 +149,9 @@ def create_prior_string(name, distribution): extra_params = set(params.keys()) - set(required_params) if extra_params: - warnings.warn(f"Distribution parameters {extra_params} are not used by {dist_type} distribution and will be ignored") + warnings.warn( + f"Distribution parameters {extra_params} are not used by {dist_type} distribution and will be ignored" + ) params = {k: params[k] for k in required_params if k in params} diff --git a/nmma/em/utils.py b/nmma/em/utils.py index e7da76a98..fbda1d79a 100644 --- a/nmma/em/utils.py +++ b/nmma/em/utils.py @@ -24,12 +24,14 @@ import matplotlib.pyplot as plt from nmma.em.training import SVDTrainingModel -try: - from m4opt.missions import uvex - M4OPT_INSTALLED = True -except: - M4OPT_INSTALLED = False - print("Install m4opt if you want to use uvex filters") +# FIXME : to be remove +# try: +# from m4opt.missions import uvex +# M4OPT_INSTALLED = True +# except: +# M4OPT_INSTALLED = False +# print("Install m4opt if you want to use uvex filters") + try: import afterglowpy @@ -86,6 +88,25 @@ def inner(func): "ultrasat", ] +# Global mission name storage for M4OPT mission of specific filters +# When a mission name is provided (ULTRASAT or UVEX), this global variable +# stores it so that get_default_filts_lambdas() can dynamically add the +# appropriate bandpass filters (NUV for ULTRASAT, FUV+NUV for UVEX). +# Set once at the start of analysis() via set_mission_name(), +# then accessed throughout the code via get_mission_name(). +_MISSION_NAME = None + + +def set_mission_name(mission_name): + """Store the mission name globally for M4OPT mission filter processing""" + global _MISSION_NAME + _MISSION_NAME = mission_name + + +def get_mission_name(): + """Retrieve the stored mission name to determine which filters to use""" + return _MISSION_NAME + def extinctionFactorP92SMC(nu, Ebv, z, cutoff_hi=2e16): @@ -354,15 +375,38 @@ def get_default_filts_lambdas(filters=None): filts = filts + [band.name for band in bandpasses] lambdas = np.concatenate([lambdas, [1e-10 * band.wave_eff for band in bandpasses]]) - if M4OPT_INSTALLED: - fuv_bandpass = uvex.detector.bandpasses['FUV'] - nuv_bandpass = uvex.detector.bandpasses['NUV'] - fuv_lambda = 1e-10 * fuv_bandpass.avgwave().value - nuv_lambda = 1e-10 * nuv_bandpass.avgwave().value + # To be remove + # if M4OPT_INSTALLED: + # fuv_bandpass = uvex.detector.bandpasses['FUV'] + # nuv_bandpass = uvex.detector.bandpasses['NUV'] + + # fuv_lambda = 1e-10 * fuv_bandpass.avgwave().value + # nuv_lambda = 1e-10 * nuv_bandpass.avgwave().value + + # filts = filts + ['FUV', 'NUV'] + # lambdas = np.concatenate([lambdas, [fuv_lambda, nuv_lambda]]) + + mission_name = get_mission_name() - filts = filts + ['FUV', 'NUV'] - lambdas = np.concatenate([lambdas, [fuv_lambda, nuv_lambda]]) + if mission_name: + mission = mission_name.lower() + + if mission == "ultrasat": + nuv_lambda = 1e-10 * 2599.999670463154 # ULTRASAT NUV + filts = filts + ["NUV"] + lambdas = np.concatenate([lambdas, [nuv_lambda]]) + + elif mission == "uvex": + fuv_lambda = 1e-10 * 1599.9999030773981 + nuv_lambda = 1e-10 * 2299.999825539317 + filts = filts + ["FUV", "NUV"] + lambdas = np.concatenate([lambdas, [fuv_lambda, nuv_lambda]]) + + else: + print( + f"Warning: Unknown mission '{mission_name}'. Supported missions: ULTRASAT, UVEX" + ) if filters is not None: filts_slice = [] @@ -678,16 +722,11 @@ def grb_lc(t_day, Ebv, param_dict, filters=None): # output flux density is in milliJansky try: - if isinstance(param_dict['E0'], np.ndarray): - E0 = param_dict.pop('E0') + if isinstance(param_dict["E0"], np.ndarray): + E0 = param_dict.pop("E0") vec_func = np.vectorize( - lambda i: fluxDensity( - default_time[i], - nu_0s, - E0=E0[i], - **param_dict - ), - otypes=[np.ndarray] + lambda i: fluxDensity(default_time[i], nu_0s, E0=E0[i], **param_dict), + otypes=[np.ndarray], ) mJys = vec_func(np.arange(times.shape[0])) mJys = np.stack(mJys) @@ -1780,12 +1819,12 @@ def parse_LANLfile(filename, key="band"): # the following functions are for the semi-analytic model using Hotokezaka & Nakar heating rate def heating_rate_HoNa(t, eth=0.5): """Computes the nuclear specific heating rate over time. - + This implementation is based on a model from Korobkin et al. 2012 (DOI: 10.1111/j.1365-2966.2012.21859.x), derived from nucleosynthesis simulations in compact binary merger ejecta. The model uses these parameters: eps0 = 2e18, t0 = 1.3, sig = 0.11, alpha = 1.3. - + Args: t: float or numpy.ndarray Time(s) in rest-frame to evaluate the light curve. Can be an array @@ -1793,21 +1832,22 @@ def heating_rate_HoNa(t, eth=0.5): eth: float or numpy.ndarray, default=0.5 Efficiency parameter representing the fraction of nuclear power retained in the matter, as defined by Korobkin et al. 2012. - + Returns: float or numpy.ndarray: Nuclear specific heating rate in erg/g/s (units implied but not explicitly used). """ # Define model constants eps0 = 2e18 # erg/g/s - t0 = 1.3 # s - sig = 0.11 # s + t0 = 1.3 # s + sig = 0.11 # s alpha = 1.3 # dimensionless # Calculate the time evolution term - time_term = 0.5 - 1.0 / np.pi * np.arctan((t-t0) / sig) + time_term = 0.5 - 1.0 / np.pi * np.arctan((t - t0) / sig) # Return the heating rate return eps0 * np.power(time_term, alpha) * eth / 0.5 + def luminosity_HoNa(E, t, td, be): # Calculate diffusion time ratio t_dif = td / t @@ -1818,6 +1858,7 @@ def luminosity_HoNa(E, t, td, be): # Return luminosity using complementary error function return scipy.special.erfc(ymax) * E / tesc + def dEdt_HoNa(t, E, dM, td, be): # Calculate heating contribution heat = dM * heating_rate_HoNa(t) @@ -1831,12 +1872,12 @@ def lightcurve_HoNa(t, mass, velocities, opacities, n): # check input quantity type assert ( - isinstance(mass, float) - and isinstance(velocities, list) - and all(isinstance(x, float) for x in velocities) - and isinstance(opacities, list) - and all(isinstance(x, float) for x in opacities) - ), 'Expected: mass=float, velocities/opacities=list[float]' + isinstance(mass, float) + and isinstance(velocities, list) + and all(isinstance(x, float) for x in velocities) + and isinstance(opacities, list) + and all(isinstance(x, float) for x in opacities) + ), "Expected: mass=float, velocities/opacities=list[float]" # define constants c = astropy.constants.c sigSB = astropy.constants.sigma_sb @@ -1850,13 +1891,11 @@ def lightcurve_HoNa(t, mass, velocities, opacities, n): # Validate arguments t0 = 5e-2 * astropy.units.day opacities = np.atleast_1d(opacities) - - assert np.all(t > t0), ( - f'Times must be > {t0}' - ) - assert len(velocities) == len(opacities) + 1, ( - 'len(velocities) must be len(opacities) + 1' - ) + + assert np.all(t > t0), f"Times must be > {t0}" + assert ( + len(velocities) == len(opacities) + 1 + ), "len(velocities) must be len(opacities) + 1" # convert to internal units - using vectorized operations t = t.to_value(astropy.units.s) t0 = t0.to_value(astropy.units.s) @@ -1864,68 +1903,73 @@ def lightcurve_HoNa(t, mass, velocities, opacities, n): bej = (velocities / c).to_value(astropy.units.dimensionless_unscaled) vej_0 = velocities[0].to_value(astropy.units.cm / astropy.units.s) kappas = opacities.to_value(astropy.units.cm**2 / astropy.units.g) - + # Prepare velocity shells n_shells = 100 be_0 = bej[0] be_max = bej[-1] - rho_0 = mej * (n - 3) / (4 * np.pi * vej_0**3) / (1 - (be_max/be_0)**(3 - n)) - + rho_0 = mej * (n - 3) / (4 * np.pi * vej_0**3) / (1 - (be_max / be_0) ** (3 - n)) + # Use inverse log spacing for velocity steps - simplified with direct calculation bes = be_max + be_0 - np.geomspace(be_0, be_max, n_shells) bes = np.flipud(bes)[:-1] # Flip and remove last element in one operation dbe = np.diff(np.append(bes, be_max)) # Calculate diff by appending be_max - + # Calculate optical depths more efficiently i = np.searchsorted(bej, bes) - + # Calculate power factors once for reuse - bej_power = (bej / be_0)**(1 - n) - bes_power = (bes / be_0)**(1 - n) - + bej_power = (bej / be_0) ** (1 - n) + bes_power = (bes / be_0) ** (1 - n) + # Vectorized calculation of tau_accum tau_accum = -np.cumsum((kappas * np.diff(bej_power))[::-1])[::-1] - tau_accum = np.append(tau_accum, 0) + tau_accum = np.append(tau_accum, 0) # Vectorized calculation of taus taus = tau_accum[i] + kappas[i - 1] * (bes_power - bej_power[i]) taus *= vej_0 * rho_0 / (n - 1) - + # Mass and time delay calculations - bes_power_2n = (bes / be_0)**(2 - n) # Calculate power once - dMs = 4. * np.pi * vej_0**3 * rho_0 * bes_power_2n * dbe / be_0 + bes_power_2n = (bes / be_0) ** (2 - n) # Calculate power once + dMs = 4.0 * np.pi * vej_0**3 * rho_0 * bes_power_2n * dbe / be_0 tds = taus * bes - + # Prepare arrays for solve_ivp - use broadcasting directly bes_col = bes[:, np.newaxis] tds_col = tds[:, np.newaxis] dMs_col = dMs[:, np.newaxis] - + # Evolve in time out = scipy.integrate.solve_ivp( - dEdt_HoNa, (t0, t.max()), np.zeros(len(bes)), first_step=t0, - args=(dMs_col, tds_col, bes_col), vectorized=True) - + dEdt_HoNa, + (t0, t.max()), + np.zeros(len(bes)), + first_step=t0, + args=(dMs_col, tds_col, bes_col), + vectorized=True, + ) + # Total luminosity calculation LL = luminosity_HoNa(out.y, out.t[np.newaxis, :], tds_col, bes_col).sum(0) - + # Log-log space interpolation - preserve only necessary portion log_t = np.log(out.t[1:]) log_LL = np.log(LL[1:]) - log_L_interp = interp.interp1d(log_t, log_LL, kind='cubic', assume_sorted=True) - + log_L_interp = interp.interp1d(log_t, log_LL, kind="cubic", assume_sorted=True) + # Calculate final results in vectorized operations L = np.exp(log_L_interp(np.log(t))) * (astropy.units.erg / astropy.units.s) - + # Effective radius - use vectorized log operations log_taus = np.log(taus[::-1]) log_bes = np.log(bes[::-1]) log_t_doubled = 2 * np.log(t) be = np.exp(np.interp(log_t_doubled, log_taus, log_bes)) r = be * t * (c * astropy.units.s) - + # Effective temperature - use broadcasting for squaring - T = ((L / (4 * np.pi * sigSB * r**2))**0.25).to(astropy.units.K) - + T = ((L / (4 * np.pi * sigSB * r**2)) ** 0.25).to(astropy.units.K) + # Return results return L, T, r.to(astropy.units.cm) @@ -1933,7 +1977,7 @@ def lightcurve_HoNa(t, mass, velocities, opacities, n): def running_in_ci(): """ Check if the code is running in a GitHub Actions. - + Returns: -------- bool: True if running in CI, False otherwise. diff --git a/nmma/em/utils_lbol.py b/nmma/em/utils_lbol.py index 352772c60..d17e29355 100644 --- a/nmma/em/utils_lbol.py +++ b/nmma/em/utils_lbol.py @@ -24,7 +24,7 @@ def arnett_lc_get_int_B_non_vec(x, y, s): def arnett_lc(t_day, param_dict): day = 86400.0 # in seconds ts = t_day * day - Mni = 10**param_dict["log10_mni"] + Mni = 10 ** param_dict["log10_mni"] Mni *= astropy.constants.M_sun.cgs.value # in g tau_m = param_dict["tau_m"] * day @@ -40,8 +40,11 @@ def arnett_lc(t_day, param_dict): int_A = arnett_lc_get_int_A(x, y) int_B = arnett_lc_get_int_B(x, y, s) - Ls = Mni * np.exp(-x**2) * ( - (epsilon_ni - epsilon_co) * int_A + epsilon_co * int_B) + Ls = ( + Mni + * np.exp(-(x**2)) + * ((epsilon_ni - epsilon_co) * int_A + epsilon_co * int_B) + ) return Ls @@ -49,7 +52,7 @@ def arnett_lc(t_day, param_dict): def arnett_modified_lc(t_day, param_dict): day = 86400.0 # in seconds ts = t_day * day # in seconds - Mni = 10**param_dict["log10_mni"] + Mni = 10 ** param_dict["log10_mni"] Mni *= astropy.constants.M_sun.cgs.value # in g tau_m = param_dict["tau_m"] * day t0 = param_dict["t_0"] * day @@ -66,7 +69,11 @@ def arnett_modified_lc(t_day, param_dict): int_A = arnett_lc_get_int_A(x, y) int_B = arnett_lc_get_int_B(x, y, s) - Ls = Mni * np.exp(-x**2) * ((epsilon_ni - epsilon_co) * int_A + epsilon_co * int_B) - Ls_modified = Ls * (1. - np.exp(-1. * (t0 / ts)**2)) + Ls = ( + Mni + * np.exp(-(x**2)) + * ((epsilon_ni - epsilon_co) * int_A + epsilon_co * int_B) + ) + Ls_modified = Ls * (1.0 - np.exp(-1.0 * (t0 / ts) ** 2)) return Ls_modified diff --git a/nmma/em/validate_lightcurve.py b/nmma/em/validate_lightcurve.py index 342968542..51e2b6dd8 100644 --- a/nmma/em/validate_lightcurve.py +++ b/nmma/em/validate_lightcurve.py @@ -6,6 +6,7 @@ from .utils import DEFAULT_FILTERS from .io import loadEvent + def get_parser(**kwargs): """ Arguments provided when function is called from command line @@ -47,7 +48,8 @@ def get_parser(**kwargs): ) return parser -def validate_lightcurve(data,filters=None,min_obs=3,cutoff_time=0,silent=False): + +def validate_lightcurve(data, filters=None, min_obs=3, cutoff_time=0, silent=False): """ Evaluates whether the lightcurve has the requisite user-defined number of observations in the filters provided within the defined time window from the first observation. In the case where one wants to check that at least one filter has the requisite number of observations, the function should be called multiple times with different filter arguments. @@ -70,34 +72,57 @@ def validate_lightcurve(data,filters=None,min_obs=3,cutoff_time=0,silent=False): filters_to_check = list(data.keys()) ## determine time window to consider - min_time = np.min([np.min(array[:,0]) for array in data.values()]) - max_time = min_time + cutoff_time if cutoff_time > 0 else np.max([np.max(array[:,0]) for array in data.values()]) + min_time = np.min([np.min(array[:, 0]) for array in data.values()]) + max_time = ( + min_time + cutoff_time + if cutoff_time > 0 + else np.max([np.max(array[:, 0]) for array in data.values()]) + ) ## evaluate lightcurve for each filter for filter in filters_to_check: if filter not in DEFAULT_FILTERS: raise ValueError(f"Filter {filter} not in supported filter list") elif filter not in data.keys(): - print(f"{filter} not present in data file, cannot validate") if not silent else None + print( + f"{filter} not present in data file, cannot validate" + ) if not silent else None return False - filter_data_indices = np.where(data[filter][:,0] <= max_time)[0] + filter_data_indices = np.where(data[filter][:, 0] <= max_time)[0] filter_data = [data[filter][i] for i in filter_data_indices] - + ## evaluate the number of detections - num_observations = sum(1 for value in filter_data if value[1] != np.inf and not np.isnan(value[1])) - num_detections = sum(1 for value in filter_data if value[2] != np.inf and not np.isnan(value[2]) and value[2] != 99) + num_observations = sum( + 1 for value in filter_data if value[1] != np.inf and not np.isnan(value[1]) + ) + num_detections = sum( + 1 + for value in filter_data + if value[2] != np.inf and not np.isnan(value[2]) and value[2] != 99 + ) if num_detections < min_obs: - print(f"{filter} filter has {num_detections} detections, less than the required {min_obs}") if not silent else None + print( + f"{filter} filter has {num_detections} detections, less than the required {min_obs}" + ) if not silent else None return False else: continue - print(f"Lightcurve has at least {min_obs} detections in the filters within the first {max_time-min_time} days") if not silent else None + print( + f"Lightcurve has at least {min_obs} detections in the filters within the first {max_time-min_time} days" + ) if not silent else None return True + def main(args=None): if args is None: parser = get_parser() args = parser.parse_args(args) - validation_bool = validate_lightcurve(data=args.data,filters=args.filters,min_obs=args.min_obs,cutoff_time=args.cutoff_time,silent=args.silent) - return validation_bool \ No newline at end of file + validation_bool = validate_lightcurve( + data=args.data, + filters=args.filters, + min_obs=args.min_obs, + cutoff_time=args.cutoff_time, + silent=args.silent, + ) + return validation_bool diff --git a/nmma/eos/create_injection.py b/nmma/eos/create_injection.py index 4cb9b817b..d548a8408 100644 --- a/nmma/eos/create_injection.py +++ b/nmma/eos/create_injection.py @@ -33,8 +33,9 @@ def file_to_dataframe( table = Table.read(injection_file, format="ligolw", tablename="sim_inspiral") elif injection_file.endswith(".dat"): table = Table.read(injection_file, format="csv", delimiter="\t") - elif injection_file.endswith('.ecsv'): + elif injection_file.endswith(".ecsv"): from astropy.table import Table as astro_Table + table = astro_Table.read(injection_file) else: raise ValueError("Only understand xml, dat and ecsv") @@ -77,7 +78,10 @@ def file_to_dataframe( if "geocent_end_time" in row.colnames: if "geocent_end_time_ns" in row.colnames: - injection_values["geocent_time"].append(float(row["geocent_end_time"]) + float(row["geocent_end_time_ns"]) * (10 ** -9)) + injection_values["geocent_time"].append( + float(row["geocent_end_time"]) + + float(row["geocent_end_time_ns"]) * (10**-9) + ) else: injection_values["geocent_time"].append(float(row["geocent_end_time"])) else: @@ -269,21 +273,21 @@ def get_parser(): ) parser.add_argument( "--snr-cutoff", - default=0., + default=0.0, type=float, - help="The value of SNR cutoff for the injection generation (default: 0)" + help="The value of SNR cutoff for the injection generation (default: 0)", ) parser.add_argument( "--ifos", default="H1,L1,V1", type=str, - help="Comma seperated list of ifos to be used for SNR calculation" + help="Comma seperated list of ifos to be used for SNR calculation", ) parser.add_argument( "--ifos-psd", default="", type=str, - help="Comma seperated list of psds files to be used (default: using bilby default)" + help="Comma seperated list of psds files to be used (default: using bilby default)", ) parser.add_argument("-d", "--detections-file", type=str) parser.add_argument("-i", "--indices-file", type=str) @@ -389,13 +393,23 @@ def main(args=None): ) dataframe_from_prior["timeshift"] = timeshifts - # remove the gps time from the prior dataframe + # remove the gps time from the prior dataframe # in case that is in the injection table inj_columns = set(dataframe_from_inj.columns.tolist()) - if 'geocent_time' in inj_columns: - dataframe_from_prior.drop(columns=['geocent_time',], inplace=True) - elif 'geocent_time_x' in inj_columns: - dataframe_from_prior.drop(columns=['geocent_time_x',], inplace=True) + if "geocent_time" in inj_columns: + dataframe_from_prior.drop( + columns=[ + "geocent_time", + ], + inplace=True, + ) + elif "geocent_time_x" in inj_columns: + dataframe_from_prior.drop( + columns=[ + "geocent_time_x", + ], + inplace=True, + ) prior_columns = set(dataframe_from_prior.columns.tolist()) @@ -468,14 +482,14 @@ def main(args=None): if args.eject: if args.binary_type == "BNS": ejectaFitting = BNSEjectaFitting() - dataframe = dataframe[dataframe['lambda_1'] != 0.0] - dataframe = dataframe[dataframe['lambda_2'] != 0.0] + dataframe = dataframe[dataframe["lambda_1"] != 0.0] + dataframe = dataframe[dataframe["lambda_2"] != 0.0] dataframe.reset_index(drop=True, inplace=True) elif args.binary_type == "NSBH": ejectaFitting = NSBHEjectaFitting() - dataframe = dataframe[dataframe['lambda_1'] == 0.0] - dataframe = dataframe[dataframe['lambda_2'] != 0.0] + dataframe = dataframe[dataframe["lambda_1"] == 0.0] + dataframe = dataframe[dataframe["lambda_2"] != 0.0] dataframe.reset_index(drop=True, inplace=True) else: @@ -516,12 +530,14 @@ def main(args=None): if args.snr_cutoff: print(f"Removing injection with SNR less than {args.snr_cutoff}") import bilby + bilby.core.utils.setup_logger(log_level="ERROR") - IFOs = bilby.gw.detector.InterferometerList(args.ifos.split(',')) + IFOs = bilby.gw.detector.InterferometerList(args.ifos.split(",")) if len(args.ifos_psd) > 0: from bilby.gw.detector import PowerSpectralDensity - psds = args.ifos_psd.split(',') + + psds = args.ifos_psd.split(",") for ifo, psd in zip(IFOs, psds): ifo.power_spectral_density = PowerSpectralDensity(psd_file=psd) @@ -530,7 +546,7 @@ def main(args=None): sampling_frequency=2048, frequency_domain_source_model=bilby.gw.source.lal_binary_neutron_star, parameter_conversion=bilby.gw.conversion.convert_to_lal_binary_neutron_star_parameters, - waveform_arguments = { + waveform_arguments={ "waveform_approximant": "IMRPhenomPv2_NRTidalv2", "reference_frequency": 20, "minimum_frequency": 20, # Ensure consistency with PSD @@ -541,8 +557,9 @@ def main(args=None): for _, inj in tqdm(dataframe.iterrows()): # create data just to please the bilby sanity check IFOs.set_strain_data_from_power_spectral_densities( - sampling_frequency=2048, duration=256, - start_time=inj['geocent_time'] - (256 - 2.) + sampling_frequency=2048, + duration=256, + start_time=inj["geocent_time"] - (256 - 2.0), ) # Inject the waveform into the detectors snr = 0.0 @@ -552,10 +569,10 @@ def main(args=None): parameters=inj.to_dict(), raise_error=False, ) - snr += ifo.meta_data['optimal_SNR'] + snr += ifo.meta_data["optimal_SNR"] snr_list.append(np.sqrt(snr)) index_taken = np.where(np.array(snr_list) >= args.snr_cutoff)[0] - dataframe['optimal_network_snr'] = np.array(snr_list) + dataframe["optimal_network_snr"] = np.array(snr_list) dataframe = dataframe.take(index_taken) print("{0} injections left after snr filtering".format(len(index_taken))) @@ -563,7 +580,7 @@ def main(args=None): # re-index the simulation_id if not args.injection_file: print("Re-indexing the simulation_id") - dataframe['simulation_id'] = np.arange(len(dataframe)) + dataframe["simulation_id"] = np.arange(len(dataframe)) # dump the whole thing back into a json injection file injection_creator.write_injection_dataframe( dataframe, args.filename, args.extension diff --git a/nmma/eos/eos.py b/nmma/eos/eos.py index 7c637ee44..d7d0d34c0 100644 --- a/nmma/eos/eos.py +++ b/nmma/eos/eos.py @@ -22,31 +22,40 @@ class EOS_with_CSE(object): seed: int, seed of random draw extension (default: 42) """ - def __init__(self, low_density_eos, n_connect=0.16, n_lim=2., N_seg=5, cs2_limit=1., seed=42, extension_scheme='peter', low_density_eos_stiff=None): + def __init__( + self, + low_density_eos, + n_connect=0.16, + n_lim=2.0, + N_seg=5, + cs2_limit=1.0, + seed=42, + extension_scheme="peter", + low_density_eos_stiff=None, + ): self.seed = seed if not low_density_eos_stiff: - self.n_low = low_density_eos['n'] - self.p_low = low_density_eos['p'] - self.e_low = low_density_eos['e'] + self.n_low = low_density_eos["n"] + self.p_low = low_density_eos["p"] + self.e_low = low_density_eos["e"] else: - assert len(low_density_eos) == len(low_density_eos_stiff), ('This requires ' - 'interpolation. ' - 'Will be added ' - 'in the future.') + assert len(low_density_eos) == len(low_density_eos_stiff), ( + "This requires " "interpolation. " "Will be added " "in the future." + ) # fix the seed np.random.seed(self.seed) alpha = np.random.uniform() - self.n_low_soft = low_density_eos['n'] - self.p_low_soft = low_density_eos['p'] - self.e_low_soft = low_density_eos['e'] + self.n_low_soft = low_density_eos["n"] + self.p_low_soft = low_density_eos["p"] + self.e_low_soft = low_density_eos["e"] - self.n_low_stiff = low_density_eos_stiff['n'] - self.p_low_stiff = low_density_eos_stiff['p'] - self.e_low_stiff = low_density_eos_stiff['e'] + self.n_low_stiff = low_density_eos_stiff["n"] + self.p_low_stiff = low_density_eos_stiff["p"] + self.e_low_stiff = low_density_eos_stiff["e"] diff_e = self.e_low_stiff - self.e_low_soft diff_p = self.p_low_stiff - self.p_low_soft @@ -61,8 +70,11 @@ def __init__(self, low_density_eos, n_connect=0.16, n_lim=2., N_seg=5, cs2_limit self.e_at_n_connect = np.exp(log_e_of_log_n_low(np.log(n_connect))) self.p_at_n_connect = np.exp(log_p_of_log_n_low(np.log(n_connect))) - self.cs2_at_n_connect = self.p_at_n_connect / self.e_at_n_connect *\ - log_p_of_log_e_low.derivative()(np.log(self.e_at_n_connect)) + self.cs2_at_n_connect = ( + self.p_at_n_connect + / self.e_at_n_connect + * log_p_of_log_e_low.derivative()(np.log(self.e_at_n_connect)) + ) self.n_connect = n_connect self.n_lim = n_lim @@ -70,9 +82,9 @@ def __init__(self, low_density_eos, n_connect=0.16, n_lim=2., N_seg=5, cs2_limit self.N_seg = N_seg self.cs2_limit = cs2_limit - if extension_scheme == 'peter': + if extension_scheme == "peter": self.__extend() - elif extension_scheme == 'rahul': + elif extension_scheme == "rahul": self.__extend_v1() self.__calculate_pseudo_enthalpy() @@ -91,18 +103,21 @@ def __extend(self): # draw cs2(n_node) randomly for node_index in range(1, self.N_seg + 1): n_val_lower_bound = cs2_draw[node_index - 1, 0] - n_val_upper_bound = min(n_val_lower_bound + 1.5 * self.n_extend_range / self.N_seg, self.n_lim) + n_val_upper_bound = min( + n_val_lower_bound + 1.5 * self.n_extend_range / self.N_seg, self.n_lim + ) n_val = np.random.uniform(n_val_lower_bound, n_val_upper_bound) - cs2_val = np.random.uniform(0., self.cs2_limit) + cs2_val = np.random.uniform(0.0, self.cs2_limit) cs2_draw[node_index] = [n_val, cs2_val] # add the speed for sound value at n_lim - cs2_at_n_lim = np.random.uniform(0., self.cs2_limit) + cs2_at_n_lim = np.random.uniform(0.0, self.cs2_limit) cs2_draw[-1] = [self.n_lim, cs2_at_n_lim] # interpolation for cs2(n) - cs2_extent = interp1d(cs2_draw[:, 0], cs2_draw[:, 1], kind='linear', - fill_value="extrapolate") + cs2_extent = interp1d( + cs2_draw[:, 0], cs2_draw[:, 1], kind="linear", fill_value="extrapolate" + ) # construct the extended EOS # do the integration in log-space for stability @@ -120,15 +135,23 @@ def dydt(t, y, cs2): # dedn = (p + e) / n # dpdn = cs2(n) * dedn - dloge_dlogn = 1. + p / e - dlogp_dlogn = cs2(n) * (e / p + 1.) + dloge_dlogn = 1.0 + p / e + dlogp_dlogn = cs2(n) * (e / p + 1.0) return [dlogp_dlogn, dloge_dlogn] y0 = (np.log(self.p_at_n_connect), np.log(self.e_at_n_connect)) - sol = solve_ivp(dydt, (np.log(self.n_connect), np.log(self.n_lim)), y0, - method='DOP853', t_eval=np.log(n_high), args=(cs2_extent,), rtol=1e-7, atol=0.) + sol = solve_ivp( + dydt, + (np.log(self.n_connect), np.log(self.n_lim)), + y0, + method="DOP853", + t_eval=np.log(n_high), + args=(cs2_extent,), + rtol=1e-7, + atol=0.0, + ) logp_high = sol.y[0] loge_high = sol.y[1] @@ -146,8 +169,9 @@ def __extend_v1(self): # fix the seed np.random.seed(self.seed) - n_ext_grid = np.linspace(self.n_connect + 1e-4 * self.n_connect, self.n_lim, - num=self.N_seg + 1) + n_ext_grid = np.linspace( + self.n_connect + 1e-4 * self.n_connect, self.n_lim, num=self.N_seg + 1 + ) c2_ext_grid = [np.random.uniform(0, self.cs2_limit) for i in n_ext_grid] c2_ext_grid[0] = self.cs2_at_n_connect @@ -156,8 +180,10 @@ def __extend_v1(self): mu_ext_grid[0] = self.mu_at_n_connect num = 50 - n_high = [np.linspace(n_ext_grid[i], n_ext_grid[i + 1], endpoint=False, num=num) - for i in range(n_ext_grid.size - 1)] + n_high = [ + np.linspace(n_ext_grid[i], n_ext_grid[i + 1], endpoint=False, num=num) + for i in range(n_ext_grid.size - 1) + ] n_high = np.array(n_high) c2_high = np.zeros_like(n_high) mu_high = np.zeros_like(n_high) @@ -165,16 +191,37 @@ def __extend_v1(self): # Integrates the sound speed to compute the chemical potential # Fills in all elements of mu_ext for i in range(n_ext_grid.size - 1): - slope = (c2_ext_grid[i + 1] - c2_ext_grid[i]) / (n_ext_grid[i + 1] - n_ext_grid[i]) + slope = (c2_ext_grid[i + 1] - c2_ext_grid[i]) / ( + n_ext_grid[i + 1] - n_ext_grid[i] + ) c2_high[i, :] = slope * (n_high[i, :] - n_ext_grid[i]) + c2_ext_grid[i] - mu_ext_grid[i + 1] = mu_ext_grid[i] * np.exp(slope * (n_ext_grid[i + 1] - n_ext_grid[i] - n_ext_grid[i] * np.log(n_ext_grid[i + 1] / n_ext_grid[i])) + c2_ext_grid[i] * np.log(n_ext_grid[i + 1] / n_ext_grid[i])) - mu_high[i, :] = mu_ext_grid[i] * np.exp(slope * (n_high[i, :] - n_ext_grid[i] - n_ext_grid[i] * np.log(n_high[i, :] / n_ext_grid[i])) + c2_ext_grid[i] * np.log(n_high[i, :] / n_ext_grid[i])) + mu_ext_grid[i + 1] = mu_ext_grid[i] * np.exp( + slope + * ( + n_ext_grid[i + 1] + - n_ext_grid[i] + - n_ext_grid[i] * np.log(n_ext_grid[i + 1] / n_ext_grid[i]) + ) + + c2_ext_grid[i] * np.log(n_ext_grid[i + 1] / n_ext_grid[i]) + ) + mu_high[i, :] = mu_ext_grid[i] * np.exp( + slope + * ( + n_high[i, :] + - n_ext_grid[i] + - n_ext_grid[i] * np.log(n_high[i, :] / n_ext_grid[i]) + ) + + c2_ext_grid[i] * np.log(n_high[i, :] / n_ext_grid[i]) + ) self.n_high = n_high.flatten() self.c2_high = c2_high.flatten() self.mu_high = mu_high.flatten() - self.e_high = cumulative_trapezoid(self.mu_high, self.n_high, initial=0) + self.e_at_n_connect + self.e_high = ( + cumulative_trapezoid(self.mu_high, self.n_high, initial=0) + + self.e_at_n_connect + ) self.p_high = self.mu_high * self.n_high - self.e_high n_low = self.n_low @@ -185,67 +232,108 @@ def __extend_v1(self): def __calculate_pseudo_enthalpy(self): intergrand = self.p_array / (self.e_array + self.p_array) - self.h_array = cumulative_trapezoid(intergrand, np.log(self.p_array), initial=0) + intergrand[0] + self.h_array = ( + cumulative_trapezoid(intergrand, np.log(self.p_array), initial=0) + + intergrand[0] + ) def __construct_all_interpolation(self): - self.log_energy_density_from_log_pressure = interp1d(np.log(self.p_array), - np.log(self.e_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - self.log_energy_density_from_log_pseudo_enthalpy = interp1d(np.log(self.h_array), - np.log(self.e_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - self.log_energy_density_from_log_number_density = interp1d(np.log(self.n_array), - np.log(self.e_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - - self.log_pressure_from_log_energy_density = interp1d(np.log(self.e_array), - np.log(self.p_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - self.log_pressure_from_log_number_density = interp1d(np.log(self.n_array), - np.log(self.p_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - self.log_pressure_from_log_pseudo_enthalpy = interp1d(np.log(self.h_array), - np.log(self.p_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - - self.log_number_density_from_log_pressure = interp1d(np.log(self.p_array), - np.log(self.n_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - self.log_number_density_from_log_pseudo_enthalpy = interp1d(np.log(self.h_array), - np.log(self.n_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - self.log_number_density_from_log_energy_density = interp1d(np.log(self.e_array), - np.log(self.n_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - - self.log_pseudo_enthalpy_from_log_pressure = interp1d(np.log(self.p_array), - np.log(self.h_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - self.log_pseudo_enthalpy_from_log_energy_density = interp1d(np.log(self.e_array), - np.log(self.h_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - self.log_pseudo_enthalpy_from_log_number_density = interp1d(np.log(self.n_array), - np.log(self.h_array), kind='linear', - fill_value='extrapolate', - assume_sorted=True) - - self.log_dedp_from_log_pressure = interp1d(np.log(self.p_array), - np.gradient(np.log(self.e_array), np.log(self.p_array)), - kind='linear', - fill_value='extrapolate', - assume_sorted=True) + self.log_energy_density_from_log_pressure = interp1d( + np.log(self.p_array), + np.log(self.e_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + self.log_energy_density_from_log_pseudo_enthalpy = interp1d( + np.log(self.h_array), + np.log(self.e_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + self.log_energy_density_from_log_number_density = interp1d( + np.log(self.n_array), + np.log(self.e_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + + self.log_pressure_from_log_energy_density = interp1d( + np.log(self.e_array), + np.log(self.p_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + self.log_pressure_from_log_number_density = interp1d( + np.log(self.n_array), + np.log(self.p_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + self.log_pressure_from_log_pseudo_enthalpy = interp1d( + np.log(self.h_array), + np.log(self.p_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + + self.log_number_density_from_log_pressure = interp1d( + np.log(self.p_array), + np.log(self.n_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + self.log_number_density_from_log_pseudo_enthalpy = interp1d( + np.log(self.h_array), + np.log(self.n_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + self.log_number_density_from_log_energy_density = interp1d( + np.log(self.e_array), + np.log(self.n_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + + self.log_pseudo_enthalpy_from_log_pressure = interp1d( + np.log(self.p_array), + np.log(self.h_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + self.log_pseudo_enthalpy_from_log_energy_density = interp1d( + np.log(self.e_array), + np.log(self.h_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + self.log_pseudo_enthalpy_from_log_number_density = interp1d( + np.log(self.n_array), + np.log(self.h_array), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) + + self.log_dedp_from_log_pressure = interp1d( + np.log(self.p_array), + np.gradient(np.log(self.e_array), np.log(self.p_array)), + kind="linear", + fill_value="extrapolate", + assume_sorted=True, + ) def energy_density_from_pressure(self, p): return np.exp(self.log_energy_density_from_log_pressure(np.log(p))) @@ -316,9 +404,10 @@ def construct_family(self, ndat=100): if i != ndat - 1: # build a interpolation for logpc-mass to get logpc at max mass - f = interp1d(logpcs, -ms, kind='linear') - res = minimize_scalar(f, method='bounded', - bounds=(logpcs[0] * 1.001, logpcs[-1] * 0.999)) + f = interp1d(logpcs, -ms, kind="linear") + res = minimize_scalar( + f, method="bounded", bounds=(logpcs[0] * 1.001, logpcs[-1] * 0.999) + ) logpmax = res.x pmax = np.exp(logpmax) mmax, rmax, kmax = TOVSolver(self, pmax) @@ -338,12 +427,12 @@ def construct_family(self, ndat=100): rs /= 1e3 # calculate the tidal deformability - lambdas = 2. / 3. * ks * np.power(cs, -5.) + lambdas = 2.0 / 3.0 * ks * np.power(cs, -5.0) # build the mass-lambda interpolation - self.lambda_m_interp = interp1d(ms, lambdas, kind='linear') + self.lambda_m_interp = interp1d(ms, lambdas, kind="linear") # build the mass-radius interpolation - self.radius_m_interp = interp1d(ms, rs, kind='linear') + self.radius_m_interp = interp1d(ms, rs, kind="linear") return diff --git a/nmma/eos/utils.py b/nmma/eos/utils.py index 0a951717c..b9f1dec8e 100644 --- a/nmma/eos/utils.py +++ b/nmma/eos/utils.py @@ -9,28 +9,33 @@ class PulsarConstraint(object): - def __init__(self, pulsar_masses, pulsar_masses_error, EOSpath, Neos): - self.pulsar_masses = pulsar_masses # list of pulsar masses - self.pulsar_masses_error = pulsar_masses_error # list of errors of pulsar masses - self.EOSpath = EOSpath # path of EOS data files - self.Neos = Neos # number of EOS files + self.pulsar_masses = pulsar_masses # list of pulsar masses + self.pulsar_masses_error = ( + pulsar_masses_error # list of errors of pulsar masses + ) + self.EOSpath = EOSpath # path of EOS data files + self.Neos = Neos # number of EOS files if len(self.pulsar_masses) != len(self.pulsar_masses_error): - raise ValueError('Number of masses and mass errors are inconsistent') + raise ValueError("Number of masses and mass errors are inconsistent") def __repr__(self): - return self.__class__.__name__ + ' with {} pulsars of masses {} and errors of {}'\ - .format(len(self.pulsar_masses), self.pulsar_masses, self.pulsar_masses_error) + return ( + self.__class__.__name__ + + " with {} pulsars of masses {} and errors of {}".format( + len(self.pulsar_masses), self.pulsar_masses, self.pulsar_masses_error + ) + ) def pulsar_logweight(self, pulsar_masses, pulsar_masses_error, EOSpath, Neos): - + logweight = [] - for EOSIdx in range(1,Neos+1): - EOSdata = '{0}/{1}.dat'.format(EOSpath, EOSIdx) - R, M = np.loadtxt(EOSdata, usecols=[0,1], unpack=True) + for EOSIdx in range(1, Neos + 1): + EOSdata = "{0}/{1}.dat".format(EOSpath, EOSIdx) + R, M = np.loadtxt(EOSdata, usecols=[0, 1], unpack=True) mMax = M[-1] - logl = 0. + logl = 0.0 for i in range(len(pulsar_masses)): mass = pulsar_masses[i] error = pulsar_masses_error[i] @@ -39,131 +44,173 @@ def pulsar_logweight(self, pulsar_masses, pulsar_masses_error, EOSpath, Neos): logweight -= scipy.special.logsumexp(logweight) pulsar_weight = np.exp(logweight) - np.savetxt('pulsar_weight.dat', np.c_[pulsar_weight]) + np.savetxt("pulsar_weight.dat", np.c_[pulsar_weight]) return pulsar_weight -class MTOVUpperConstraint(object): +class MTOVUpperConstraint(object): def __init__(self, maxMTOV, maxMTOV_error, EOSpath, Neos): - self.maxMTOV = maxMTOV # list of maximum TOV masses - self.maxMTOV_error = maxMTOV_error # list of maximum TOV mass errors - self.EOSpath = EOSpath # path of EOS data files - self.Neos = Neos # number of EOS files + self.maxMTOV = maxMTOV # list of maximum TOV masses + self.maxMTOV_error = maxMTOV_error # list of maximum TOV mass errors + self.EOSpath = EOSpath # path of EOS data files + self.Neos = Neos # number of EOS files def __repr__(self): - return self.__class__.__name__ + ' maxTOV of {} and error of {}'\ - .format(self.maxMTOV, self.maxMTOV_error) + return self.__class__.__name__ + " maxTOV of {} and error of {}".format( + self.maxMTOV, self.maxMTOV_error + ) def maxTOVmass_logweight(self, maxMTOV, maxMTOV_error, EOSpath, Neos): - + logweight = [] - for EOSIdx in range(1,Neos+1): - EOSdata = '{0}/{1}.dat'.format(EOSpath, EOSIdx) - R, M = np.loadtxt(EOSdata, usecols=[0,1], unpack=True) + for EOSIdx in range(1, Neos + 1): + EOSdata = "{0}/{1}.dat".format(EOSpath, EOSIdx) + R, M = np.loadtxt(EOSdata, usecols=[0, 1], unpack=True) mMax = M[-1] - logl = 0. + logl = 0.0 for i in range(len(maxMTOV)): mass = maxMTOV[i] error = maxMTOV_error[i] logl += scipy.stats.norm.logsf(mMax, loc=mass, scale=error) logweight.append(logl) - + logweight -= scipy.special.logsumexp(logweight) maxTOVmass_weight = np.exp(logweight) - np.savetxt('maxTOVmass_weight.dat', np.c_[maxTOVmass_weight]) + np.savetxt("maxTOVmass_weight.dat", np.c_[maxTOVmass_weight]) return maxTOVmass_weight -class NICERConstraint(object): +class NICERConstraint(object): def __init__(self, NICER_path, EOSpath, Neos): - self.NICER_path = NICER_path # path with NICER data files - self.EOSpath = EOSpath # path of EOS data files - self.Neos = Neos # number of EOS files + self.NICER_path = NICER_path # path with NICER data files + self.EOSpath = EOSpath # path of EOS data files + self.Neos = Neos # number of EOS files def __repr__(self): - return self.__class__.__name__ + ' with the observational NICER data taken from {}'\ - .format(self.NICER_path) - + return ( + self.__class__.__name__ + + " with the observational NICER data taken from {}".format(self.NICER_path) + ) + def NICER_logweight(self, NICER_path, EOSpath, Neos): - - #NICER I - NICER_radius, NICER_mass = np.loadtxt('{0}/J0030_3spot_RM.txt'.format(NICER_path), skiprows=1, usecols=[0, 1], unpack=True) + + # NICER I + NICER_radius, NICER_mass = np.loadtxt( + "{0}/J0030_3spot_RM.txt".format(NICER_path), + skiprows=1, + usecols=[0, 1], + unpack=True, + ) NICER_radius = NICER_radius[::100] NICER_mass = NICER_mass[::100] - NICER_KDE = scipy.stats.gaussian_kde((NICER_radius,NICER_mass)) - - #NICER II - NICER_radius, NICER_mass = np.loadtxt('{0}/NICER_x_XMM_J0740_XPSI_STU_NSX_FIH_radius_mass.txt'.format(NICER_path), skiprows=1, usecols=[0, 1], unpack=True) + NICER_KDE = scipy.stats.gaussian_kde((NICER_radius, NICER_mass)) + + # NICER II + NICER_radius, NICER_mass = np.loadtxt( + "{0}/NICER_x_XMM_J0740_XPSI_STU_NSX_FIH_radius_mass.txt".format(NICER_path), + skiprows=1, + usecols=[0, 1], + unpack=True, + ) NICER_radius = NICER_radius[::42] NICER_mass = NICER_mass[::42] - NICER_KDE_XPSI = scipy.stats.gaussian_kde((NICER_radius,NICER_mass)) - - NICER_radius, NICER_mass = np.loadtxt('{0}/NICER+XMM_J0740_RM.txt'.format(NICER_path), skiprows=1, usecols=[0, 1], unpack=True) + NICER_KDE_XPSI = scipy.stats.gaussian_kde((NICER_radius, NICER_mass)) + + NICER_radius, NICER_mass = np.loadtxt( + "{0}/NICER+XMM_J0740_RM.txt".format(NICER_path), + skiprows=1, + usecols=[0, 1], + unpack=True, + ) NICER_radius = NICER_radius[::160] NICER_mass = NICER_mass[::160] - NICER_KDE_Maryland = scipy.stats.gaussian_kde((NICER_radius,NICER_mass)) + NICER_KDE_Maryland = scipy.stats.gaussian_kde((NICER_radius, NICER_mass)) logweight = [] for EOSIdx in tqdm(range(1, Neos + 1)): - EOSdata = '{0}/{1}.dat'.format(EOSpath, EOSIdx) - R, M = np.loadtxt(EOSdata, usecols=[0,1], unpack=True) + EOSdata = "{0}/{1}.dat".format(EOSpath, EOSIdx) + R, M = np.loadtxt(EOSdata, usecols=[0, 1], unpack=True) mMax = M[-1] - + # NICER I - R_M_interp = scipy.interpolate.interp1d(M, R, kind='linear') - Mrange = np.linspace(1., 2.0, num=100) + R_M_interp = scipy.interpolate.interp1d(M, R, kind="linear") + Mrange = np.linspace(1.0, 2.0, num=100) Mrange = Mrange[Mrange < mMax] - NICER_logpdf = scipy.special.logsumexp(NICER_KDE.logpdf((R_M_interp(Mrange),Mrange))) + NICER_logpdf = scipy.special.logsumexp( + NICER_KDE.logpdf((R_M_interp(Mrange), Mrange)) + ) # NICER II Mrange = np.linspace(1.5, 2.5, num=100) Mrange = Mrange[Mrange < mMax] - NICER_XPSI_logpdf = scipy.special.logsumexp(NICER_KDE_XPSI.logpdf((R_M_interp(Mrange),Mrange))) - NICER_Maryland_logpdf = scipy.special.logsumexp(NICER_KDE_Maryland.logpdf((R_M_interp(Mrange),Mrange))) - - NICERII_logpdf = scipy.special.logsumexp([NICER_XPSI_logpdf, NICER_Maryland_logpdf]) # adding the likelihood from XPSI and Maryland (therefore, an average) - - logweight.append(NICER_logpdf + NICERII_logpdf) + NICER_XPSI_logpdf = scipy.special.logsumexp( + NICER_KDE_XPSI.logpdf((R_M_interp(Mrange), Mrange)) + ) + NICER_Maryland_logpdf = scipy.special.logsumexp( + NICER_KDE_Maryland.logpdf((R_M_interp(Mrange), Mrange)) + ) + + NICERII_logpdf = scipy.special.logsumexp( + [NICER_XPSI_logpdf, NICER_Maryland_logpdf] + ) # adding the likelihood from XPSI and Maryland (therefore, an average) + + logweight.append(NICER_logpdf + NICERII_logpdf) logweight -= scipy.special.logsumexp(logweight) NICER_weight = np.exp(logweight) - np.savetxt('NICER_weight.dat', np.c_[NICER_weight]) + np.savetxt("NICER_weight.dat", np.c_[NICER_weight]) return NICER_weight class JointConstraint(object): + def __init__( + self, + EOSpath, + Neos, + pulsar_masses, + pulsar_masses_error, + pulsar_weight, + maxMTOV, + maxMTOV_error, + maxMTOV_weight, + NICER_path, + NICER_weight, + ): + + self.EOSpath = EOSpath # path of EOS data files + self.Neos = Neos # number of EOS files + self.pulsar_masses = pulsar_masses # list of pulsar masses + self.pulsar_masses_error = ( + pulsar_masses_error # list of errors of pulsar masses + ) + self.pulsar_weight = pulsar_weight # pulsar weights computed with JointConstraint.pulsar_logweight() + self.maxMTOV = maxMTOV # list of maximum TOV masses + self.maxMTOV_error = maxMTOV_error # list of maximum TOV mass errors + self.maxMTOV_weight = maxMTOV_weight # maximum TOV mass weights computed with JointConstraint.maxTOVmass_logweight() + self.NICER_path = NICER_path # path with NICER data files + self.NICER_weight = NICER_weight # NICER weights computed with JointConstraint.NICER_logweight() - def __init__(self, EOSpath, Neos, pulsar_masses, pulsar_masses_error, pulsar_weight, maxMTOV, maxMTOV_error, maxMTOV_weight, NICER_path, NICER_weight): - - self.EOSpath = EOSpath # path of EOS data files - self.Neos = Neos # number of EOS files - self.pulsar_masses = pulsar_masses # list of pulsar masses - self.pulsar_masses_error = pulsar_masses_error # list of errors of pulsar masses - self.pulsar_weight = pulsar_weight # pulsar weights computed with JointConstraint.pulsar_logweight() - self.maxMTOV = maxMTOV # list of maximum TOV masses - self.maxMTOV_error = maxMTOV_error # list of maximum TOV mass errors - self.maxMTOV_weight = maxMTOV_weight # maximum TOV mass weights computed with JointConstraint.maxTOVmass_logweight() - self.NICER_path = NICER_path # path with NICER data files - self.NICER_weight = NICER_weight # NICER weights computed with JointConstraint.NICER_logweight() - def __repr__(self): - return self.__class__.__name__ + ' for pulsar, maximum TOV mass and NICER weights computation. NICER data taken from {}'\ - .format(self.NICER_path) - + return ( + self.__class__.__name__ + + " for pulsar, maximum TOV mass and NICER weights computation. NICER data taken from {}".format( + self.NICER_path + ) + ) def pulsar_logweight(self, pulsar_masses, pulsar_masses_error, EOSpath, Neos): - + logweight = [] - for EOSIdx in range(1,Neos+1): - EOSdata = '{0}/{1}.dat'.format(EOSpath, EOSIdx) - R, M = np.loadtxt(EOSdata, usecols=[0,1], unpack=True) + for EOSIdx in range(1, Neos + 1): + EOSdata = "{0}/{1}.dat".format(EOSpath, EOSIdx) + R, M = np.loadtxt(EOSdata, usecols=[0, 1], unpack=True) mMax = M[-1] - logl = 0. + logl = 0.0 for i in range(len(pulsar_masses)): mass = pulsar_masses[i] error = pulsar_masses_error[i] @@ -173,14 +220,14 @@ def pulsar_logweight(self, pulsar_masses, pulsar_masses_error, EOSpath, Neos): return logweight def maxTOVmass_logweight(self, maxMTOV, maxMTOV_error, EOSpath, Neos): - + logweight = [] - for EOSIdx in range(1,Neos+1): - EOSdata = '{0}/{1}.dat'.format(EOSpath, EOSIdx) - R, M = np.loadtxt(EOSdata, usecols=[0,1], unpack=True) + for EOSIdx in range(1, Neos + 1): + EOSdata = "{0}/{1}.dat".format(EOSpath, EOSIdx) + R, M = np.loadtxt(EOSdata, usecols=[0, 1], unpack=True) mMax = M[-1] - logl = 0. + logl = 0.0 for i in range(len(maxMTOV)): mass = maxMTOV[i] error = maxMTOV_error[i] @@ -190,57 +237,82 @@ def maxTOVmass_logweight(self, maxMTOV, maxMTOV_error, EOSpath, Neos): return logweight def NICER_logweight(self, NICER_path, EOSpath, Neos): - - #NICER I - NICER_radius, NICER_mass = np.loadtxt('{0}/J0030_3spot_RM.txt'.format(NICER_path), skiprows=1, usecols=[0, 1], unpack=True) + + # NICER I + NICER_radius, NICER_mass = np.loadtxt( + "{0}/J0030_3spot_RM.txt".format(NICER_path), + skiprows=1, + usecols=[0, 1], + unpack=True, + ) NICER_radius = NICER_radius[::100] NICER_mass = NICER_mass[::100] - NICER_KDE = scipy.stats.gaussian_kde((NICER_radius,NICER_mass)) - - #NICER II - NICER_radius, NICER_mass = np.loadtxt('{0}/NICER_x_XMM_J0740_XPSI_STU_NSX_FIH_radius_mass.txt'.format(NICER_path), skiprows=1, usecols=[0, 1], unpack=True) + NICER_KDE = scipy.stats.gaussian_kde((NICER_radius, NICER_mass)) + + # NICER II + NICER_radius, NICER_mass = np.loadtxt( + "{0}/NICER_x_XMM_J0740_XPSI_STU_NSX_FIH_radius_mass.txt".format(NICER_path), + skiprows=1, + usecols=[0, 1], + unpack=True, + ) NICER_radius = NICER_radius[::42] NICER_mass = NICER_mass[::42] - NICER_KDE_XPSI = scipy.stats.gaussian_kde((NICER_radius,NICER_mass)) - - NICER_radius, NICER_mass = np.loadtxt('{0}/NICER+XMM_J0740_RM.txt'.format(NICER_path), skiprows=1, usecols=[0, 1], unpack=True) + NICER_KDE_XPSI = scipy.stats.gaussian_kde((NICER_radius, NICER_mass)) + + NICER_radius, NICER_mass = np.loadtxt( + "{0}/NICER+XMM_J0740_RM.txt".format(NICER_path), + skiprows=1, + usecols=[0, 1], + unpack=True, + ) NICER_radius = NICER_radius[::160] NICER_mass = NICER_mass[::160] - NICER_KDE_Maryland = scipy.stats.gaussian_kde((NICER_radius,NICER_mass)) + NICER_KDE_Maryland = scipy.stats.gaussian_kde((NICER_radius, NICER_mass)) logweight = [] for EOSIdx in tqdm(range(1, Neos + 1)): - EOSdata = '{0}/{1}.dat'.format(EOSpath, EOSIdx) - R, M = np.loadtxt(EOSdata, usecols=[0,1], unpack=True) + EOSdata = "{0}/{1}.dat".format(EOSpath, EOSIdx) + R, M = np.loadtxt(EOSdata, usecols=[0, 1], unpack=True) mMax = M[-1] - + # NICER I - R_M_interp = scipy.interpolate.interp1d(M, R, kind='linear') - Mrange = np.linspace(1., 2.0, num=100) + R_M_interp = scipy.interpolate.interp1d(M, R, kind="linear") + Mrange = np.linspace(1.0, 2.0, num=100) Mrange = Mrange[Mrange < mMax] - NICER_logpdf = scipy.special.logsumexp(NICER_KDE.logpdf((R_M_interp(Mrange),Mrange))) + NICER_logpdf = scipy.special.logsumexp( + NICER_KDE.logpdf((R_M_interp(Mrange), Mrange)) + ) # NICER II Mrange = np.linspace(1.5, 2.5, num=100) Mrange = Mrange[Mrange < mMax] - NICER_XPSI_logpdf = scipy.special.logsumexp(NICER_KDE_XPSI.logpdf((R_M_interp(Mrange),Mrange))) - NICER_Maryland_logpdf = scipy.special.logsumexp(NICER_KDE_Maryland.logpdf((R_M_interp(Mrange),Mrange))) - - NICERII_logpdf = scipy.special.logsumexp([NICER_XPSI_logpdf, NICER_Maryland_logpdf]) # adding the likelihood from XPSI and Maryland (therefore, an average) - - logweight.append(NICER_logpdf + NICERII_logpdf) + NICER_XPSI_logpdf = scipy.special.logsumexp( + NICER_KDE_XPSI.logpdf((R_M_interp(Mrange), Mrange)) + ) + NICER_Maryland_logpdf = scipy.special.logsumexp( + NICER_KDE_Maryland.logpdf((R_M_interp(Mrange), Mrange)) + ) + + NICERII_logpdf = scipy.special.logsumexp( + [NICER_XPSI_logpdf, NICER_Maryland_logpdf] + ) # adding the likelihood from XPSI and Maryland (therefore, an average) + + logweight.append(NICER_logpdf + NICERII_logpdf) return logweight - def total_logweight(self, EOSpath, Neos, pulsar_weight, maxMTOV_weight, NICER_weight): - #eos weighting with all information: pulsar, mtov, NICER + def total_logweight( + self, EOSpath, Neos, pulsar_weight, maxMTOV_weight, NICER_weight + ): + # eos weighting with all information: pulsar, mtov, NICER total_weight = list(map(add, pulsar_weight, maxMTOV_weight)) total_weight = list(map(add, total_weight, NICER_weight)) total_weight -= scipy.special.logsumexp(total_weight) weight = np.exp(total_weight) - np.savetxt('joint_weight.dat', np.c_[weight]) - + np.savetxt("joint_weight.dat", np.c_[weight]) + return weight def pulsar_mtov_weight(self, EOSpath, Neos, pulsar_weight, maxMTOV_weight): @@ -248,25 +320,27 @@ def pulsar_mtov_weight(self, EOSpath, Neos, pulsar_weight, maxMTOV_weight): total_weight = list(map(add, pulsar_weight, maxMTOV_weight)) total_weight -= scipy.special.logsumexp(total_weight) weight = np.exp(total_weight) - np.savetxt('pulsar_mtov_weight.dat', np.c_[weight]) - + np.savetxt("pulsar_mtov_weight.dat", np.c_[weight]) + return weight + def EOSConstraintsLoglikelihood(eos_path, Neos, Constraint): parameters = {} logLs = [] for eos in range(1, Neos + 1): - radius, mass, Lambda = np.loadtxt('{0}/{1}.dat'.format(eos_path, eos), - unpack=True, usecols=[0, 1, 2]) + radius, mass, Lambda = np.loadtxt( + "{0}/{1}.dat".format(eos_path, eos), unpack=True, usecols=[0, 1, 2] + ) interp_mass_lambda = scipy.interpolate.interp1d(mass, Lambda) interp_mass_radius = scipy.interpolate.interp1d(mass, radius) - parameters['TOV_mass'] = mass[-1] - parameters['TOV_radius'] = radius[-1] - parameters['interp_mass_radius'] = interp_mass_radius - parameters['interp_mass_lambda'] = interp_mass_lambda + parameters["TOV_mass"] = mass[-1] + parameters["TOV_radius"] = radius[-1] + parameters["interp_mass_radius"] = interp_mass_radius + parameters["interp_mass_lambda"] = interp_mass_lambda logLs.append(Constraint.log_likelihood(parameters)) @@ -275,14 +349,17 @@ def EOSConstraintsLoglikelihood(eos_path, Neos, Constraint): return logLs - def EOSSorting(eos_path, out_path, Neos, weights): sortIdx = weights.argsort() for sortedIdx in range(0, Neos): originalIdx = sortIdx[sortedIdx] + 1 - os.system("cp {0}/{1}.dat {2}/{3}.dat".format(eos_path, originalIdx, out_path, sortedIdx + 1)) + os.system( + "cp {0}/{1}.dat {2}/{3}.dat".format( + eos_path, originalIdx, out_path, sortedIdx + 1 + ) + ) return @@ -321,6 +398,6 @@ def EOSConstraints2Prior(eos_path, out_path, Neos, Constraint): xx = np.arange(0, Neos + 1) yy = np.concatenate((sorted_weights, sorted_weights[-1])) - prior = Interped(xx, yy, minimum=0, maximum=Neos, name='EOS') + prior = Interped(xx, yy, minimum=0, maximum=Neos, name="EOS") return prior, logNorm diff --git a/nmma/gw/likelihood.py b/nmma/gw/likelihood.py index d907fe0d5..18e81605f 100644 --- a/nmma/gw/likelihood.py +++ b/nmma/gw/likelihood.py @@ -4,16 +4,19 @@ from ..joint.conversion import ( MultimessengerConversion, - MultimessengerConversionWithLambdas - ) + MultimessengerConversionWithLambdas, +) -from bilby.gw.likelihood import GravitationalWaveTransient, ROQGravitationalWaveTransient +from bilby.gw.likelihood import ( + GravitationalWaveTransient, + ROQGravitationalWaveTransient, +) from bilby.core.likelihood import Likelihood from bilby.core.prior import Interped class GravitationalWaveTransientLikelihoodwithEOS(Likelihood): - """ A GravitationalWaveTransient likelihood object + """A GravitationalWaveTransient likelihood object This likelihood uses the usual gravitational-wave transient but include an EOS handling for parameter conversion. @@ -93,33 +96,55 @@ class GravitationalWaveTransientLikelihoodwithEOS(Likelihood): """ - def __init__(self, interferometers, waveform_generator, - eos_path, Neos, eos_weight_path, binary_type, gw_likelihood_type, - priors, with_eos=True, - roq_weights=None, roq_params=None, roq_scale_factor=None, - time_marginalization=False, distance_marginalization=False, - phase_marginalization=False, distance_marginalization_lookup_table=None, - jitter_time=True, reference_frame="sky", time_reference="geocenter"): + def __init__( + self, + interferometers, + waveform_generator, + eos_path, + Neos, + eos_weight_path, + binary_type, + gw_likelihood_type, + priors, + with_eos=True, + roq_weights=None, + roq_params=None, + roq_scale_factor=None, + time_marginalization=False, + distance_marginalization=False, + phase_marginalization=False, + distance_marginalization_lookup_table=None, + jitter_time=True, + reference_frame="sky", + time_reference="geocenter", + ): # construct the eos prior if with_eos: xx = np.arange(0, Neos + 1) eos_weight = np.loadtxt(eos_weight_path) yy = np.concatenate((eos_weight, [eos_weight[-1]])) - eos_prior = Interped(xx, yy, minimum=0, maximum=Neos, name='EOS') - priors['EOS'] = eos_prior + eos_prior = Interped(xx, yy, minimum=0, maximum=Neos, name="EOS") + priors["EOS"] = eos_prior # construct the eos conversion - parameter_conversion_class = MultimessengerConversion(eos_data_path=eos_path, - Neos=Neos, - binary_type=binary_type, - with_ejecta=False) + parameter_conversion_class = MultimessengerConversion( + eos_data_path=eos_path, + Neos=Neos, + binary_type=binary_type, + with_ejecta=False, + ) else: - parameter_conversion_class = MultimessengerConversionWithLambdas(binary_type=binary_type, - with_ejecta=False) + parameter_conversion_class = MultimessengerConversionWithLambdas( + binary_type=binary_type, with_ejecta=False + ) - priors.conversion_function = parameter_conversion_class.priors_conversion_function - parameter_conversion = parameter_conversion_class.convert_to_multimessenger_parameters + priors.conversion_function = ( + parameter_conversion_class.priors_conversion_function + ) + parameter_conversion = ( + parameter_conversion_class.convert_to_multimessenger_parameters + ) waveform_generator.parameter_conversion = parameter_conversion # initialize the GW likelihood @@ -134,14 +159,19 @@ def __init__(self, interferometers, waveform_generator, reference_frame=reference_frame, time_reference=time_reference, ) - if gw_likelihood_type == 'GravitationalWaveTransient': + if gw_likelihood_type == "GravitationalWaveTransient": GWLikelihood = GravitationalWaveTransient(**gw_likelihood_kwargs) - elif gw_likelihood_type == 'ROQGravitationalWaveTransient': + elif gw_likelihood_type == "ROQGravitationalWaveTransient": gw_likelihood_kwargs.pop("time_marginalization", None) gw_likelihood_kwargs.pop("jitter_time", None) - gw_likelihood_kwargs.update(dict(weights=roq_weights, roq_params=roq_params, - roq_scale_factor=roq_scale_factor)) + gw_likelihood_kwargs.update( + dict( + weights=roq_weights, + roq_params=roq_params, + roq_scale_factor=roq_scale_factor, + ) + ) GWLikelihood = ROQGravitationalWaveTransient(**gw_likelihood_kwargs) super(GravitationalWaveTransientLikelihoodwithEOS, self).__init__(parameters={}) @@ -154,7 +184,7 @@ def __init__(self, interferometers, waveform_generator, self.__sync_parameters() def __repr__(self): - return self.__class__.__name__ + ' with ' + self.GWLikelihood.__repr__() + return self.__class__.__name__ + " with " + self.GWLikelihood.__repr__() def __sync_parameters(self): self.GWLikelihood.parameters = self.parameters diff --git a/nmma/joint/constraints.py b/nmma/joint/constraints.py index ed91d197e..981d90ad6 100644 --- a/nmma/joint/constraints.py +++ b/nmma/joint/constraints.py @@ -7,19 +7,22 @@ class PulsarConstraint(object): - def __init__(self, pulsar_masses, pulsar_masses_error): self.pulsar_masses = pulsar_masses self.pulsar_masses_error = pulsar_masses_error if len(self.pulsar_masses) != len(self.pulsar_masses_error): - raise ValueError('Number of masses and mass errors are inconsistent') + raise ValueError("Number of masses and mass errors are inconsistent") def __repr__(self): - return self.__class__.__name__ + ' with {} pulsars of masses {} and errors of {}'\ - .format(len(self.pulsar_masses), self.pulsar_masses, self.pulsar_masses_error) + return ( + self.__class__.__name__ + + " with {} pulsars of masses {} and errors of {}".format( + len(self.pulsar_masses), self.pulsar_masses, self.pulsar_masses_error + ) + ) def log_likelihood(self, parameters): - MTOV = parameters['TOV_mass'] + MTOV = parameters["TOV_mass"] logl = 0 for mass, error in zip(self.pulsar_masses, self.pulsar_masses_error): logl += scipy.stats.norm.logcdf(MTOV, loc=mass, scale=error) @@ -27,27 +30,28 @@ def log_likelihood(self, parameters): class MTOVUpperConstraint(object): - def __init__(self, maxMTOV, maxMTOV_error): self.maxMTOV = maxMTOV self.maxMTOV_error = maxMTOV_error def __repr__(self): - return self.__class__.__name__ + ' maxTOV of {} and error of {}'\ - .format(self.maxMTOV, self.maxMTOV_error) + return self.__class__.__name__ + " maxTOV of {} and error of {}".format( + self.maxMTOV, self.maxMTOV_error + ) def log_likelihood(self, parameters): - MTOV = parameters['TOV_mass'] + MTOV = parameters["TOV_mass"] logl = scipy.stats.norm.logsf(MTOV, loc=self.maxMTOV, scale=self.maxMTOV_error) return logl class NICERConstraint(object): - def __init__(self, NICER_path=None): if NICER_path is None: - self.NICER_path = os.path.join(os.path.dirname(__file__), 'NICER/J0030_3spot_RM.txt') + self.NICER_path = os.path.join( + os.path.dirname(__file__), "NICER/J0030_3spot_RM.txt" + ) else: self.NICER_path = NICER_path @@ -61,14 +65,18 @@ def __init__(self, NICER_path=None): self.KDE = scipy.stats.gaussian_kde((self.radius, self.mass)) def __repr__(self): - return self.__class__.__name__ + ' with the observation of PSR J0030+0451 ' +\ - 'with three potentially overlapping ovals. Data taken from {}'\ - .format(self.NICER_path) + return ( + self.__class__.__name__ + + " with the observation of PSR J0030+0451 " + + "with three potentially overlapping ovals. Data taken from {}".format( + self.NICER_path + ) + ) def log_likelihood(self, parameters): - interp = parameters['interp_mass_radius'] - masses = np.linspace(1., 2., num=1000) + interp = parameters["interp_mass_radius"] + masses = np.linspace(1.0, 2.0, num=1000) masses = masses[masses < interp.x[-1]] logl = scipy.special.logsumexp(self.KDE.logpdf((interp(masses), masses))) @@ -78,12 +86,11 @@ def log_likelihood(self, parameters): class JointConstraint(object): - def __init__(self, *constraints): self.constraints = constraints def log_likelihood(self, parameters): - logl = 0. + logl = 0.0 for constraint in self.constraints: logl += constraint.log_likelihood(parameters) diff --git a/nmma/joint/conversion.py b/nmma/joint/conversion.py index 559aad8f1..6f0e2f19f 100644 --- a/nmma/joint/conversion.py +++ b/nmma/joint/conversion.py @@ -22,7 +22,10 @@ def Hubble_constant_to_distance(converted_parameters, added_keys): # FIXME for future detection with high redshift # a proper cosmological model is needed - if "redshift" in converted_parameters.keys() and "Hubble_constant" in converted_parameters.keys(): + if ( + "redshift" in converted_parameters.keys() + and "Hubble_constant" in converted_parameters.keys() + ): Hubble_constant = converted_parameters["Hubble_constant"] redshift = converted_parameters["redshift"] # redshift is supposed to be dimensionless @@ -38,17 +41,26 @@ def source_frame_masses(converted_parameters, added_keys): if "redshift" not in converted_parameters.keys(): distance = converted_parameters["luminosity_distance"] - if hasattr(distance, '__len__') and len(distance)>50: #luminosity_distance_to_redshift gets really slow if too many distances are put in at once - from astropy import units - from astropy import cosmology as cosmo - cosmology = cosmo.Planck15 - zmin = cosmo.z_at_value(cosmology.luminosity_distance, distance.min() * units.Mpc) - zmax = cosmo.z_at_value(cosmology.luminosity_distance, distance.max() * units.Mpc) - zgrid = np.geomspace(zmin, zmax, 50) - distance_grid = cosmology.luminosity_distance(zgrid).value - converted_parameters["redshift"] = np.interp(distance, distance_grid, zgrid).value + if ( + hasattr(distance, "__len__") and len(distance) > 50 + ): # luminosity_distance_to_redshift gets really slow if too many distances are put in at once + from astropy import units + from astropy import cosmology as cosmo + + cosmology = cosmo.Planck15 + zmin = cosmo.z_at_value( + cosmology.luminosity_distance, distance.min() * units.Mpc + ) + zmax = cosmo.z_at_value( + cosmology.luminosity_distance, distance.max() * units.Mpc + ) + zgrid = np.geomspace(zmin, zmax, 50) + distance_grid = cosmology.luminosity_distance(zgrid).value + converted_parameters["redshift"] = np.interp( + distance, distance_grid, zgrid + ).value else: - converted_parameters["redshift"] = luminosity_distance_to_redshift(distance) + converted_parameters["redshift"] = luminosity_distance_to_redshift(distance) added_keys = added_keys + ["redshift"] if "mass_1_source" not in converted_parameters.keys(): @@ -64,22 +76,19 @@ def source_frame_masses(converted_parameters, added_keys): return converted_parameters, added_keys -def EOS2Parameters( - mass_val, radius_val, Lambda_val, mass_1_source, mass_2_source -): +def EOS2Parameters(mass_val, radius_val, Lambda_val, mass_1_source, mass_2_source): TOV_mass = mass_val.max() TOV_radius = radius_val[np.argmax(mass_val)] minimum_mass = mass_val.min() - if mass_1_source < minimum_mass or mass_1_source > TOV_mass: lambda_1 = np.array([0.0]) radius_1 = np.array([2.0 * mass_1_source * lal.MRSUN_SI / 1e3]) else: lambda_1 = np.array(np.interp(mass_1_source, mass_val, Lambda_val)).reshape(1) radius_1 = np.array(np.interp(mass_1_source, mass_val, radius_val)).reshape(1) - + if mass_2_source < minimum_mass or mass_2_source > TOV_mass: lambda_2 = np.array([0.0]) radius_2 = np.array([2.0 * mass_2_source * lal.MRSUN_SI / 1e3]) @@ -92,16 +101,17 @@ def EOS2Parameters( return TOV_mass, TOV_radius, lambda_1, lambda_2, radius_1, radius_2, R_14, R_16 + class NSBHEjectaFitting(object): def __init__(self): pass def chibh2risco(self, chi_bh): - Z1 = 1.0 + (1.0 - chi_bh ** 2) ** (1.0 / 3) * ( + Z1 = 1.0 + (1.0 - chi_bh**2) ** (1.0 / 3) * ( (1 + chi_bh) ** (1.0 / 3) + (1 - chi_bh) ** (1.0 / 3) ) - Z2 = np.sqrt(3.0 * chi_bh ** 2 + Z1 ** 2.0) + Z2 = np.sqrt(3.0 * chi_bh**2 + Z1**2.0) return 3.0 + Z2 - np.sign(chi_bh) * np.sqrt((3.0 - Z1) * (3.0 + Z1 + 2.0 * Z2)) @@ -123,9 +133,9 @@ def remnant_disk_mass_fitting( c=0.25512517, d=0.761250847, ): - ''' + """ equation (4) in https://arxiv.org/pdf/1807.00011 - ''' + """ mass_ratio_invert = mass_1_source / mass_2_source symm_mass_ratio = mass_ratio_invert / np.power(1.0 + mass_ratio_invert, 2.0) @@ -217,7 +227,9 @@ def ejecta_parameter_conversion(self, converted_parameters, added_keys): converted_parameters["ratio_zeta"] ) + np.log10(mdisk_fit) # total eject mass - total_ejeta_mass = 10**log10_mej_dyn + 10**converted_parameters["log10_mej_wind"] + total_ejeta_mass = ( + 10**log10_mej_dyn + 10 ** converted_parameters["log10_mej_wind"] + ) log10_mej = np.log10(total_ejeta_mass) converted_parameters["log10_mej"] = log10_mej @@ -236,7 +248,6 @@ def ejecta_parameter_conversion(self, converted_parameters, added_keys): converted_parameters["log10_mej_wind"] = -np.inf converted_parameters["log10_mej"] = -np.inf - added_keys = added_keys + ["log10_mej_dyn", "log10_mej_wind", "log10_mej"] return converted_parameters, added_keys @@ -335,8 +346,8 @@ def ejecta_parameter_conversion(self, converted_parameters, added_keys): # prevent the output message flooded by these warning messages old = np.seterr() - np.seterr(invalid='ignore') - np.seterr(divide='ignore') + np.seterr(invalid="ignore") + np.seterr(divide="ignore") mass_1_source = converted_parameters["mass_1_source"] mass_2_source = converted_parameters["mass_2_source"] @@ -452,7 +463,12 @@ def ejecta_parameter_conversion(self, converted_parameters, added_keys): converted_parameters["log10_mej"] = float(log10_mej) converted_parameters["log10_E0"] = float(log10_E0) - added_keys = added_keys + ["log10_mej_dyn", "log10_mej_wind", "log10_mej", "log10_E0"] + added_keys = added_keys + [ + "log10_mej_dyn", + "log10_mej_wind", + "log10_mej", + "log10_E0", + ] np.seterr(**old) @@ -509,13 +525,14 @@ def convert_to_multimessenger_parameters(self, parameters): R_14_list = [] R_16_list = [] EOSID = np.array(converted_parameters["EOS"]).astype(int) + 1 - EOSdata = [None]*self.Neos + EOSdata = [None] * self.Neos for j in np.unique(EOSID): - EOSdata[j-1] = np.loadtxt(f"{self.eos_data_path}/{j}.dat", usecols = [0,1,2]).T - + EOSdata[j - 1] = np.loadtxt( + f"{self.eos_data_path}/{j}.dat", usecols=[0, 1, 2] + ).T for i in range(0, len(EOSID)): - radius_val, mass_val, Lambda_val = EOSdata[EOSID[i]-1] + radius_val, mass_val, Lambda_val = EOSdata[EOSID[i] - 1] ( TOV_mass, @@ -525,7 +542,7 @@ def convert_to_multimessenger_parameters(self, parameters): radius_1, radius_2, R_14, - R_16 + R_16, ) = EOS2Parameters( mass_val, radius_val, @@ -533,7 +550,7 @@ def convert_to_multimessenger_parameters(self, parameters): mass_1_source[i], mass_2_source[i], ) - + TOV_radius_list.append(TOV_radius) TOV_mass_list.append(TOV_mass) lambda_1_list.append(lambda_1[0]) @@ -557,7 +574,9 @@ def convert_to_multimessenger_parameters(self, parameters): else: EOSID = int(converted_parameters["EOS"]) + 1 - radius_val, mass_val, Lambda_val = np.loadtxt(f"{self.eos_data_path}/{EOSID}.dat", unpack = True, usecols=[0,1,2]) + radius_val, mass_val, Lambda_val = np.loadtxt( + f"{self.eos_data_path}/{EOSID}.dat", unpack=True, usecols=[0, 1, 2] + ) ( TOV_mass, @@ -567,15 +586,11 @@ def convert_to_multimessenger_parameters(self, parameters): radius_1, radius_2, R_14, - R_16 + R_16, ) = EOS2Parameters( - mass_val, - radius_val, - Lambda_val, - mass_1_source, - mass_2_source + mass_val, radius_val, Lambda_val, mass_1_source, mass_2_source ) - + converted_parameters["TOV_radius"] = TOV_radius converted_parameters["TOV_mass"] = TOV_mass @@ -614,13 +629,12 @@ def convert_to_multimessenger_parameters(self, parameters): ) added_keys = added_keys + ["KNtheta", "inclination_EM"] - + added_keys = [ key for key in converted_parameters.keys() if key not in original_keys ] return converted_parameters, added_keys - def generate_all_parameters(self, sample, likelihood=None, priors=None, npool=1): waveform_defaults = { diff --git a/nmma/joint/maximum_mass_constraint.py b/nmma/joint/maximum_mass_constraint.py index 4be3a6474..a45ca320a 100644 --- a/nmma/joint/maximum_mass_constraint.py +++ b/nmma/joint/maximum_mass_constraint.py @@ -7,124 +7,120 @@ import os - def get_parser(): parser = argparse.ArgumentParser( description="Inference on the maximum mass constraint of the EOS when a joint posterior for binary components, ejecta and EOS is provided." ) + parser.add_argument("--outdir", metavar="PATH", type=str, required=True) parser.add_argument( - "--outdir", - metavar="PATH", - type=str, - required=True - ) - parser.add_argument( - "--joint-posterior", - metavar="PATH", - type=str, - required=True, - help="Posterior file with chirp mass, eta_star, EOSID, log10_mdisk, log10_mej_dyn as columns." + "--joint-posterior", + metavar="PATH", + type=str, + required=True, + help="Posterior file with chirp mass, eta_star, EOSID, log10_mdisk, log10_mej_dyn as columns.", ) parser.add_argument( - "--prior", - metavar="PATH", - type=str, - required=True, - help="Prior specification for chirp mass, eta_star, log10_mdisk, log10_mej_dyn. If use_M_kepler is True, prior file must also contain ratio_R and delta.", + "--prior", + metavar="PATH", + type=str, + required=True, + help="Prior specification for chirp mass, eta_star, log10_mdisk, log10_mej_dyn. If use_M_kepler is True, prior file must also contain ratio_R and delta.", ) parser.add_argument( - "--eos-path-macro", - metavar="PATH", - type=str, - required=True, - help="EOS folder with EOS files in [R, M, Lambda, p_central] format." + "--eos-path-macro", + metavar="PATH", + type=str, + required=True, + help="EOS folder with EOS files in [R, M, Lambda, p_central] format.", ) parser.add_argument( - "--eos-path-micro", - metavar="PATH", - type=str, - required=True, - help="EOS folder with EOS files in [n, energy density, pressure, speed of sound squared] format." + "--eos-path-micro", + metavar="PATH", + type=str, + required=True, + help="EOS folder with EOS files in [n, energy density, pressure, speed of sound squared] format.", ) parser.add_argument( - "--use-M-Kepler", - action="store_true", - help="If set, it is assumed that the BNS remnant collapsed to a black hole above the Kepler limit.", - + "--use-M-Kepler", + action="store_true", + help="If set, it is assumed that the BNS remnant collapsed to a black hole above the Kepler limit.", ) parser.add_argument( - "--nlive", - metavar="nlive", - type=int, - required=False, - default=1024 + "--nlive", metavar="nlive", type=int, required=False, default=1024 ) - + return parser + def maximum_mass_resampling(args): try: - os.makedirs(args.outdir + "/pm/") + os.makedirs(args.outdir + "/pm/") except: - pass - - posterior_samples = pd.read_csv(args.joint_posterior, header = 0, delimiter = " ") + pass + + posterior_samples = pd.read_csv(args.joint_posterior, header=0, delimiter=" ") prior = bilby.gw.prior.PriorDict(args.prior) Neos = len(next(os.walk(args.eos_path_macro))[2]) - - + if args.use_M_Kepler: - n_dims = 7 - if len(prior.keys())!= n_dims-1: - raise Exception("If use_M_kepler is True, you need to provide a prior fo ratio_R and delta to be used in the quasi-universal relation.") + n_dims = 7 + if len(prior.keys()) != n_dims - 1: + raise Exception( + "If use_M_kepler is True, you need to provide a prior fo ratio_R and delta to be used in the quasi-universal relation." + ) else: - n_dims = 5 - - - - + n_dims = 5 + pymulti_kwargs = dict( - outputfiles_basename=args.outdir + "/pm/", - n_dims=n_dims, - n_live_points=args.nlive, - verbose=True, - resume=True, - seed=42, - importance_nested_sampling=False, - use_MPI = True, - ) - - solution = utils.PostmergerInference(prior, posterior_samples, Neos, args.eos_path_macro, args.eos_path_micro, args.use_M_Kepler, **pymulti_kwargs) - - + outputfiles_basename=args.outdir + "/pm/", + n_dims=n_dims, + n_live_points=args.nlive, + verbose=True, + resume=True, + seed=42, + importance_nested_sampling=False, + use_MPI=True, + ) + + solution = utils.PostmergerInference( + prior, + posterior_samples, + Neos, + args.eos_path_macro, + args.eos_path_micro, + args.use_M_Kepler, + **pymulti_kwargs + ) + samples = solution.samples.T posterior = dict() - + posterior["chirp_mass"] = samples[0] posterior["eta_star"] = samples[1] posterior["EOS"] = samples[2] posterior["log10_mdisk"] = samples[3] posterior["log10_mej_dyn"] = samples[4] - + if args.use_M_Kepler: - posterior["ratio_R"] = samples[5] - posterior["delta"] = samples[6] - - + posterior["ratio_R"] = samples[5] + posterior["delta"] = samples[6] + posterior = pd.DataFrame.from_dict(posterior) - posterior.to_csv(args.outdir+"/posterior_samples.dat", sep = " ", index = False) - -def main(args = None): - - if args is None: - parser = get_parser() - args = parser.parse_args() - - maximum_mass_resampling(args) - - else: - maximum_mass_resampling(args) + posterior.to_csv(args.outdir + "/posterior_samples.dat", sep=" ", index=False) + + +def main(args=None): + + if args is None: + parser = get_parser() + args = parser.parse_args() + + maximum_mass_resampling(args) + + else: + maximum_mass_resampling(args) + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/nmma/joint/maximum_mass_constraint_utils.py b/nmma/joint/maximum_mass_constraint_utils.py index 3b9422ae3..48acd9046 100644 --- a/nmma/joint/maximum_mass_constraint_utils.py +++ b/nmma/joint/maximum_mass_constraint_utils.py @@ -14,104 +14,131 @@ import sys import contextlib + def fileno(file_or_fd): - fd = getattr(file_or_fd, 'fileno', lambda: file_or_fd)() + fd = getattr(file_or_fd, "fileno", lambda: file_or_fd)() if not isinstance(fd, int): raise ValueError("Expected a file (`.fileno()`) or a file descriptor") return fd + @contextlib.contextmanager def stdout_redirected(to=os.devnull, stdout=None): """ https://stackoverflow.com/a/22434262/190597 (J.F. Sebastian) """ if stdout is None: - stdout = sys.stdout + stdout = sys.stdout stdout_fd = fileno(stdout) # copy stdout_fd before it is overwritten - #NOTE: `copied` is inheritable on Windows when duplicating a standard stream - with os.fdopen(os.dup(stdout_fd), 'wb') as copied: + # NOTE: `copied` is inheritable on Windows when duplicating a standard stream + with os.fdopen(os.dup(stdout_fd), "wb") as copied: stdout.flush() # flush library buffers that dup2 knows nothing about try: os.dup2(fileno(to), stdout_fd) # $ exec >&to except ValueError: # filename - with open(to, 'wb') as to_file: + with open(to, "wb") as to_file: os.dup2(to_file.fileno(), stdout_fd) # $ exec > to try: - yield stdout # allow code to be run with the redirected stdout + yield stdout # allow code to be run with the redirected stdout finally: # restore stdout to its previous value - #NOTE: dup2 makes stdout_fd inheritable unconditionally + # NOTE: dup2 makes stdout_fd inheritable unconditionally stdout.flush() os.dup2(copied.fileno(), stdout_fd) # $ exec >&copied - def baryonic_mass(gravitational_mass, EOS, eos_path_macro, eos_path_micro): - Msun_to_MeV = (1*u.Msun*constants.c**2).to(u.MeV).value #1.1154E+60 MeV are one solar mass - MeV_per_fm3_to_Msun_per_km3 = 1/Msun_to_MeV * (1.*u.fm**(-3)).to(u.km**(-3)).value # 1 MeV/fm**3 is 8.9653E-7 Msun/km**3 - G = (constants.G/constants.c**2).to(u.km/u.Msun).value #G/c^2 is 1.4766 km/Msun - - - R, M, L, P0 = np.loadtxt(eos_path_macro+f"/{EOS}.dat", unpack= True, skiprows = 0) - N, EPS, P, CS2 = np.loadtxt(eos_path_micro+f"/{EOS}.dat", unpack = True, skiprows = 0) - eps_of_p = scipy.interpolate.interp1d(P, EPS, kind = 'linear', bounds_error = False, fill_value =0) - + Msun_to_MeV = ( + (1 * u.Msun * constants.c**2).to(u.MeV).value + ) # 1.1154E+60 MeV are one solar mass + MeV_per_fm3_to_Msun_per_km3 = ( + 1 / Msun_to_MeV * (1.0 * u.fm ** (-3)).to(u.km ** (-3)).value + ) # 1 MeV/fm**3 is 8.9653E-7 Msun/km**3 + G = ( + (constants.G / constants.c**2).to(u.km / u.Msun).value + ) # G/c^2 is 1.4766 km/Msun + + R, M, L, P0 = np.loadtxt(eos_path_macro + f"/{EOS}.dat", unpack=True, skiprows=0) + N, EPS, P, CS2 = np.loadtxt(eos_path_micro + f"/{EOS}.dat", unpack=True, skiprows=0) + eps_of_p = scipy.interpolate.interp1d( + P, EPS, kind="linear", bounds_error=False, fill_value=0 + ) def TOVeq(y, x): - p, m = y - eps = eps_of_p(p) - - Dp = -G*m*eps/x**2 * (1+p/eps) * (1+4*np.pi*(x**3* p)*MeV_per_fm3_to_Msun_per_km3/m) * (1-2*G*m/x)**(-1) - Dm = (4*np.pi*x**2*eps)*MeV_per_fm3_to_Msun_per_km3 - - return [Dp, Dm] - - dr =0.001 + p, m = y + eps = eps_of_p(p) + + Dp = ( + -G + * m + * eps + / x**2 + * (1 + p / eps) + * (1 + 4 * np.pi * (x**3 * p) * MeV_per_fm3_to_Msun_per_km3 / m) + * (1 - 2 * G * m / x) ** (-1) + ) + Dm = (4 * np.pi * x**2 * eps) * MeV_per_fm3_to_Msun_per_km3 + + return [Dp, Dm] + + dr = 0.001 r = np.interp(gravitational_mass, M, R) - p0= np.interp(gravitational_mass, M, P0) + p0 = np.interp(gravitational_mass, M, P0) eps0 = eps_of_p(p0) - m0 = (eps0*4*np.pi/3*dr**3)*MeV_per_fm3_to_Msun_per_km3 - x = np.arange(dr, r+dr, dr) + m0 = (eps0 * 4 * np.pi / 3 * dr**3) * MeV_per_fm3_to_Msun_per_km3 + x = np.arange(dr, r + dr, dr) y0 = [p0, m0] with stdout_redirected(): - p_solv, m_solv = scipy.integrate.odeint(TOVeq, y0=y0, t = x).T + p_solv, m_solv = scipy.integrate.odeint(TOVeq, y0=y0, t=x).T n_solv = np.interp(p_solv, P, N) - if np.any(np.isnan(p_solv)): - cut = np.where(np.isnan(p_solv))[0][0] - if np.any(np.isnan(m_solv)): - cut = min(cut, np.where(np.isnan(m_solv))[0][0]) - n_solv = n_solv[:cut]; m_solv = m_solv[:cut]; x= x[:cut] - - n_solv = n_solv *(1.*u.fm**(-3)).to(u.km**(-3)) #convert from fm**(-3) to km**(-3) - particle_mass = constants.m_p.to(u.Msun).value #get proton mass in Msun - m_baryonic = particle_mass*4*np.pi*scipy.integrate.simpson(y=(n_solv)*x**2/np.sqrt(1-2*G*m_solv/x), x=x) #get baryonic mass in Msun - + cut = np.where(np.isnan(p_solv))[0][0] + if np.any(np.isnan(m_solv)): + cut = min(cut, np.where(np.isnan(m_solv))[0][0]) + n_solv = n_solv[:cut] + m_solv = m_solv[:cut] + x = x[:cut] + + n_solv = n_solv * (1.0 * u.fm ** (-3)).to( + u.km ** (-3) + ) # convert from fm**(-3) to km**(-3) + particle_mass = constants.m_p.to(u.Msun).value # get proton mass in Msun + m_baryonic = ( + particle_mass + * 4 + * np.pi + * scipy.integrate.simpson( + y=(n_solv) * x**2 / np.sqrt(1 - 2 * G * m_solv / x), x=x + ) + ) # get baryonic mass in Msun + if np.isnan(m_baryonic): - import warnings - warnings.warn("baryonic_mass returned nan") + import warnings + + warnings.warn("baryonic_mass returned nan") return m_baryonic def baryonic_Kepler_mass(mTOV, R_14, ratio_R, delta): - """ - see https://arxiv.org/abs/2307.03225 and https://arxiv.org/abs/1905.03784 - """ - m_max = ratio_R*mTOV - m_max_b = m_max + 0.78/R_14 * m_max**2 - m_max_b *= (1+delta) + """ + see https://arxiv.org/abs/2307.03225 and https://arxiv.org/abs/1905.03784 + """ + m_max = ratio_R * mTOV + m_max_b = m_max + 0.78 / R_14 * m_max**2 + m_max_b *= 1 + delta - return m_max_b + return m_max_b from pymultinest.solve import Solver + class PostmergerInference(Solver): """ Class to sample over a joint GW+EM posterior and to determine the remnant mass of the system. @@ -119,22 +146,30 @@ class PostmergerInference(Solver): See https://arxiv.org/abs/2402.04172 for details. """ - def __init__(self, prior, posterior_samples, Neos, eos_path_macro, eos_path_micro, use_M_max = False, **kwargs): + def __init__( + self, + prior, + posterior_samples, + Neos, + eos_path_macro, + eos_path_micro, + use_M_max=False, + **kwargs, + ): self.posterior_samples = posterior_samples self.use_M_max = use_M_max - - self.priors = {'chirp_mass': prior['chirp_mass'], - 'eta_star': prior['eta_star'], - 'EOS': Uniform(name = "EOS", minimum=0, maximum=Neos), - 'log10_mdisk': prior['log10_mdisk'], - 'log10_mej_dyn': prior['log10_mej_dyn'] - } + self.priors = { + "chirp_mass": prior["chirp_mass"], + "eta_star": prior["eta_star"], + "EOS": Uniform(name="EOS", minimum=0, maximum=Neos), + "log10_mdisk": prior["log10_mdisk"], + "log10_mej_dyn": prior["log10_mej_dyn"], + } if self.use_M_max: self.priors["ratio_R"] = prior["ratio_R"] self.priors["delta"] = prior["delta"] - self.priors = PriorDict(self.priors) self._search_parameter_keys = self.priors.keys() self.Neos = Neos @@ -147,7 +182,9 @@ def __init__(self, prior, posterior_samples, Neos, eos_path_macro, eos_path_micr log10_mej_dyn = self.posterior_samples.log10_mej_dyn.to_numpy() log10_mdisk = self.posterior_samples.log10_mdisk.to_numpy() - self.KDE = scipy.stats.gaussian_kde((chirp_mass, eta_star, EOS, log10_mdisk, log10_mej_dyn)) + self.KDE = scipy.stats.gaussian_kde( + (chirp_mass, eta_star, EOS, log10_mdisk, log10_mej_dyn) + ) Solver.__init__(self, **kwargs) @@ -157,36 +194,49 @@ def Prior(self, x): def LogLikelihood(self, x): if self.use_M_max: - chirp_mass, eta_star, EOS, log10_mdisk, log10_mej_dyn, ratio_R, delta = x + chirp_mass, eta_star, EOS, log10_mdisk, log10_mej_dyn, ratio_R, delta = x else: - chirp_mass, eta_star, EOS, log10_mdisk, log10_mej_dyn = x + chirp_mass, eta_star, EOS, log10_mdisk, log10_mej_dyn = x - - logprior = self.KDE.logpdf((chirp_mass, eta_star, EOS, log10_mdisk, log10_mej_dyn)) #use the joint GW+EM posterior as "prior" here + logprior = self.KDE.logpdf( + (chirp_mass, eta_star, EOS, log10_mdisk, log10_mej_dyn) + ) # use the joint GW+EM posterior as "prior" here EOS = int(EOS) + 1 - R, M = np.loadtxt(self.eos_path_macro+f"/{EOS}.dat", unpack = True, usecols = [0,1], skiprows = 0) + R, M = np.loadtxt( + self.eos_path_macro + f"/{EOS}.dat", unpack=True, usecols=[0, 1], skiprows=0 + ) mTOV = M.max() R_14 = np.interp(1.4, M, R) - q = conversion.symmetric_mass_ratio_to_mass_ratio(0.25-np.exp(eta_star)) - mass_1, mass_2 = conversion.chirp_mass_and_mass_ratio_to_component_masses(chirp_mass, q) + q = conversion.symmetric_mass_ratio_to_mass_ratio(0.25 - np.exp(eta_star)) + mass_1, mass_2 = conversion.chirp_mass_and_mass_ratio_to_component_masses( + chirp_mass, q + ) mdisk = 10**log10_mdisk mej_dyn = 10**log10_mej_dyn - m_rem_b = baryonic_mass(mass_1, EOS, self.eos_path_macro, self.eos_path_micro) + baryonic_mass(mass_2, EOS, self.eos_path_macro, self.eos_path_micro) - mdisk - mej_dyn #calculate the baryonic remnant mass + m_rem_b = ( + baryonic_mass(mass_1, EOS, self.eos_path_macro, self.eos_path_micro) + + baryonic_mass(mass_2, EOS, self.eos_path_macro, self.eos_path_micro) + - mdisk + - mej_dyn + ) # calculate the baryonic remnant mass if self.use_M_max: - m_threshold = baryonic_Kepler_mass(mTOV, R_14, ratio_R, delta) #if the Kepler limit is the threshold, use the quasiuniversal relation + m_threshold = baryonic_Kepler_mass( + mTOV, R_14, ratio_R, delta + ) # if the Kepler limit is the threshold, use the quasiuniversal relation else: - m_threshold = baryonic_mass(mTOV, EOS, self.eos_path_macro, self.eos_path_micro) #if the TOV mass is the limit just determine its baryonic mass - + m_threshold = baryonic_mass( + mTOV, EOS, self.eos_path_macro, self.eos_path_micro + ) # if the TOV mass is the limit just determine its baryonic mass if m_threshold > m_rem_b: - loglikelihood = np.nan_to_num(-np.inf) + loglikelihood = np.nan_to_num(-np.inf) else: - loglikelihood = 0 + loglikelihood = 0 - return logprior + loglikelihood \ No newline at end of file + return logprior + loglikelihood diff --git a/nmma/mlmodel/dataprocessing.py b/nmma/mlmodel/dataprocessing.py index 947165187..adb2ec904 100644 --- a/nmma/mlmodel/dataprocessing.py +++ b/nmma/mlmodel/dataprocessing.py @@ -28,11 +28,11 @@ time_step = 0.25 def open_json(file_name, dir_path): - ''' - Opens a json file, loads the data as a dictionary, and closes the file + ''' + Opens a json file, loads the data as a dictionary, and closes the file Inputs: file_name = /name of json file.json - dir_path = directory containing json files + dir_path = directory containing json files Returns: data = dictionary containing json content ''' @@ -42,14 +42,14 @@ def open_json(file_name, dir_path): return data def get_names(path, label, set, num): - ''' + ''' Gets the file path for the fixed data Inputs: path = string, directory to point to label = string, label assigned during nmma light curve generation set = int, number in directory name num = int, number of files to unpack - Returns: + Returns: file_names = list, contains full path file names ''' file_names = [0] * num @@ -59,7 +59,7 @@ def get_names(path, label, set, num): return file_names def json_to_df(file_names, num_sims, detection_limit=detection_limit, bands=bands): - ''' + ''' Flattens json files into a dataframe Inputs: file_names = list, contains full path file names as strings @@ -145,7 +145,7 @@ def pad_the_data(actual_df, column_list, desired_count=num_points, filler_time_s Outputs: cat_df: padded dataframe ''' - actual_df.iloc[:, actual_df.columns.get_loc('t')] = actual_df.iloc[:, actual_df.columns.get_loc('t')].apply(lambda x: x - t_min) + actual_df.iloc[:, actual_df.columns.get_loc('t')] = actual_df.iloc[:, actual_df.columns.get_loc('t')].apply(lambda x: x - t_min) cat_df = actual_df cat_count = len(cat_df) prepended_count = 0 @@ -156,8 +156,8 @@ def pad_the_data(actual_df, column_list, desired_count=num_points, filler_time_s cat_df = pd.concat([prepend_filler_df, actual_df], ignore_index=True) cat_count = len(cat_df) append_count = desired_count - cat_count - if append_count > 0: - max_t = cat_df['t'].max() + if append_count > 0: + max_t = cat_df['t'].max() steps_per_count = 1/filler_time_step filler_min_time = int(max_t*steps_per_count)/steps_per_count + filler_time_step # start at next time step append_filler_df = gen_append_filler(column_list, filler_data, filler_min_time, append_count) @@ -169,7 +169,7 @@ def pad_the_data(actual_df, column_list, desired_count=num_points, filler_time_s def pad_all_dfs(df_list): ''' Pads multiple dataframes at a time - Inputs: + Inputs: df_list: list of dataframes to pad Outputs: padded_df_list: list of dataframes after padding @@ -194,7 +194,7 @@ def load_in_data(data_dir, name, csv_no, num_points=num_points, num_repeats=num_ num_points: number of data points per light curve num_repeats: repeats of injection parameters to determine batches Outputs: - data_df: single dataframe containing the data + data_df: single dataframe containing the data ''' data_list = [] for i in range (0, csv_no): @@ -218,7 +218,7 @@ def load_in_data(data_dir, name, csv_no, num_points=num_points, num_repeats=num_ batch_no += 1 data_df['batch_id'] = batch_list return data_df - + def match_fix_to_var(data_dir, name1, name2, start, stop, num_points=num_points, num_repeats=num_repeats): ''' Matches the shifted injection light curve data to its fixed counterpart @@ -279,7 +279,7 @@ def match_fix_to_var(data_dir, name1, name2, start, stop, num_points=num_points, batch_no += 1 fixed_data_df['batch_id'] = batch_list varied_data_df['batch_id'] = batch_list - + return fixed_data_df, varied_data_df def matched(data_dir, name1, name2, start, stop, num_points=num_points, num_repeats=num_repeats): @@ -339,14 +339,14 @@ def add_batch_sim_nums_all(df, num_points=num_points, num_repeats=num_repeats): df['sim_id'] = sim_list_split def get_test_names(path, label, set, num): - ''' + ''' Gets the file path for the fixed data Inputs: path = string, directory to point to label = string, label assigned during nmma light curve generation set = int, number in directory name num = int, number of files to unpack - Returns: + Returns: list, contains full path file names ''' file_names = [0] * num @@ -373,13 +373,13 @@ def repeated_df_to_tensor(df_varied, df_fixed, batches): param_shifted_list = [] param_unshifted_list = [] for idx in tqdm(range(0, batches)): - data_shifted = torch.tensor(df_varied.loc[df_varied['batch_id'] == idx].iloc[:, 1:4].values.reshape(num_repeats, num_points, num_channels), + data_shifted = torch.tensor(df_varied.loc[df_varied['batch_id'] == idx].iloc[:, 1:4].values.reshape(num_repeats, num_points, num_channels), dtype=torch.float32).transpose(1, 2) - data_unshifted = torch.tensor(df_fixed.loc[df_fixed['batch_id'] == idx].iloc[:, 1:4].values.reshape(num_repeats, num_points, num_channels), + data_unshifted = torch.tensor(df_fixed.loc[df_fixed['batch_id'] == idx].iloc[:, 1:4].values.reshape(num_repeats, num_points, num_channels), dtype=torch.float32).transpose(1, 2) - param_shifted = torch.tensor(df_varied.loc[df_varied['batch_id'] == idx].iloc[::num_points, 6:11].values, + param_shifted = torch.tensor(df_varied.loc[df_varied['batch_id'] == idx].iloc[::num_points, 6:11].values, dtype=torch.float32).unsqueeze(2).transpose(1,2) - param_unshifted = torch.tensor(df_fixed.loc[df_fixed['batch_id'] == idx].iloc[::num_points, 5:10].values, + param_unshifted = torch.tensor(df_fixed.loc[df_fixed['batch_id'] == idx].iloc[::num_points, 5:10].values, dtype=torch.float32).unsqueeze(2).transpose(1,2) data_shifted_list.append(data_shifted) data_unshifted_list.append(data_unshifted) @@ -404,7 +404,7 @@ def __len__(self): def __getitem__(self, idx): if torch.is_tensor(idx): idx = idx.tolist() - + return ( self.param_shifted_paper[idx].to(device), self.param_unshifted_paper[idx].to(device), diff --git a/nmma/mlmodel/embedding.py b/nmma/mlmodel/embedding.py index 98fbafe7c..3f7a7a44f 100644 --- a/nmma/mlmodel/embedding.py +++ b/nmma/mlmodel/embedding.py @@ -8,19 +8,21 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -bands = ['ztfg', 'ztfr', 'ztfi'] +bands = ["ztfg", "ztfr", "ztfi"] detection_limit = 22.0 num_repeats = 50 num_channels = 3 num_points = 121 + class VICRegLoss(nn.Module): - ''' + """ Variance-Invariance-Covariance Regularization Loss Function - citation: A. Bardes, J. Ponce, and Y. LeCun. Vicreg: - Variance-invariance-covariance regularization - for self-supervised learning, 2022. - ''' + citation: A. Bardes, J. Ponce, and Y. LeCun. Vicreg: + Variance-invariance-covariance regularization + for self-supervised learning, 2022. + """ + def forward(self, x, y, wt_repr=1.0, wt_cov=1.0, wt_std=1.0): repr_loss = F.mse_loss(x, y) x = x - x.mean(dim=0) @@ -30,8 +32,8 @@ def forward(self, x, y, wt_repr=1.0, wt_cov=1.0, wt_std=1.0): std_x = torch.sqrt(x.var(dim=0) + 0.0001) std_y = torch.sqrt(y.var(dim=0) + 0.0001) std_loss = torch.mean(F.relu(1 - std_x)) / 2 + torch.mean(F.relu(1 - std_y)) / 2 - x = (x-x.mean(dim=0))/std_x - y = (y-y.mean(dim=0))/std_y + x = (x - x.mean(dim=0)) / std_x + y = (y - y.mean(dim=0)) / std_y # transpose dims 1 and 2; keep batch dim i.e. 0, unchanged x = x.unsqueeze(1) y = y.unsqueeze(1) @@ -39,27 +41,34 @@ def forward(self, x, y, wt_repr=1.0, wt_cov=1.0, wt_std=1.0): cov_y = (y.transpose(1, 2) @ y) / (N - 1) cov_loss = self.off_diagonal(cov_x).pow_(2).sum().div(D) cov_loss += self.off_diagonal(cov_y).pow_(2).sum().div(D) - s = wt_repr*repr_loss + wt_cov*cov_loss + wt_std*std_loss - + s = wt_repr * repr_loss + wt_cov * cov_loss + wt_std * std_loss + return s, repr_loss, cov_loss, std_loss - def off_diagonal(self,cov): + def off_diagonal(self, cov): num_batch, n, m = cov.shape assert n == m # All off diagonal elements from complete batch flattened # import pdb; pdb.set_trace() - - return cov.flatten(start_dim=1)[...,:-1].view(num_batch, n - 1, n + 1)[...,1:].flatten() - + + return ( + cov.flatten(start_dim=1)[..., :-1] + .view(num_batch, n - 1, n + 1)[..., 1:] + .flatten() + ) + + class ConvResidualBlock(nn.Module): - def __init__(self, - channels, - kernel_size, - activation=F.relu, - dropout_probability=0.1, - use_batch_norm=True, - zero_initialization=True): - ''' + def __init__( + self, + channels, + kernel_size, + activation=F.relu, + dropout_probability=0.1, + use_batch_norm=True, + zero_initialization=True, + ): + """ Inputs: channels: number of separate dimensions, here it is photometric bands kernel_size: (odd int) input field size of the CNN @@ -67,13 +76,20 @@ def __init__(self, dropout_probability: (float) probability an element will be zeroed use_batch_norm: (bool) determines if data will be normalized by re-scaling and re-centering zero_initialization: (bool) initializes weights with zeros - ''' + """ super().__init__() self.activation = activation self.use_batch_norm = use_batch_norm if use_batch_norm: - self.batch_norm_layers = nn.ModuleList([nn.BatchNorm1d(channels, eps=1e-3) for _ in range(2)]) - self.conv_layers = nn.ModuleList([nn.Conv1d(channels, channels, kernel_size=kernel_size, padding='same') for _ in range(2)]) + self.batch_norm_layers = nn.ModuleList( + [nn.BatchNorm1d(channels, eps=1e-3) for _ in range(2)] + ) + self.conv_layers = nn.ModuleList( + [ + nn.Conv1d(channels, channels, kernel_size=kernel_size, padding="same") + for _ in range(2) + ] + ) self.dropout = nn.Dropout(p=dropout_probability) if zero_initialization: nn.init.uniform_(self.conv_layers[-1].weight, -1e-3, 1e-3) @@ -90,22 +106,25 @@ def forward(self, inputs): temps = self.activation(temps) temps = self.dropout(temps) temps = self.conv_layers[1](temps) - + return inputs + temps + class ConvResidualNet(nn.Module): - ''' Convolutional Neural Network composed of many Convolutional Residual Blocks''' - def __init__(self, - in_channels, - out_channels, - hidden_channels, - num_blocks, - kernel_size, - activation=F.relu, - dropout_probability=0.1, - use_batch_norm=True, - ): - ''' + """Convolutional Neural Network composed of many Convolutional Residual Blocks""" + + def __init__( + self, + in_channels, + out_channels, + hidden_channels, + num_blocks, + kernel_size, + activation=F.relu, + dropout_probability=0.1, + use_batch_norm=True, + ): + """ Inputs: in_channels: starting number of channels, ie number of photometric bands out_channels: end number of channels @@ -115,56 +134,74 @@ def __init__(self, activation: (function) calculates node output dropout_probability: (float) probability of resetting weights use_batch_norm: (bool) determines if batch is normalized - ''' + """ super().__init__() self.hidden_channels = hidden_channels - self.initial_layer = nn.Conv1d(in_channels=in_channels, out_channels=hidden_channels, kernel_size=kernel_size, padding='same') + self.initial_layer = nn.Conv1d( + in_channels=in_channels, + out_channels=hidden_channels, + kernel_size=kernel_size, + padding="same", + ) self.blocks = nn.ModuleList( - [ConvResidualBlock( - channels=hidden_channels, - activation=activation, - dropout_probability=dropout_probability, - use_batch_norm=use_batch_norm, - kernel_size=kernel_size) - for _ in range(num_blocks)]) - self.final_layer = nn.Conv1d(hidden_channels, out_channels, kernel_size=kernel_size, padding='same') + [ + ConvResidualBlock( + channels=hidden_channels, + activation=activation, + dropout_probability=dropout_probability, + use_batch_norm=use_batch_norm, + kernel_size=kernel_size, + ) + for _ in range(num_blocks) + ] + ) + self.final_layer = nn.Conv1d( + hidden_channels, out_channels, kernel_size=kernel_size, padding="same" + ) def forward(self, inputs): temps = self.initial_layer(inputs) for block in self.blocks: temps = block(temps) outputs = self.final_layer(temps) - + return outputs + class SimilarityEmbedding(nn.Module): - ''' + """ A fully connective neural network with a ResNet layer f and an expander layer h - ''' - def __init__(self, - num_dim=3, - num_hidden_layers_f=1, - num_hidden_layers_h=1, - num_blocks=4, - kernel_size=5, - num_dim_final=10, - activation=torch.tanh, - num_channels=num_channels, - num_points=num_points - ): + """ + + def __init__( + self, + num_dim=3, + num_hidden_layers_f=1, + num_hidden_layers_h=1, + num_blocks=4, + kernel_size=5, + num_dim_final=10, + activation=torch.tanh, + num_channels=num_channels, + num_points=num_points, + ): super(SimilarityEmbedding, self).__init__() self.layer_norm = nn.LayerNorm([num_channels, num_points]) self.num_hidden_layers_f = num_hidden_layers_f self.num_hidden_layers_h = num_hidden_layers_h - self.layers_f = ResNet(num_ifos=[3,None], layers=[2,2], kernel_size=kernel_size, context_dim=100) + self.layers_f = ResNet( + num_ifos=[3, None], layers=[2, 2], kernel_size=kernel_size, context_dim=100 + ) self.contraction_layer = nn.Linear(in_features=100, out_features=num_dim) # self.layers_f = ConvResidualNet(in_channels=num_channels, out_channels=1, hidden_channels=20, num_blocks=num_blocks, kernel_size=kernel_size) # self.contraction_layer = nn.Linear(in_features=in_features, out_features=num_dim) self.expander_layer = nn.Linear(num_dim, 20) - self.layers_h = nn.ModuleList([nn.Linear(20, 20) for _ in range(num_hidden_layers_h)]) + self.layers_h = nn.ModuleList( + [nn.Linear(20, 20) for _ in range(num_hidden_layers_h)] + ) self.final_layer = nn.Linear(20, num_dim_final) self.activation = activation - + def forward(self, x): x = self.layers_f(x) x = self.contraction_layer(x) @@ -174,21 +211,23 @@ def forward(self, x): x = layer(x) x = self.activation(x) x = self.final_layer(x) - + return x, representation -def train_one_epoch_se(epoch_index, - tb_writer, - data_loader, - similarity_embedding, - optimizer, - verbose, - vicreg_loss, - **vicreg_kwargs, - ): - ''' + +def train_one_epoch_se( + epoch_index, + tb_writer, + data_loader, + similarity_embedding, + optimizer, + verbose, + vicreg_loss, + **vicreg_kwargs, +): + """ Training function - Inputs: + Inputs: epoch_index: current epoch number tb_writer: writes to tensorboard data_loader: validation data in tensor format @@ -199,23 +238,21 @@ def train_one_epoch_se(epoch_index, **vicreg_kwargs: additional loss function parameters to change loss weights Outputs: last_sim_loss: final loss calculation - ''' - running_sim_loss = 0. - last_sim_loss = 0. + """ + running_sim_loss = 0.0 + last_sim_loss = 0.0 for idx, val in enumerate(data_loader, 1): - augmented_shift, unshifted_shift, augmented_data, unshifted_data = val - augmented_shift = augmented_shift.reshape((-1,)+augmented_shift.shape[2:]) - unshifted_shift = unshifted_shift.reshape((-1,)+unshifted_shift.shape[2:]) - augmented_data = augmented_data.reshape((-1,)+augmented_data.shape[2:]) - unshifted_data = unshifted_data.reshape((-1,)+unshifted_data.shape[2:]) + augmented_shift, unshifted_shift, augmented_data, unshifted_data = val + augmented_shift = augmented_shift.reshape((-1,) + augmented_shift.shape[2:]) + unshifted_shift = unshifted_shift.reshape((-1,) + unshifted_shift.shape[2:]) + augmented_data = augmented_data.reshape((-1,) + augmented_data.shape[2:]) + unshifted_data = unshifted_data.reshape((-1,) + unshifted_data.shape[2:]) embedded_values_aug, _ = similarity_embedding(augmented_data) embedded_values_orig, _ = similarity_embedding(unshifted_data) similar_embedding_loss, _repr, _cov, _std = vicreg_loss( - embedded_values_aug, - embedded_values_orig, - **vicreg_kwargs + embedded_values_aug, embedded_values_orig, **vicreg_kwargs ) optimizer.zero_grad() similar_embedding_loss.backward() @@ -227,23 +264,30 @@ def train_one_epoch_se(epoch_index, if idx % n == 0: last_sim_loss = running_sim_loss / n if verbose == True: - print(' Avg. train loss/batch after {} batches = {:.4f}'.format(idx, last_sim_loss)) - print(f'Last {_repr.item():.2f}; {_cov.item():.2f}; {_std.item():.2f}') + print( + " Avg. train loss/batch after {} batches = {:.4f}".format( + idx, last_sim_loss + ) + ) + print(f"Last {_repr.item():.2f}; {_cov.item():.2f}; {_std.item():.2f}") tb_x = epoch_index * len(data_loader) + idx - tb_writer.add_scalar('SimLoss/train', last_sim_loss, tb_x) - running_sim_loss = 0. - + tb_writer.add_scalar("SimLoss/train", last_sim_loss, tb_x) + running_sim_loss = 0.0 + return last_sim_loss -def val_one_epoch_se(epoch_index, - tb_writer, - data_loader, - similarity_embedding, - vicreg_loss, - **vicreg_kwargs): - ''' + +def val_one_epoch_se( + epoch_index, + tb_writer, + data_loader, + similarity_embedding, + vicreg_loss, + **vicreg_kwargs, +): + """ Validation training function - Inputs: + Inputs: epoch_index: current epoch number tb_writer: writes to tensorboard data_loader: validation data in tensor format @@ -252,23 +296,21 @@ def val_one_epoch_se(epoch_index, **vicreg_kwargs: additional loss function parameters to change loss weights Outputs: last_sim_loss: final loss calculation - ''' - running_sim_loss = 0. - last_sim_loss = 0. + """ + running_sim_loss = 0.0 + last_sim_loss = 0.0 for idx, val in enumerate(data_loader, 1): augmented_shift, unshifted_shift, augmented_data, unshifted_data = val - augmented_shift = augmented_shift.reshape((-1,)+augmented_shift.shape[2:]) - unshifted_shift = unshifted_shift.reshape((-1,)+unshifted_shift.shape[2:]) - augmented_data = augmented_data.reshape((-1,)+augmented_data.shape[2:]) - unshifted_data = unshifted_data.reshape((-1,)+unshifted_data.shape[2:]) + augmented_shift = augmented_shift.reshape((-1,) + augmented_shift.shape[2:]) + unshifted_shift = unshifted_shift.reshape((-1,) + unshifted_shift.shape[2:]) + augmented_data = augmented_data.reshape((-1,) + augmented_data.shape[2:]) + unshifted_data = unshifted_data.reshape((-1,) + unshifted_data.shape[2:]) embedded_values_aug, unshifted_shift = similarity_embedding(augmented_data) embedded_values_orig, unshifted_shift = similarity_embedding(unshifted_data) similar_embedding_loss, _repr, _cov, _std = vicreg_loss( - embedded_values_aug, - embedded_values_orig, - **vicreg_kwargs + embedded_values_aug, embedded_values_orig, **vicreg_kwargs ) running_sim_loss += similar_embedding_loss.item() @@ -276,9 +318,9 @@ def val_one_epoch_se(epoch_index, if idx % n == 0: last_sim_loss = running_sim_loss / n tb_x = epoch_index * len(data_loader) + idx + 1 - tb_writer.add_scalar('SimLoss/val', last_sim_loss, tb_x) + tb_writer.add_scalar("SimLoss/val", last_sim_loss, tb_x) tb_writer.flush() - running_sim_loss = 0. + running_sim_loss = 0.0 tb_writer.flush() return last_sim_loss diff --git a/nmma/mlmodel/inference.py b/nmma/mlmodel/inference.py index 2af5dbc5e..ff5f9a656 100644 --- a/nmma/mlmodel/inference.py +++ b/nmma/mlmodel/inference.py @@ -22,34 +22,35 @@ num_points = 121 + def cast_as_bilby_result(samples, truth, priors): - - posterior = dict.fromkeys({'log10_mej', 'log10_vej', 'log10_Xlan'}) + + posterior = dict.fromkeys({"log10_mej", "log10_vej", "log10_Xlan"}) samples_numpy = samples.numpy() - posterior['log10_mej'] = samples_numpy.T[0].flatten() - posterior['log10_vej'] = samples_numpy.T[1].flatten() - posterior['log10_Xlan'] = samples_numpy.T[2].flatten() + posterior["log10_mej"] = samples_numpy.T[0].flatten() + posterior["log10_vej"] = samples_numpy.T[1].flatten() + posterior["log10_Xlan"] = samples_numpy.T[2].flatten() posterior = pd.DataFrame(posterior) - + if truth == None: result = bilby.result.Result( label="test_data", posterior=posterior, search_parameter_keys=list(posterior.keys()), - priors=priors + priors=priors, ) - else: - injections = dict.fromkeys({'log10_mej', 'log10_vej', 'log10_Xlan'}) - injections['log10_mej'] = float(truth.numpy()[0]) - injections['log10_vej'] = float(truth.numpy()[1]) - injections['log10_Xlan'] = float(truth.numpy()[2]) + else: + injections = dict.fromkeys({"log10_mej", "log10_vej", "log10_Xlan"}) + injections["log10_mej"] = float(truth.numpy()[0]) + injections["log10_vej"] = float(truth.numpy()[1]) + injections["log10_Xlan"] = float(truth.numpy()[2]) result = bilby.result.Result( label="test_data", injection_parameters=injections, posterior=posterior, search_parameter_keys=list(posterior.keys()), - priors=priors + priors=priors, ) - + return result diff --git a/nmma/mlmodel/normalizingflows.py b/nmma/mlmodel/normalizingflows.py index 9ed532558..d4edd69e7 100644 --- a/nmma/mlmodel/normalizingflows.py +++ b/nmma/mlmodel/normalizingflows.py @@ -7,6 +7,7 @@ import os, sys, time, glob import json import warnings + # pytorch import torch import torch.nn as nn @@ -14,6 +15,7 @@ from torch.utils.data import Dataset, DataLoader, TensorDataset, random_split import torch.nn.functional as F from torch.utils.tensorboard import SummaryWriter + # nflows from nflows.nn.nets.resnet import ResidualNet from nflows import transforms, distributions, flows @@ -22,6 +24,7 @@ from nflows.transforms.autoregressive import MaskedAffineAutoregressiveTransform from nflows.transforms import CompositeTransform, RandomPermutation import nflows.utils as torchutils + # functions from .embedding import SimilarityEmbedding @@ -31,9 +34,11 @@ similarity_embedding = SimilarityEmbedding() num_points = 121 + class EmbeddingNet(nn.Module): """Wrapper around the similarity embedding defined above""" - #def __init__(self, *args, **kwargs): + + # def __init__(self, *args, **kwargs): def __init__(self, similarity_embedding, *args, **kwargs): super().__init__(*args, **kwargs) self.representation_net = similarity_embedding @@ -41,54 +46,63 @@ def __init__(self, similarity_embedding, *args, **kwargs): # self.num_dim = num_dim # the expander network is unused and hence don't track gradients for name, param in self.representation_net.named_parameters(): - if 'expander_layer' in name or 'layers_h' in name or 'final_layer' in name: + if "expander_layer" in name or "layers_h" in name or "final_layer" in name: param.requires_grad = False - # set freeze status of part of the conv layer of embedding_net - elif 'layers_f' in name: + # set freeze status of part of the conv layer of embedding_net + elif "layers_f" in name: param.requires_grad = True - else: + else: param.requires_grad = True self.context_layer = nn.Sequential( - nn.Linear(num_dim, 1000), - nn.ReLU(), - nn.Linear(1000, num_dim) - ) + nn.Linear(num_dim, 1000), nn.ReLU(), nn.Linear(1000, num_dim) + ) + def forward(self, x): batch_size, channels, dims = x.shape - _, rep = self.representation_net(x) # 500, 3, 191 -> 500, 1, 2 + _, rep = self.representation_net(x) # 500, 3, 191 -> 500, 1, 2 rep = rep.reshape(batch_size, num_dim) return self.context_layer(rep) -def normflow_params(similarity_embedding, num_transforms, num_blocks, hidden_features, context_features, num_dim): + +def normflow_params( + similarity_embedding, + num_transforms, + num_blocks, + hidden_features, + context_features, + num_dim, +): base_dist = StandardNormal([3]) transforms = [] features = 3 for _ in range(num_transforms): - block = [MaskedAffineAutoregressiveTransform( - features=3, + block = [ + MaskedAffineAutoregressiveTransform( + features=3, hidden_features=hidden_features, # 80 context_features=context_features, # 5 - num_blocks=num_blocks, # 5 + num_blocks=num_blocks, # 5 activation=torch.tanh, use_batch_norm=False, use_residual_blocks=True, dropout_probability=0.01, - # integrand_net_layers=[20, 20] + # integrand_net_layers=[20, 20] ), - RandomPermutation(features=features) + RandomPermutation(features=features), ] transforms += block transform = CompositeTransform(transforms) embedding_net = EmbeddingNet(similarity_embedding) return transform, base_dist, embedding_net + def train_one_epoch(epoch_index, tb_writer, data_loader, flow, optimizer, flatten_dim): - running_loss = 0. - last_loss = 0. + running_loss = 0.0 + last_loss = 0.0 for idx, val in enumerate(data_loader, 1): augmented_shift, unshifted_shift, augmented_data, unshifted_data = val - augmented_shift = augmented_shift[...,0:3].to(device) + augmented_shift = augmented_shift[..., 0:3].to(device) augmented_shift = augmented_shift.flatten(0, flatten_dim).to(device) augmented_data = augmented_data.reshape(-1, 3, num_points).to(device) loss = 0 @@ -101,19 +115,24 @@ def train_one_epoch(epoch_index, tb_writer, data_loader, flow, optimizer, flatte n = 10 if idx % n == 0: last_loss = running_loss / n - print(' Avg. train loss/batch after {} batches = {:.4f}'.format(idx, last_loss)) + print( + " Avg. train loss/batch after {} batches = {:.4f}".format( + idx, last_loss + ) + ) tb_x = epoch_index * len(data_loader) + idx - tb_writer.add_scalar('Flow Loss/train', last_loss, tb_x) + tb_writer.add_scalar("Flow Loss/train", last_loss, tb_x) tb_writer.flush() - running_loss = 0. + running_loss = 0.0 return last_loss + def val_one_epoch(epoch_index, tb_writer, data_loader, flow, flatten_dim): - running_loss = 0. - last_loss = 0. + running_loss = 0.0 + last_loss = 0.0 for idx, val in enumerate(data_loader, 1): augmented_shift, unshifted_shift, augmented_data, unshifted_data = val - augmented_shift = augmented_shift[...,0:3].to(device) + augmented_shift = augmented_shift[..., 0:3].to(device) augmented_shift = augmented_shift.flatten(0, flatten_dim).to(device) augmented_data = augmented_data.reshape(-1, 3, num_points).to(device) loss = 0 @@ -123,10 +142,14 @@ def val_one_epoch(epoch_index, tb_writer, data_loader, flow, flatten_dim): n = 1 if idx % n == 0: last_loss = running_loss / n - print(' Avg. train loss/batch after {} batches = {:.4f}'.format(idx, last_loss)) + print( + " Avg. train loss/batch after {} batches = {:.4f}".format( + idx, last_loss + ) + ) tb_x = epoch_index * len(data_loader) + idx + 1 - tb_writer.add_scalar('Flow Loss/val', last_loss, tb_x) + tb_writer.add_scalar("Flow Loss/val", last_loss, tb_x) tb_writer.flush() - running_loss = 0. + running_loss = 0.0 tb_writer.flush() return last_loss diff --git a/nmma/mlmodel/resnet.py b/nmma/mlmodel/resnet.py index 533096c4a..0488a26d5 100644 --- a/nmma/mlmodel/resnet.py +++ b/nmma/mlmodel/resnet.py @@ -38,9 +38,7 @@ def forward(self, x): # reshape back to full channel dimension if self.num_groups != self.num_channels: mean = torch.stack([mean, sq_mean], dim=1) - mean = mean.reshape( - -1, 2, self.num_groups, self.channels_per_group - ) + mean = mean.reshape(-1, 2, self.num_groups, self.channels_per_group) mean = mean.mean(-1, keepdims=True) mean = mean.expand(-1, -1, -1, self.channels_per_group) mean = mean.reshape(-1, 2, self.num_channels, 1) @@ -90,9 +88,7 @@ def convN( def conv1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv1d: """kernel-size 1 convolution""" - return nn.Conv1d( - in_planes, out_planes, kernel_size=1, stride=stride, bias=False - ) + return nn.Conv1d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) class BasicBlock(nn.Module): @@ -115,13 +111,9 @@ def __init__( if norm_layer is None: norm_layer = nn.BatchNorm1d if groups != 1 or base_width != 64: - raise ValueError( - "BasicBlock only supports groups=1 and base_width=64" - ) + raise ValueError("BasicBlock only supports groups=1 and base_width=64") if dilation > 1: - raise NotImplementedError( - "Dilation > 1 not supported in BasicBlock" - ) + raise NotImplementedError("Dilation > 1 not supported in BasicBlock") # Both self.conv1 and self.downsample layers # downsample the input when stride != 1 @@ -367,9 +359,7 @@ def __init__( for m in self.modules(): if isinstance(m, nn.Conv1d): - nn.init.kaiming_normal_( - m.weight, mode="fan_out", nonlinearity="relu" - ) + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") elif isinstance(m, (nn.BatchNorm1d, nn.GroupNorm)): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) diff --git a/nmma/pbilby/analysis/analysis_run.py b/nmma/pbilby/analysis/analysis_run.py index c92b943bf..797866a79 100644 --- a/nmma/pbilby/analysis/analysis_run.py +++ b/nmma/pbilby/analysis/analysis_run.py @@ -79,8 +79,8 @@ def __init__( logger.setLevel(logging.WARNING) # split the likelihood for difference inference_favour - assert inference_favour in ['nmma', 'nmma_gw'], "Invalid inference_favour" - if inference_favour == 'nmma': + assert inference_favour in ["nmma", "nmma_gw"], "Invalid inference_favour" + if inference_favour == "nmma": light_curve_data = data_dump["light_curve_data"] likelihood, priors = setup_nmma_likelihood( interferometers=ifo_list, @@ -90,7 +90,7 @@ def __init__( args=args, ) - elif inference_favour == 'nmma_gw': + elif inference_favour == "nmma_gw": likelihood, priors = setup_nmma_gw_likelihood( interferometers=ifo_list, waveform_generator=waveform_generator, diff --git a/nmma/pbilby/analysis/main.py b/nmma/pbilby/analysis/main.py index ff35c5299..328516fdf 100644 --- a/nmma/pbilby/analysis/main.py +++ b/nmma/pbilby/analysis/main.py @@ -29,8 +29,8 @@ from ..parser import ( create_nmma_analysis_parser, create_nmma_gw_analysis_parser, - parse_analysis_args - ) + parse_analysis_args, +) from .analysis_run import AnalysisRun @@ -270,7 +270,7 @@ def analysis_runner( nested_samples, sampler_kwargs, sampling_time, - rejection_sample_posterior=rejection_sample_posterior + rejection_sample_posterior=rejection_sample_posterior, ) posterior = conversion.fill_from_fixed_priors( @@ -281,10 +281,10 @@ def analysis_runner( "Generating posterior from marginalized parameters for" f" nsamples={len(posterior)}, POOL={pool.size}" ) - #fill_args = [ + # fill_args = [ # (ii, row, run.likelihood) for ii, row in posterior.iterrows() - #] - #samples = pool.map(fill_sample, fill_args) + # ] + # samples = pool.map(fill_sample, fill_args) posterior, _ = run.likelihood.parameter_conversion( posterior, ) @@ -343,7 +343,7 @@ def main_nmma(): input_args = parse_analysis_args(analysis_parser, cli_args=cli_args) # Run the analysis - analysis_runner(**vars(input_args), inference_favour='nmma') + analysis_runner(**vars(input_args), inference_favour="nmma") def main_nmma_gw(): @@ -360,4 +360,4 @@ def main_nmma_gw(): input_args = parse_analysis_args(analysis_parser, cli_args=cli_args) # Run the analysis - analysis_runner(**vars(input_args), inference_favour='nmma_gw') + analysis_runner(**vars(input_args), inference_favour="nmma_gw") diff --git a/nmma/pbilby/parser/__init__.py b/nmma/pbilby/parser/__init__.py index 6291cde8a..a13c48927 100644 --- a/nmma/pbilby/parser/__init__.py +++ b/nmma/pbilby/parser/__init__.py @@ -1,10 +1,10 @@ from .analysis import ( create_nmma_analysis_parser, create_nmma_gw_analysis_parser, - parse_analysis_args - ) + parse_analysis_args, +) from .generation import ( create_nmma_generation_parser, create_nmma_gw_generation_parser, - parse_generation_args - ) + parse_generation_args, +) diff --git a/nmma/pbilby/parser/generation.py b/nmma/pbilby/parser/generation.py index bd48f236f..e8ae770b6 100644 --- a/nmma/pbilby/parser/generation.py +++ b/nmma/pbilby/parser/generation.py @@ -4,8 +4,8 @@ from .shared import ( _create_base_nmma_gw_parser, _create_base_nmma_parser, - _add_slurm_settings_to_parser - ) + _add_slurm_settings_to_parser, +) logger = bilby.core.utils.logger diff --git a/nmma/tests/analysis.py b/nmma/tests/analysis.py index 174128afd..161623c1f 100644 --- a/nmma/tests/analysis.py +++ b/nmma/tests/analysis.py @@ -110,7 +110,7 @@ def test_analysis_systematics_with_time_and_filters(args): def test_analysis_systematics_without_time(args): - + args.filters = "ztfr" args.systematics_file = f"{DATA_DIR}/systematics_without_time.yaml" analysis.main(args) @@ -135,7 +135,7 @@ def test_nn_analysis(args): args.sampler = "neuralnet" args.prior = "priors/Ka2017.prior" args.dt = 0.25 - args.dt_inj= 0.25 + args.dt_inj = 0.25 args.filters = "ztfg,ztfr,ztfi" args.local_only = False args.injection = f"{DATA_DIR}/Ka2017_injection.json" diff --git a/nmma/tests/analysis_lbol.py b/nmma/tests/analysis_lbol.py index e56642a1e..174b907c1 100644 --- a/nmma/tests/analysis_lbol.py +++ b/nmma/tests/analysis_lbol.py @@ -16,7 +16,7 @@ def args(): outdir="outdir", label="lbol_test", trigger_time=60168.79041667, - data='example_files/lbol/ztf23bqun/23bqun_bbdata.csv', + data="example_files/lbol/ztf23bqun/23bqun_bbdata.csv", prior="example_files/lbol/ztf23bqun/Arnett_modified.priors", tmin=0.005, tmax=20.0, @@ -45,7 +45,7 @@ def args(): ra=None, dec=None, fetch_Ebv_from_dustmap=False, - systematics_file=None + systematics_file=None, ) return args diff --git a/nmma/tests/data/multi_config_analysis/injection.json b/nmma/tests/data/multi_config_analysis/injection.json index ac74590d4..fd61af7b6 100644 --- a/nmma/tests/data/multi_config_analysis/injection.json +++ b/nmma/tests/data/multi_config_analysis/injection.json @@ -28,4 +28,4 @@ ] } } -} \ No newline at end of file +} diff --git a/nmma/tests/injections.py b/nmma/tests/injections.py index 087815270..6d8f05524 100644 --- a/nmma/tests/injections.py +++ b/nmma/tests/injections.py @@ -34,12 +34,12 @@ def lightcurveInjectionTest(model_name, model_lightcurve_function): "current working directory: ", os.getcwd() ) # assumes run in root nmma folder, will need to modify if this is not true - workingDir=os.path.dirname(__file__) - dataDir = os.path.join(workingDir, 'data') + workingDir = os.path.dirname(__file__) + dataDir = os.path.join(workingDir, "data") test_directory = os.path.join(dataDir, model_name) - priorDir=os.path.join(workingDir, '../../priors/') - svdmodels=os.path.join(workingDir, '../../svdmodels/') - + priorDir = os.path.join(workingDir, "../../priors/") + svdmodels = os.path.join(workingDir, "../../svdmodels/") + if os.path.isdir(test_directory): shutil.rmtree(test_directory, ignore_errors=True) os.makedirs(test_directory, exist_ok=True) @@ -59,11 +59,13 @@ def create_injection_from_command_line(model_name): """ if model_name == "nugent-hyper": - prior_path = os.path.join( priorDir, "sncosmo-generic" + ".prior") + prior_path = os.path.join(priorDir, "sncosmo-generic" + ".prior") elif model_name == "TrPi2018": - prior_path = os.path.join( dataDir, "TrPi2018_pinned_parameters" + ".prior") #pinning the parameter svalues in the prior file + prior_path = os.path.join( + dataDir, "TrPi2018_pinned_parameters" + ".prior" + ) # pinning the parameter svalues in the prior file else: - prior_path = os.path.join( priorDir, model_name + ".prior") + prior_path = os.path.join(priorDir, model_name + ".prior") assert os.path.exists(prior_path), "prior file does not exist" injection_name = os.path.join(test_directory, model_name + "_injection.json") @@ -303,30 +305,41 @@ def test_injections(): for model_name, model_lightcurve_function in lightcurve_models.items(): lightcurveInjectionTest(model_name, model_lightcurve_function) + def test_validate_lightcurves(): print("validate_lightcurve test") ## initialize args, check a file that is known to have 3 observations in the ztf g filter and 1 in the ztf r filter. All detections occur within 9 days of the original observation. args = Namespace( - data='example_files/candidate_data/ZTF20abwysqy.dat', + data="example_files/candidate_data/ZTF20abwysqy.dat", filters="ztfg", min_obs=3, cutoff_time=0, silent=False, ) - assert validate_lightcurve(**vars(args)) == True, "Test for 3 observations in the ztf g filter failed" + assert ( + validate_lightcurve(**vars(args)) == True + ), "Test for 3 observations in the ztf g filter failed" args.filters = "ztfr" args.min_obs = 1 - assert validate_lightcurve(**vars(args)) == True, "Test for 1 observation in the ztf r filter failed" + assert ( + validate_lightcurve(**vars(args)) == True + ), "Test for 1 observation in the ztf r filter failed" args.filters = "ztfg,ztfr" - assert validate_lightcurve(**vars(args)) == True, "Test for passing multiple filters failed" + assert ( + validate_lightcurve(**vars(args)) == True + ), "Test for passing multiple filters failed" args.filters = "" args.min_obs = 0 - assert validate_lightcurve(**vars(args)) == True, "Test for automatic filter selection failed" + assert ( + validate_lightcurve(**vars(args)) == True + ), "Test for automatic filter selection failed" args.cutoff_time = 1 args.min_obs = 1 - assert validate_lightcurve(**vars(args)) == False, "Test for setting cutoff time failed" + assert ( + validate_lightcurve(**vars(args)) == False + ), "Test for setting cutoff time failed" diff --git a/nmma/tests/maximum_mass.py b/nmma/tests/maximum_mass.py index 403307b70..364a8765e 100644 --- a/nmma/tests/maximum_mass.py +++ b/nmma/tests/maximum_mass.py @@ -6,6 +6,7 @@ from ..joint import maximum_mass_constraint + @pytest.fixture(scope="module") def args(): workingDir = os.path.dirname(__file__) @@ -15,17 +16,17 @@ def args(): args = Namespace( outdir="outdir", - prior = f"{priorDir}/maximum_mass_resampling.prior", - joint_posterior = f"{dataDir}/GW+KN+GRB_posterior", - eos_path_macro = f"{dataDir}/eos_macro", - eos_path_micro = f"{dataDir}/eos_micro", - nlive = 32, - use_M_Kepler = False + prior=f"{priorDir}/maximum_mass_resampling.prior", + joint_posterior=f"{dataDir}/GW+KN+GRB_posterior", + eos_path_macro=f"{dataDir}/eos_macro", + eos_path_micro=f"{dataDir}/eos_micro", + nlive=32, + use_M_Kepler=False, ) return args def test_maximum_mass_resampling(args): - + maximum_mass_constraint.main(args) diff --git a/nmma/tests/multi_config_analysis.py b/nmma/tests/multi_config_analysis.py index 5797546d2..b14443982 100644 --- a/nmma/tests/multi_config_analysis.py +++ b/nmma/tests/multi_config_analysis.py @@ -9,7 +9,6 @@ WORKING_DIR = os.path.dirname(__file__) - @pytest.fixture(scope="module") def args(): WORKING_DIR = os.path.dirname(__file__) @@ -24,6 +23,7 @@ def args(): def test_analysis_multi(args): multi_config_analysis.main(args) + # Randomly fails on GitHub runners # def test_analysis_parallel(args): # args.parallel = True diff --git a/nmma/tests/systematics.py b/nmma/tests/systematics.py index 2191c9f19..a6315a8ed 100644 --- a/nmma/tests/systematics.py +++ b/nmma/tests/systematics.py @@ -10,7 +10,7 @@ handle_withoutTime, main, ALLOWED_FILTERS, - ALLOWED_DISTRIBUTIONS + ALLOWED_DISTRIBUTIONS, ) @@ -48,8 +48,12 @@ def test_validate_only_one_true_valid(sample_yaml_file): def test_validate_only_one_true_invalid(): - invalid_yaml = {"config": {"withTime": {"value": True}, "withoutTime": {"value": True}}} - with pytest.raises(ValidationError, match="Only one configuration key can be set to True at a time"): + invalid_yaml = { + "config": {"withTime": {"value": True}, "withoutTime": {"value": True}} + } + with pytest.raises( + ValidationError, match="Only one configuration key can be set to True at a time" + ): validate_only_one_true(invalid_yaml) @@ -65,7 +69,13 @@ def test_validate_filters_invalid(): def test_handle_withTime(): - values = {"type": "Uniform", "minimum": 0.0, "maximum": 1.0, "time_nodes": 2, "filters": [["bessellb", "bessellv"], "ztfr"]} + values = { + "type": "Uniform", + "minimum": 0.0, + "maximum": 1.0, + "time_nodes": 2, + "filters": [["bessellb", "bessellv"], "ztfr"], + } result = handle_withTime(values) assert "sys_err_bessellb___bessellv1" in result[0] assert "sys_err_ztfr2" in result[3] @@ -75,7 +85,10 @@ def test_handle_withoutTime(): values = {"type": "Uniform", "minimum": 0.0, "maximum": 1.0} result = handle_withoutTime(values) assert len(result) == 1 - assert "sys_err = Uniform(minimum=0.0, maximum=1.0, name='sys_err', latex_label='sys_err', unit=None, boundary=None)" in result[0] + assert ( + "sys_err = Uniform(minimum=0.0, maximum=1.0, name='sys_err', latex_label='sys_err', unit=None, boundary=None)" + in result[0] + ) def test_main(sample_yaml_file): @@ -91,26 +104,37 @@ def test_main_invalid_yaml(tmp_path): def test_validate_only_one_true_all_false(): - invalid_yaml = {"config": {"withTime": {"value": False}, "withoutTime": {"value": False}}} - with pytest.raises(ValidationError, match="At least one configuration key must be set to True"): + invalid_yaml = { + "config": {"withTime": {"value": False}, "withoutTime": {"value": False}} + } + with pytest.raises( + ValidationError, match="At least one configuration key must be set to True" + ): validate_only_one_true(invalid_yaml) def test_validate_only_one_true_missing_value(): invalid_yaml = {"config": {"withTime": {}, "withoutTime": {"value": False}}} - with pytest.raises(ValidationError, match="'value' key must be present and be a boolean"): + with pytest.raises( + ValidationError, match="'value' key must be present and be a boolean" + ): validate_only_one_true(invalid_yaml) def test_validate_filters_duplicate_in_group(): invalid_filters = [["bessellb", "bessellb"], "ztfr"] - with pytest.raises(ValidationError, match="Duplicate filter value 'bessellb' within the same group"): + with pytest.raises( + ValidationError, match="Duplicate filter value 'bessellb' within the same group" + ): validate_filters(invalid_filters) def test_validate_filters_duplicate_across_groups(): invalid_filters = [["bessellb", "bessellv"], "bessellb"] - with pytest.raises(ValidationError, match="Duplicate filter value 'bessellb'. A filter can only be used in one group"): + with pytest.raises( + ValidationError, + match="Duplicate filter value 'bessellb'. A filter can only be used in one group", + ): validate_filters(invalid_filters) @@ -134,18 +158,30 @@ def test_validate_distribution_invalid(): def test_validate_distribution_case_sensitive(): with pytest.raises(KeyError): - assert (ALLOWED_DISTRIBUTIONS["uniform"]) # Should be "Uniform" + assert ALLOWED_DISTRIBUTIONS["uniform"] # Should be "Uniform" def test_handle_withTime_single_filter(): - values = {"type": "Uniform", "minimum": 0.0, "maximum": 1.0, "time_nodes": 2, "filters": ["ztfr"]} + values = { + "type": "Uniform", + "minimum": 0.0, + "maximum": 1.0, + "time_nodes": 2, + "filters": ["ztfr"], + } result = handle_withTime(values) assert len(result) == 2 assert all("sys_err_ztfr" in line for line in result) def test_handle_withTime_all_filters(): - values = {"type": "Uniform", "minimum": 0.0, "maximum": 1.0, "time_nodes": 1, "filters": [None]} + values = { + "type": "Uniform", + "minimum": 0.0, + "maximum": 1.0, + "time_nodes": 1, + "filters": [None], + } result = handle_withTime(values) assert len(result) == 1 assert "sys_err_all1" in result[0] @@ -179,13 +215,21 @@ def test_main_empty_config(tmp_path): """ yaml_file = tmp_path / "empty_config.yaml" yaml_file.write_text(yaml_content) - with pytest.raises(ValidationError, match="At least one configuration key must be set to True"): + with pytest.raises( + ValidationError, match="At least one configuration key must be set to True" + ): main(yaml_file) @pytest.mark.parametrize("filter_name", ALLOWED_FILTERS) def test_all_allowed_filters(filter_name): - values = {"type": "Uniform", "minimum": 0.0, "maximum": 1.0, "time_nodes": 1, "filters": [filter_name]} + values = { + "type": "Uniform", + "minimum": 0.0, + "maximum": 1.0, + "time_nodes": 1, + "filters": [filter_name], + } result = handle_withTime(values) assert len(result) == 1 assert f"sys_err_{filter_name}1" in result[0] diff --git a/nmma/tests/utils.py b/nmma/tests/utils.py index 4214500f9..98a28b6ca 100644 --- a/nmma/tests/utils.py +++ b/nmma/tests/utils.py @@ -15,9 +15,9 @@ def test_lightcurve_HoNa_basic(): n = 3.5 L, T, r = lightcurve_HoNa(t, mass, velocities, opacities, n) - assert L.shape == T.shape == r.shape == t.shape, ( - "Output shapes must match input time array" - ) + assert ( + L.shape == T.shape == r.shape == t.shape + ), "Output shapes must match input time array" assert np.all(L > 0), "Luminosity must be positive" assert np.all(T > 0), "Temperature must be positive" assert np.all(r > 0), "Radius must be positive" @@ -55,6 +55,6 @@ def test_lightcurve_HoNa_types(): n = 3.5 with pytest.raises( AssertionError, - match=re.escape("Expected: mass=float, velocities/opacities=list[float]") + match=re.escape("Expected: mass=float, velocities/opacities=list[float]"), ): lightcurve_HoNa(t, mass, velocities, opacities, n) diff --git a/nmma/utils/gitlab.py b/nmma/utils/gitlab.py index 96689280d..5db917f57 100644 --- a/nmma/utils/gitlab.py +++ b/nmma/utils/gitlab.py @@ -15,18 +15,22 @@ try: from mpi4py import MPI + mpi_enabled = True except ImportError: mpi_enabled = False + def download_and_decompress(file_info): download(file_info) decompress(file_info[1]) + def mpi_barrier(comm): if mpi_enabled: comm.Barrier() + def download_models_list(models_home=None): # first we load the models list from gitlab models_home = get_models_home(models_home) @@ -36,6 +40,7 @@ def download_models_list(models_home=None): with open(Path(models_home, "models.yaml"), "wb") as f: f.write(r.content) + def load_models_list(models_home=None): models_home = get_models_home(models_home) @@ -102,6 +107,7 @@ def load_models_list(models_home=None): return models, downloaded_if_missing is False + def refresh_models_list(models_home=None): global MODELS models_home = get_models_home(models_home) @@ -115,6 +121,7 @@ def refresh_models_list(models_home=None): raise ValueError(f"Could not load models list: {str(e)}") return models + def get_model( models_home=None, model_name=None, diff --git a/nmma/utils/models.py b/nmma/utils/models.py index 2aa73efad..b0d89c1bb 100644 --- a/nmma/utils/models.py +++ b/nmma/utils/models.py @@ -2,14 +2,17 @@ try: from mpi4py import MPI + mpi_enabled = True except ImportError: mpi_enabled = False + def mpi_barrier(comm): if mpi_enabled: comm.Barrier() + def refresh_models_list(models_home=None, source=None): if source is None: diff --git a/priors/HoNa2020.prior b/priors/HoNa2020.prior index 5d4e56979..2d4184b0d 100644 --- a/priors/HoNa2020.prior +++ b/priors/HoNa2020.prior @@ -1,8 +1,8 @@ luminosity_distance = Uniform(minimum=0.0, maximum=200, name="luminosity_distance",latex_label='$D_L$') -log10_Mej = Uniform(name="log10_Mej", minimum=-3, maximum=-0.5, latex_label='$\\log_{10}(M_{\\rm{ej}})$') +log10_mej = Uniform(name="log10_mej", minimum=-3, maximum=-0.5, latex_label='$\\log_{10}(M_{\\rm{ej}})$') vej_max = Uniform(name="vej_max", minimum=0.2, maximum=0.4, latex_label='$v_{\\rm{max}}$') vej_min = Uniform(name="vej_min", minimum=0.1, maximum=0.2, latex_label='$v_{\\rm{min}}$') vej_frac = Uniform(name="vej_frac", minimum=0, maximum=1, latex_label='$f_{\\rm{ej}}$') log10_kappa_low_vej = Uniform(name="log10_kappa_low_vej", minimum=-1, maximum=0, latex_label='$\\log_{10}(\\kappa_{\\rm{low}})$') log10_kappa_high_vej = Uniform(name="log10_kappa_high_vej", minimum=0, maximum=1, latex_label='$\\log_{10}(\\kappa_{\\rm{high}})$') -timeshift = 0 \ No newline at end of file +timeshift = 0 diff --git a/priors/maximum_mass_resampling.prior b/priors/maximum_mass_resampling.prior index 2fdec3e14..0e557b30d 100644 --- a/priors/maximum_mass_resampling.prior +++ b/priors/maximum_mass_resampling.prior @@ -1,4 +1,4 @@ chirp_mass = Uniform(name = "chirp_mass", minimum = 1.184, maximum = 1.19) eta_star = Uniform(name = "eta_star", minimum = -20, maximum = -5.5) log10_mdisk = Uniform(name = "log10_mdisk", minimum = -1.4, maximum = -0.6) -log10_mej_dyn = Uniform(name = "log10_mej_dyn", minimum = -2.75, maximum = -2.05) \ No newline at end of file +log10_mej_dyn = Uniform(name = "log10_mej_dyn", minimum = -2.75, maximum = -2.05) diff --git a/requirements.txt b/requirements.txt index d78c18dbd..0dd9bee3d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,4 +15,5 @@ p_tqdm tornado notebook ligo.skymap -healpy \ No newline at end of file +healpy +dustmaps diff --git a/sampler_requirements.txt b/sampler_requirements.txt index 5a8ecaeff..6fd135210 100644 --- a/sampler_requirements.txt +++ b/sampler_requirements.txt @@ -1 +1 @@ -ultranest \ No newline at end of file +ultranest diff --git a/sklearn_requirements.txt b/sklearn_requirements.txt index d356df703..50bb24704 100644 --- a/sklearn_requirements.txt +++ b/sklearn_requirements.txt @@ -1 +1 @@ -scikit-learn>=1.0.2 \ No newline at end of file +scikit-learn>=1.0.2 diff --git a/tf_requirements.txt b/tf_requirements.txt index b3a60b6da..0f5714408 100644 --- a/tf_requirements.txt +++ b/tf_requirements.txt @@ -1 +1 @@ -tensorflow \ No newline at end of file +tensorflow diff --git a/tools/corner_plot.py b/tools/corner_plot.py index 74891042f..9df779c63 100644 --- a/tools/corner_plot.py +++ b/tools/corner_plot.py @@ -64,13 +64,17 @@ def plotting_parameters(prior_filename, filename_with_fullpath, verbose): # ignore prior if it is a fixed value if re.match(r"^\s*-?[\d.]+\s*$", key_value[1]): continue - latex_label_match = re.search(r"latex_label\s*=\s*(['\"])(.*?)\1", key_value[1]) + latex_label_match = re.search( + r"latex_label\s*=\s*(['\"])(.*?)\1", key_value[1] + ) # use latex label if it exists, otherwise use the name if latex_label_match: latex_label_value = latex_label_match.group(2) else: - latex_label_value = re.search(r"name\s*=\s*['\"]([^'\"]+)['\"]", key_value[1]).group(1) + latex_label_value = re.search( + r"name\s*=\s*['\"]([^'\"]+)['\"]", key_value[1] + ).group(1) parameters[key] = latex_label_value @@ -84,7 +88,7 @@ def plotting_parameters(prior_filename, filename_with_fullpath, verbose): common_params = list(prior_params & posterior_params) - common_params_dict = {k: parameters[k] for k in common_params} + common_params_dict = {k: parameters[k] for k in common_params} return common_params_dict @@ -113,7 +117,9 @@ def load_csv(filename_with_fullpath, prior_filename, verbose): return samples -def load_injection(prior_filename, injection_file_json, injection_num, filename_with_fullpath, verbose): +def load_injection( + prior_filename, injection_file_json, injection_num, filename_with_fullpath, verbose +): """ Load injection data from a JSON file. Parameters @@ -129,12 +135,14 @@ def load_injection(prior_filename, injection_file_json, injection_num, filename_ """ df = pd.read_json(injection_file_json) df = df.from_records(df["injections"]["content"]) - columns = plotting_parameters(prior_filename, filename_with_fullpath, verbose).keys() + columns = plotting_parameters( + prior_filename, filename_with_fullpath, verbose + ).keys() df = df[[col for col in columns if col in df.columns]] truths = np.vstack(df.iloc[injection_num].values).flatten() if verbose: print("\nLoaded Injection:") - print(f"Truths from injection: {truths}") + print(f"Truths from injection: {truths}") return truths @@ -153,7 +161,9 @@ def load_bestfit(prior_filename, bestfit_file_json, filename_with_fullpath, verb A 1D numpy array representing the bestfit params to be used as truths. """ df = pd.read_json(bestfit_file_json, typ="series") - columns = plotting_parameters(prior_filename, filename_with_fullpath, verbose).keys() + columns = plotting_parameters( + prior_filename, filename_with_fullpath, verbose + ).keys() df = df[[col for col in columns if col in df.keys()]] truths = np.vstack(df.values).flatten() if verbose: @@ -282,7 +292,9 @@ def corner_plot(data, labels, filename, truths, legendlabel, ext, verbose, **kwa lw=2, ) ) - axes[2 * int(np.sqrt(len(axes))) - 3].legend(lines, legendlabel, loc=3, frameon=True, fancybox=True) + axes[2 * int(np.sqrt(len(axes))) - 3].legend( + lines, legendlabel, loc=3, frameon=True, fancybox=True + ) if len(data) == 2: title_quantiles_1 = [] title_quantiles_2 = [] @@ -298,7 +310,9 @@ def corner_plot(data, labels, filename, truths, legendlabel, ext, verbose, **kwa axes[i].text( x=coords[0] - 0.05, y=1.05 * coords[1], - s=r"${{{0:.2f}}}_{{-{1:.2f}}}^{{+{2:.2f}}}$ ".format(*title_quantiles_1[i // np.sqrt(len(axes)).astype(int)]), + s=r"${{{0:.2f}}}_{{-{1:.2f}}}^{{+{2:.2f}}}$ ".format( + *title_quantiles_1[i // np.sqrt(len(axes)).astype(int)] + ), ha="right", color=color_array[0], transform=axes[i].transAxes, @@ -306,7 +320,9 @@ def corner_plot(data, labels, filename, truths, legendlabel, ext, verbose, **kwa axes[i].text( x=coords[0] + 0.05, y=1.05 * coords[1], - s=r"${{{0:.2f}}}_{{-{1:.2f}}}^{{+{2:.2f}}}$".format(*title_quantiles_2[i // np.sqrt(len(axes)).astype(int)]), + s=r"${{{0:.2f}}}_{{-{1:.2f}}}^{{+{2:.2f}}}$".format( + *title_quantiles_2[i // np.sqrt(len(axes)).astype(int)] + ), ha="left", color=color_array[1], transform=axes[i].transAxes, @@ -391,7 +407,8 @@ def corner_plot(data, labels, filename, truths, legendlabel, ext, verbose, **kwa "-v", default=False, action="store_true", - help="Print additional information") + help="Print additional information", + ) args = parser.parse_args() posterior_files = args.posterior_files @@ -407,11 +424,14 @@ def corner_plot(data, labels, filename, truths, legendlabel, ext, verbose, **kwa if not additional_kwargs: print("\nNo additional kwargs provided") - else: + else: print("\nRunning with the following additional kwargs:") - print("\t\n".join(f" - {key}: {value}" for key, value in additional_kwargs.items())) + print( + "\t\n".join( + f" - {key}: {value}" for key, value in additional_kwargs.items() + ) + ) - # Generate legend labels from input file names legendlabel = [] if label_name is not None: @@ -428,7 +448,9 @@ def corner_plot(data, labels, filename, truths, legendlabel, ext, verbose, **kwa posteriors.append(posterior) if injection_json is not None: - truths = load_injection(prior_filename, injection_json, injection_num, posterior_files[0], verbose) + truths = load_injection( + prior_filename, injection_json, injection_num, posterior_files[0], verbose + ) elif args.bestfit_params is not None: truths = load_bestfit(prior_filename, bestfit_json, posterior_files[0], verbose) else: @@ -442,7 +464,6 @@ def corner_plot(data, labels, filename, truths, legendlabel, ext, verbose, **kwa plot_density=False, plot_contours=True, fill_contours=True, - label_kwargs={"fontsize": 16}, levels=[0.16, 0.5, 0.84], smooth=1, @@ -450,14 +471,16 @@ def corner_plot(data, labels, filename, truths, legendlabel, ext, verbose, **kwa kwargs.update(additional_kwargs) - # the code assumes that the parameters in rest of the posterior files are the same as the first posterior file. and the prior file and posterior files have the same parameters which can be plotted + # the code assumes that the parameters in rest of the posterior files are the same as the first posterior file. and the prior file and posterior files have the same parameters which can be plotted if verbose: print(f"\nParameters and Axis labels ({len(labels)} common parameters):") for k, v in labels.items(): - print(f" - {k}: {v}") + print(f" - {k}: {v}") - corner_plot(posteriors, labels, output_filename, truths, legendlabel, ext, verbose, **kwargs) + corner_plot( + posteriors, labels, output_filename, truths, legendlabel, ext, verbose, **kwargs + ) ## Example usage # python corner_plot.py -f GRB_res12_linear2dp/injection_posterior_samples.dat GRB_res12_linear4dp/injection_posterior_samples.dat -p GRB170817A_emsys_4dp.prior -o linear2d_vs_linear4dp --kwargs "{'levels':[0.05,0.5,0.95]}" diff --git a/tutorials/README.md b/tutorials/README.md index 8b1378917..e69de29bb 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -1 +0,0 @@ -