Skip to content

Use kwargs instead of strict dict #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
TonyZhou729 merged 3 commits into from
Feb 12, 2026
Merged

Use kwargs instead of strict dict #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
00afd9c
updated to kwargs
cgiovanetti Feb 12, 2026
73490ce
partial update of example notebooks
cgiovanetti Feb 12, 2026
87617bb
small notebook update
cgiovanetti Feb 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion abcmb/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,5 +4,4 @@
A fully differentiable Boltzmann solver for the CMB.
"""

__version__ = "0.1.1.1"
__author__ = 'Zilu Zhou, Cara Giovanetti, and Hongwan Liu'
11 changes: 6 additions & 5 deletions abcmb/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,8 +81,8 @@ class Model(eqx.Module):
# In the init, all species that are present within the model should be set to True.
# All couplings present between species should be set to true.
def __init__(self,
input_specs = {},
user_species=None
user_species=None,
**kwargs
):
"""
Initialize Model instance.
Expand All @@ -92,14 +92,15 @@ def __init__(self,

Parameters:
-----------
input_specs : dict
A dictionary containing run options (expected to be static)
user_species : tuple
A tuple of user-defined fluids to be included in the cosmology
**kwargs : dict
Configuration options passed as keyword arguments.
Any unknown keys will be preserved for custom species extensibility.
"""

# Fill in all user defined and missing specs parameters
specs = model_specs.load_specs(input_specs)
specs = model_specs.load_specs(kwargs)
self.specs = specs

# Populate all species
Expand Down
5 changes: 5 additions & 0 deletions abcmb/model_specs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,11 @@ def load_specs(input_specs):
specs["scale_dop"] = input_specs.get("scale_dop", 1)
specs["scale_pol"] = input_specs.get("scale_pol", 1)

# Preserve any unknown keys for custom species extensibility
for key, value in input_specs.items():
if key not in specs:
specs[key] = value

return specs

def populate_species(user_species, specs):
Expand Down
88 changes: 26 additions & 62 deletions example_notebooks/ABCMB_Fluids.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,34 +40,7 @@
"execution_count": 2,
"id": "81ea175b-9122-40bd-a054-4ace224a81fc",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"ERROR:2026-02-11 14:49:01,512:jax._src.xla_bridge:477: Jax plugin configuration error: Exception when calling jax_plugins.xla_cuda13.initialize()\n",
"Traceback (most recent call last):\n",
" File \"/ext3/miniforge3/lib/python3.12/site-packages/jax/_src/xla_bridge.py\", line 475, in discover_pjrt_plugins\n",
" plugin_module.initialize()\n",
" File \"/ext3/miniforge3/lib/python3.12/site-packages/jax_plugins/xla_cuda13/__init__.py\", line 328, in initialize\n",
" _check_cuda_versions(raise_on_first_error=True)\n",
" File \"/ext3/miniforge3/lib/python3.12/site-packages/jax_plugins/xla_cuda13/__init__.py\", line 285, in _check_cuda_versions\n",
" local_device_count = cuda_versions.cuda_device_count()\n",
" ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
"RuntimeError: jaxlib/cuda/versions_helpers.cc:113: operation cuInit(0) failed: Unknown CUDA error 303; cuGetErrorName failed. This probably means that JAX was unable to load the CUDA libraries.\n",
"ERROR:2026-02-11 14:49:02,019:jax._src.xla_bridge:477: Jax plugin configuration error: Exception when calling jax_plugins.xla_cuda12.initialize()\n",
"Traceback (most recent call last):\n",
" File \"/ext3/miniforge3/lib/python3.12/site-packages/jax/_src/xla_bridge.py\", line 475, in discover_pjrt_plugins\n",
" plugin_module.initialize()\n",
" File \"/ext3/miniforge3/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py\", line 328, in initialize\n",
" _check_cuda_versions(raise_on_first_error=True)\n",
" File \"/ext3/miniforge3/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py\", line 285, in _check_cuda_versions\n",
" local_device_count = cuda_versions.cuda_device_count()\n",
" ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
"RuntimeError: jaxlib/cuda/versions_helpers.cc:113: operation cuInit(0) failed: Unknown CUDA error 303; cuGetErrorName failed. This probably means that JAX was unable to load the CUDA libraries.\n"
]
}
],
"outputs": [],
"source": [
"import sys\n",
"sys.path.append('..')\n",
Expand Down Expand Up @@ -102,7 +75,7 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": 3,
"id": "b94bb915-8efe-4658-b8e9-c2e5ebc2a733",
"metadata": {},
"outputs": [
Expand All @@ -126,7 +99,7 @@
],
"source": [
"LCDM = Model()\n",
"output_LCDM, aux_LCDM = LCDM.run_cosmology({})"
"output_LCDM, aux_LCDM = LCDM.run_cosmology()"
]
},
{
Expand Down Expand Up @@ -321,9 +294,7 @@
{
"cell_type": "markdown",
"id": "b1278b55",
"metadata": {
"jp-MarkdownHeadingCollapsed": true
},
"metadata": {},
"source": [
"# Defining new fluids: The Fluid Module"
]
Expand All @@ -341,7 +312,9 @@
{
"cell_type": "markdown",
"id": "06e48447-3777-4dde-8ea6-8a00c5a00160",
"metadata": {},
"metadata": {
"jp-MarkdownHeadingCollapsed": true
},
"source": [
"## The Module"
]
Expand Down Expand Up @@ -434,7 +407,9 @@
{
"cell_type": "markdown",
"id": "57bc5cd4-b80f-4a58-bedd-fddbf3565c0c",
"metadata": {},
"metadata": {
"jp-MarkdownHeadingCollapsed": true
},
"source": [
"## The Equinox Fields\n",
"\n",
Expand All @@ -454,7 +429,9 @@
{
"cell_type": "markdown",
"id": "325925db",
"metadata": {},
"metadata": {
"jp-MarkdownHeadingCollapsed": true
},
"source": [
"## The Methods\n",
"\n",
Expand Down Expand Up @@ -499,7 +476,7 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": 4,
"id": "edae6b41-71cf-4943-a66c-9f7518a40957",
"metadata": {},
"outputs": [],
Expand Down Expand Up @@ -543,7 +520,7 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 5,
"id": "15443a54-3b61-4d5c-b750-67c902d86135",
"metadata": {},
"outputs": [],
Expand Down Expand Up @@ -594,7 +571,7 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 6,
"id": "1504c7f0-67e8-4b16-bcf5-162bea81d77b",
"metadata": {},
"outputs": [
Expand All @@ -603,15 +580,11 @@
"output_type": "stream",
"text": [
"w0wa:\n",
"{'DynamicalDarkEnergy': 0, 'ColdDarkMatter': 1, 'Baryon': 2, 'Photon': 3, 'MasslessNeutrino': 4}\n"
"{'ColdDarkMatter': 0, 'Photon': 1, 'DynamicalDarkEnergy': 2, 'Baryon': 3, 'MasslessNeutrino': 4}\n"
]
}
],
"source": [
"specs = {\n",
" \"use_LCDM_species\" : False # First, ask ABCMB not to use default species\n",
"}\n",
"\n",
"# Since we're not using the default species set, we need to add all the other fluids too.\n",
"user_species = {\n",
" species.Baryon,\n",
Expand All @@ -621,7 +594,7 @@
" DynamicalDarkEnergy\n",
"}\n",
"\n",
"w0wa = Model(specs, user_species=user_species)\n",
"w0wa = Model(user_species=user_species,use_LCDM_species=False) # ask ABCMB not to use default species\n",
"print(\"w0wa:\")\n",
"print(w0wa.species_dict)"
]
Expand Down Expand Up @@ -954,11 +927,7 @@
" SIDR,\n",
")\n",
"\n",
"specs = {\n",
" \"max_steps_PE\" : 4096,\n",
"}\n",
"\n",
"SIDRmodel = Model(specs, user_species)\n",
"SIDRmodel = Model(user_species,max_steps_PE=4096)\n",
"print(\"SIDR model:\")\n",
"print(SIDRmodel.species_dict)"
]
Expand All @@ -968,7 +937,7 @@
"id": "abb08ccd-49e6-4d0d-be9c-62ef35d57ad1",
"metadata": {},
"source": [
"Here we used an extra `specs` parameters `max_steps_PE`. This parameter tells the adaptive step size solver what is the maximum number of steps it should expect to take, in order to solve the linear perturbations equations. This number is 2048 by default, but adding new physics scenarios such as the SIDR can make these equations more difficult. With this method we allow our solver to take more steps to solve, though this may increase runtime. "
"Here we used an extra keyword argument `max_steps_PE`. This parameter tells the adaptive step size solver what is the maximum number of steps it should expect to take, in order to solve the linear perturbations equations. This number is 2048 by default, but adding new physics scenarios such as the SIDR can make these equations more difficult. With this method we allow our solver to take more steps to solve, though this may increase runtime. "
]
},
{
Expand Down Expand Up @@ -1238,10 +1207,7 @@
"outputs": [],
"source": [
"lmax = 2500\n",
"specs = {\n",
" \"use_LCDM_species\" : False,\n",
" \"max_steps_PE\" : 4096,\n",
"}\n",
"\n",
"user_species = (\n",
" species.DarkEnergy, \n",
" IDM, \n",
Expand All @@ -1251,7 +1217,7 @@
" species.MasslessNeutrino\n",
")\n",
"\n",
"NADMmodel = Model(specs, user_species)"
"NADMmodel = Model(user_species,use_LCDM_species=False, max_steps_PE = 4096)"
]
},
{
Expand Down Expand Up @@ -1418,9 +1384,7 @@
{
"cell_type": "markdown",
"id": "37f91f5c-c8f4-4def-a795-8edc7cf3a8f5",
"metadata": {
"jp-MarkdownHeadingCollapsed": true
},
"metadata": {},
"source": [
"# Cheat Sheet: Define your own species!"
]
Expand Down Expand Up @@ -1463,9 +1427,9 @@
],
"metadata": {
"kernelspec": {
"display_name": "default",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "my_env"
"name": "python3"
},
"language_info": {
"codemirror_mode": {
Expand All @@ -1477,7 +1441,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.12"
"version": "3.12.11"
}
},
"nbformat": 4,
Expand Down
55 changes: 34 additions & 21 deletions example_notebooks/ABCMB_basics.ipynb
View file
Open in desktop

Large diffs are not rendered by default.

42 changes: 29 additions & 13 deletions example_notebooks/ABCMB_with_LINX.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": 1,
"id": "80e9a0dc-c54f-4a0d-888c-9b0191b4b3f2",
"metadata": {},
"outputs": [],
Expand Down Expand Up @@ -55,7 +55,7 @@
},
{
"cell_type": "code",
"execution_count": 9,
"execution_count": 2,
"id": "4a4dc262-8332-4b0b-bbda-e45786b0dc97",
"metadata": {},
"outputs": [],
Expand All @@ -70,7 +70,7 @@
" \"linx_reaction_net\" : 'key_PArthENoPE' # Can select between 'key_PArthENoPE' and \"key_PRIMAT_2023\" (default) networks\n",
"}\n",
"\n",
"model_LINX = Model(specs_LINX)"
"model_LINX = Model(**specs_LINX)"
]
},
{
Expand All @@ -83,7 +83,7 @@
},
{
"cell_type": "code",
"execution_count": 10,
"execution_count": 3,
"id": "8f054803-7bc7-45a9-99a4-18919d47166a",
"metadata": {},
"outputs": [],
Expand All @@ -98,17 +98,18 @@
},
{
"cell_type": "code",
"execution_count": 11,
"execution_count": 4,
"id": "73a01be0-d3bd-4bc4-a3b6-cc21479d727c",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"You have specified a value for N_nu_massless and/or Neff, but LINX instead expects a \n",
"parameter 'Delta_Neff_init' which will be used to compute Neff. Refer to LINX \n",
"docs or https://arxiv.org/abs/2408.14538 for more information.\n"
"You have specified a value for N_nu_massless and/or Neff and/or T_nu_massless, \n",
"but LINX instead expects a parameter 'Delta_Neff_init' which will be used to \n",
"compute Neff. Refer to LINX docs or https://arxiv.org/abs/2408.14538 for more info.\n",
"\n"
]
},
{
Expand All @@ -117,14 +118,14 @@
"output_type": "error",
"traceback": [
"An exception has occurred, use %tb to see the full traceback.\n",
"\u001b[0;31mSystemExit\u001b[0m\n"
"\u001b[31mSystemExit\u001b[39m\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/ext3/miniconda3/lib/python3.10/site-packages/IPython/core/interactiveshell.py:3513: UserWarning: To exit: use 'exit', 'quit', or Ctrl-D.\n",
"/opt/anaconda3/envs/linx312/lib/python3.12/site-packages/IPython/core/interactiveshell.py:3680: UserWarning: To exit: use 'exit', 'quit', or Ctrl-D.\n",
" warn(\"To exit: use 'exit', 'quit', or Ctrl-D.\", stacklevel=1)\n"
]
}
Expand Down Expand Up @@ -157,7 +158,7 @@
},
{
"cell_type": "code",
"execution_count": 14,
"execution_count": 5,
"id": "b7021cd7-2e27-4697-9049-0311e2b9c217",
"metadata": {},
"outputs": [],
Expand All @@ -180,17 +181,32 @@
},
{
"cell_type": "code",
"execution_count": 15,
"execution_count": 6,
"id": "a9289396-b860-4a48-84ea-51c455cb7e28",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"`\\ /´ |||| |||| ||||| |||| |||| ||||\n",
" /\\_______/\\ |||| |||| ||||||| |||| |||| ||||\n",
" ) __` ´__ ( |||| |||| |||| |||| |||| |||||||\n",
"/ `-|_|-´ \\ |||| |||| |||| |||| ||| ||||||| \n",
"/ (_x_) \\ |||||||||| |||| |||| ||||||| |||| ||||\n",
" ) `-´ ( |||||||||| |||| |||| |||||| |||| ||||\n",
" \n",
"Compiling thermodynamics model...\n",
"Compiling abundance model...\n"
]
},
{
"data": {
"text/plain": [
"Array(3.12600995, dtype=float64)"
]
},
"execution_count": 15,
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
Expand Down
Loading