Model.run_cosmology is no more

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.
 

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"
    ]

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,16 +72,15 @@
    ],
    "source": [
     "model = Model()\n",
-    "output, aux = model.run_cosmology()\n",
-    "# equivalent: output, aux = model()"
+    "output, aux = model()"
    ]
   },
   {
    "cell_type": "markdown",
    "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"

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",

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] 

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)