diff --git a/abcmb/main.py b/abcmb/main.py index 8eaaed6..a3dfe7e 100644 --- a/abcmb/main.py +++ b/abcmb/main.py @@ -59,7 +59,7 @@ class Model(eqx.Module): Methods: -------- - run_cosmology : Compute CMB angular power spectra + __call__ : Compute CMB angular power spectra get_PTBG : Get perturbation table and background cosmology get_BG : Get background cosmology add_derived_parameters : Compute derived parameters @@ -148,7 +148,7 @@ def __init__(self, # need this outside of the jit context # since we want LINX to run on CPU - def run_cosmology(self, params : dict = {}): + def __call__(self, params : dict = {}): """ Compute CMB angular power spectra for given parameters. diff --git a/example_notebooks/ABCMB_Fluids.ipynb b/example_notebooks/ABCMB_Fluids.ipynb index 0022772..8d74469 100644 --- a/example_notebooks/ABCMB_Fluids.ipynb +++ b/example_notebooks/ABCMB_Fluids.ipynb @@ -60,7 +60,9 @@ { "cell_type": "markdown", "id": "a1ebdfd1-fbb5-4b6f-a1ac-32da38d92f2e", - "metadata": {}, + "metadata": { + "jp-MarkdownHeadingCollapsed": true + }, "source": [ "# A $\\Lambda {\\rm CDM}$ Run" ] @@ -99,7 +101,7 @@ ], "source": [ "LCDM = Model()\n", - "output_LCDM, aux_LCDM = LCDM.run_cosmology()" + "output_LCDM, aux_LCDM = LCDM()" ] }, { @@ -208,7 +210,7 @@ " \"T_nu_massive\" : 0.71611, # Massive neutrino effective temperature.\n", " \"m_nu_massive\" : 0.06, # Neutrino mass, in eV\n", "}\n", - "output_mnu, aux_mnu = LCDM_mnu.run_cosmology(params_LCDM_mnu)\n", + "output_mnu, aux_mnu = LCDM_mnu(params_LCDM_mnu)\n", "Pk_mnu = output_mnu[3]\n", "tt_mnu = output_mnu[0] * l * (l+1) / 2 / jnp.pi\n", "\n", @@ -637,7 +639,7 @@ " \"waDE\" : -0.6,\n", "}\n", "\n", - "output_w0wa, aux_w0wa = w0wa.run_cosmology(params)\n", + "output_w0wa, aux_w0wa = w0wa(params)\n", "\n", "# the model recompiles here as well--because we've added new fluids\n", "# to it again" @@ -979,7 +981,7 @@ } ], "source": [ - "output_SIDR, aux_SIDR = SIDRmodel.run_cosmology(params)" + "output_SIDR, aux_SIDR = SIDRmodel(params)" ] }, { @@ -1294,7 +1296,7 @@ } ], "source": [ - "output_NADM, aux_NADM = NADMmodel.run_cosmology(params)\n", + "output_NADM, aux_NADM = NADMmodel(params)\n", "Pk_NADM = output_NADM[3]\n", "tt_NADM = output_NADM[0] * l * (l+1) / 2 / jnp.pi" ] diff --git a/example_notebooks/ABCMB_basics.ipynb b/example_notebooks/ABCMB_basics.ipynb index 1e50f58..7a0984b 100644 --- a/example_notebooks/ABCMB_basics.ipynb +++ b/example_notebooks/ABCMB_basics.ipynb @@ -41,7 +41,7 @@ "source": [ "ABCMB is object-oriented, with all calculations running through an instance of `abcmb.main.Model`. The default ABCMB model is $\\Lambda$CDM cosmology in flat space and without neutrino masses. The five fluids present are cold dark matter, baryons, photons, massless neutrinos, and the cosmological constant. \n", "\n", - "To start, we initialize an instance of `Model` with all default options. Once initialized, our instance of `Model` can be used for calculations, either with `run_cosmology` or just by calling the initialized model." + "To start, we initialize an instance of `Model` with all default options. Once initialized, our instance of `Model` can be used for calculations by calling the initialized model." ] }, { @@ -72,8 +72,7 @@ ], "source": [ "model = Model()\n", - "output, aux = model.run_cosmology()\n", - "# equivalent: output, aux = model()" + "output, aux = model()" ] }, { @@ -81,7 +80,7 @@ "id": "c817c7e2-25ee-4b59-9668-e86e5fcaa407", "metadata": {}, "source": [ - "`Model.run_cosmology` returns two quantities: `output` and `aux`. Of these, the CMB power spectra, as well as the matter power spectrum if it was requested, are stored in the tuple `output`. `aux` on the other hand contains the $\\ell$ axis matching the CMB spectra or the $k$ axis for the $P(k)$. If you request either background or perturbations to be returned (illustrated below), they are included in `aux` as well.\n", + "`model()` returns two quantities: `output` and `aux`. Of these, the CMB power spectra, as well as the matter power spectrum if it was requested, are stored in the tuple `output`. `aux` on the other hand contains the $\\ell$ axis matching the CMB spectra or the $k$ axis for the $P(k)$. If you request either background or perturbations to be returned (illustrated below), they are included in `aux` as well.\n", "\n", "We can plot the results. Note by default ABCMB outputs the true $C_{\\ell}$'s and not the conventionally defined $D_{\\ell} \\equiv \\frac{\\ell(\\ell+1)}{2\\pi}C_{\\ell}$. We should make sure to include this before plotting:" ] @@ -182,7 +181,7 @@ "id": "615fc07e-dc1e-4e00-bd0e-7c6f6f679708", "metadata": {}, "source": [ - "`Model.run_cosmology` sets up background quantities (like $x_e$ and $H$), uses those background quantities to compute the evolution of perturbations, and then finally integrates the resulting transfer function to get $C_{\\ell}$s. " + "`model()` sets up background quantities (like $x_e$ and $H$), uses those background quantities to compute the evolution of perturbations, and then finally integrates the resulting transfer function to get $C_{\\ell}$s. " ] }, { @@ -210,7 +209,7 @@ } ], "source": [ - "output, aux = model.run_cosmology()" + "output, aux = model()" ] }, { @@ -324,7 +323,7 @@ " omega_m = 0.3 * 0.7**2 \n", " omega_cdm = omega_m - omega_b\n", " params = {\"omega_b\" : omega_b, \"omega_cdm\" : omega_cdm}\n", - " return model2.run_cosmology(params)" + " return model2(params)" ] }, { @@ -426,7 +425,7 @@ } ], "source": [ - "output, aux = myModel.run_cosmology({})\n", + "output, aux = myModel()\n", "full_params = aux[0] # Full list of parameters used\n", "PT = aux[1] # A pertubations table object\n", "BG = aux[2] # A background object" diff --git a/example_notebooks/ABCMB_with_LINX.ipynb b/example_notebooks/ABCMB_with_LINX.ipynb index 6082243..98ad806 100644 --- a/example_notebooks/ABCMB_with_LINX.ipynb +++ b/example_notebooks/ABCMB_with_LINX.ipynb @@ -131,7 +131,7 @@ } ], "source": [ - "ABC_Cls_LINX = model_LINX.run_cosmology(params_wrong)\n", + "ABC_Cls_LINX = model_LINX(params_wrong)\n", "ABC_tt_LINX = ABC_Cls_LINX[0] \n", "ABC_te_LINX = ABC_Cls_LINX[1] \n", "ABC_ee_LINX = ABC_Cls_LINX[2] \n", @@ -308,7 +308,7 @@ "source": [ "# first evaluation will take a long time because\n", "# the rest of ABCMB still needs to compile.\n", - "output, aux = model_LINX.run_cosmology(params_nuisance) # run_cosmology will automatically add_derived_parameters\n", + "output, aux = model_LINX(params_nuisance) # run_cosmology will automatically add_derived_parameters\n", "\n", "ells = aux[0]\n", "ABC_tt = output[0] * ells * (ells+1) / 2 / jnp.pi\n", diff --git a/pytests/accuracy_test.py b/pytests/accuracy_test.py index da491d4..a71a518 100644 --- a/pytests/accuracy_test.py +++ b/pytests/accuracy_test.py @@ -110,7 +110,7 @@ def test_accuracy_checker(h = 0.6762): # ABCMB - data, label = model.run_cosmology(params) + data, label = model(params) ells = label[0] ABC_tt = data[0] diff --git a/time_tests.py b/time_tests.py index 2f273d6..d33f147 100644 --- a/time_tests.py +++ b/time_tests.py @@ -45,7 +45,7 @@ # } params = {} - out, aux = model.run_cosmology(params) + out, aux = model(params) print(out[0]) print(time.time()-start)