jaxace 0.1.0__tar.gz

jaxace-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 CosmologicalEmulators
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jaxace-0.1.0/PKG-INFO

@@ -0,0 +1,66 @@
+Metadata-Version: 2.3
+Name: jaxace
+Version: 0.1.0
+Summary: JAX/Flax implementation of AbstractCosmologicalEmulators.jl interface
+Author: marcobonici
+Author-email: bonici.marco@gmail.com
+Requires-Python: >=3.10,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: diffrax (>=0.7.0,<0.8.0)
+Requires-Dist: flax (>=0.10.0,<0.11.0)
+Requires-Dist: interpax (>=0.3.9,<0.4.0)
+Requires-Dist: jax (>=0.4.30,<0.5.0)
+Requires-Dist: jaxlib (>=0.4.30,<0.5.0)
+Requires-Dist: numpy (>=2.2.6,<3.0.0)
+Requires-Dist: quadax (>=0.2.9,<0.3.0)
+Description-Content-Type: text/markdown
+
+# jaxace
+
+[![Tests](https://github.com/CosmologicalEmulators/jaxace/actions/workflows/tests.yml/badge.svg)](https://github.com/CosmologicalEmulators/jaxace/actions/workflows/tests.yml)
+[![Documentation](https://img.shields.io/badge/docs-stable-blue)](https://cosmologicalemulators.github.io/jaxace/stable/)
+[![Documentation](https://img.shields.io/badge/docs-dev-blue)](https://cosmologicalemulators.github.io/jaxace/dev/)
+[![codecov](https://codecov.io/gh/CosmologicalEmulators/jaxace/graph/badge.svg?token=8DGPCJR8KX)](https://codecov.io/gh/CosmologicalEmulators/jaxace)
+
+JAX/Flax implementation of cosmological emulators with automatic JIT compilation.
+
+## Installation
+
+```bash
+pip install -e .
+```
+
+## Usage
+
+```python
+import jaxace
+import jax.numpy as jnp
+
+# Cosmology
+cosmo = jaxace.W0WaCDMCosmology(
+    ln10As=3.044, ns=0.9649, h=0.6736,
+    omega_b=0.02237, omega_c=0.1200,
+    m_nu=0.06, w0=-1.0, wa=0.0
+)
+
+# Background functions
+z = jnp.array([0.0, 0.5, 1.0])
+growth = jaxace.D_z_from_cosmo(z, cosmo)
+distance = jaxace.r_z_from_cosmo(z, cosmo)
+
+# Neural network emulator
+emulator = jaxace.init_emulator(nn_dict, weights, jaxace.FlaxEmulator)
+output = emulator(input_data)  # Auto-JIT + batch detection
+```
+
+## Features
+
+- Background cosmology (growth, distances, Hubble)
+- Neural network emulators with auto-JIT
+- Massive neutrinos and dark energy support
+- Full JAX integration (grad, vmap, jit)
+

jaxace-0.1.0/README.md

@@ -0,0 +1,44 @@
+# jaxace
+
+[![Tests](https://github.com/CosmologicalEmulators/jaxace/actions/workflows/tests.yml/badge.svg)](https://github.com/CosmologicalEmulators/jaxace/actions/workflows/tests.yml)
+[![Documentation](https://img.shields.io/badge/docs-stable-blue)](https://cosmologicalemulators.github.io/jaxace/stable/)
+[![Documentation](https://img.shields.io/badge/docs-dev-blue)](https://cosmologicalemulators.github.io/jaxace/dev/)
+[![codecov](https://codecov.io/gh/CosmologicalEmulators/jaxace/graph/badge.svg?token=8DGPCJR8KX)](https://codecov.io/gh/CosmologicalEmulators/jaxace)
+
+JAX/Flax implementation of cosmological emulators with automatic JIT compilation.
+
+## Installation
+
+```bash
+pip install -e .
+```
+
+## Usage
+
+```python
+import jaxace
+import jax.numpy as jnp
+
+# Cosmology
+cosmo = jaxace.W0WaCDMCosmology(
+    ln10As=3.044, ns=0.9649, h=0.6736,
+    omega_b=0.02237, omega_c=0.1200,
+    m_nu=0.06, w0=-1.0, wa=0.0
+)
+
+# Background functions
+z = jnp.array([0.0, 0.5, 1.0])
+growth = jaxace.D_z_from_cosmo(z, cosmo)
+distance = jaxace.r_z_from_cosmo(z, cosmo)
+
+# Neural network emulator
+emulator = jaxace.init_emulator(nn_dict, weights, jaxace.FlaxEmulator)
+output = emulator(input_data)  # Auto-JIT + batch detection
+```
+
+## Features
+
+- Background cosmology (growth, distances, Hubble)
+- Neural network emulators with auto-JIT
+- Massive neutrinos and dark energy support
+- Full JAX integration (grad, vmap, jit)

jaxace-0.1.0/jaxace/__init__.py

@@ -0,0 +1,72 @@
+"""
+JAX AbstractCosmologicalEmulators (jaxace)
+
+A JAX/Flax implementation of the AbstractCosmologicalEmulators.jl interface,
+providing foundational neural network emulator infrastructure for cosmological
+computations.
+"""
+
+from .core import (
+    AbstractTrainedEmulator,
+    FlaxEmulator,
+    run_emulator,
+    get_emulator_description
+)
+
+from .initialization import (
+    init_emulator,
+    MLP
+)
+
+from .utils import (
+    maximin,
+    inv_maximin,
+    validate_nn_dict_structure,
+    validate_parameter_ranges,
+    validate_layer_structure,
+    safe_dict_access
+)
+
+# Import background cosmology functions
+from .background import (
+    W0WaCDMCosmology,
+    a_z, E_a, E_z, dlogEdloga, Ωma,
+    D_z, f_z, D_f_z,
+    D_z_from_cosmo, f_z_from_cosmo, D_f_z_from_cosmo,
+    r_z, r_z_from_cosmo,
+    dA_z, dA_z_from_cosmo,
+    dL_z, dL_z_from_cosmo,
+    ρc_z, Ωtot_z
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Core types and functions
+    "AbstractTrainedEmulator",
+    "FlaxEmulator", 
+    "run_emulator",
+    "get_emulator_description",
+    
+    # Initialization
+    "init_emulator",
+    "MLP",
+    
+    # Utilities
+    "maximin",
+    "inv_maximin",
+    "validate_nn_dict_structure",
+    "validate_parameter_ranges",
+    "validate_layer_structure",
+    "safe_dict_access",
+    
+    # Background cosmology
+    "W0WaCDMCosmology",
+    "a_z", "E_a", "E_z", "dlogEdloga", "Ωma",
+    "D_z", "f_z", "D_f_z",
+    "D_z_from_cosmo", "f_z_from_cosmo", "D_f_z_from_cosmo",
+    "r_z", "r_z_from_cosmo",
+    "dA_z", "dA_z_from_cosmo", 
+    "dL_z", "dL_z_from_cosmo",
+    "ρc_z", "Ωtot_z"
+]