AxModelFitter_fit_model
Fit custom JAX mathematical models to experimental data using numerical optimization. Requires data file, model function, and parameter bounds.
Instructions
Fit a custom JAX mathematical model against experimental data.
This tool fits user-defined mathematical models to data using numerical optimization.
All data MUST be provided via files (CSV, Excel, JSON, Parquet) - no direct data input.
REQUIRED INPUTS:
1. data_file: Path to your data file (e.g., "/path/to/data.csv")
2. input_data: Maps file columns to input variables
3. output_data: Maps file columns to output variables
4. function_source: JAX function code using jnp operations
5. parameters: Initial parameter guesses with units
6. bounds: Bounds for ALL parameters, inputs, and outputs
DATA MAPPING EXAMPLE:
- data_file: "/Users/data/experiment.csv"
- input_data: [{"column": "time_col", "name": "t", "unit": "second"}]
- output_data: {"columns": ["voltage"], "name": "v", "unit": "volt"}
FUNCTION REQUIREMENTS:
- MUST use JAX operations: jnp.exp(-rate*t), jnp.sin(freq*t), jnp.sqrt(x)
- Valid pint units: 'dimensionless', 'second', 'volt', 'meter', etc.
- All variables (parameters, inputs, outputs) need bounds
RETURNS: Optimized parameters, R², execution time, and result files
Input Schema
| Name | Required | Description | Default |
|---|---|---|---|
| bounds | Yes | ALL parameter/input/output bounds: [{'name': 'a', 'lower': {'magnitude': 0, 'unit': 'dimensionless'}, 'upper': {'magnitude': 10, 'unit': 'dimensionless'}}] | |
| max_time | No | Maximum optimization time in seconds | |
| constants | No | Fixed constants: [{'name': 'c', 'value': {'magnitude': 3.0, 'unit': 'meter'}}] | |
| data_file | Yes | Path to data file (CSV, Excel, JSON, Parquet). All data must be provided via file. | |
| docstring | No | Brief description of the model | |
| input_data | Yes | Input column mappings: [{'column': 'time', 'name': 't', 'unit': 'second'}, {'column': 'x_col', 'name': 'x', 'unit': 'meter'}] | |
| model_name | Yes | Model name (e.g., 'ExponentialDecay', 'RingResonator') | |
| parameters | Yes | Initial parameter guesses: [{'name': 'a', 'value': {'magnitude': 2.0, 'unit': 'dimensionless'}}] | |
| file_format | No | File format: 'csv', 'excel', 'json', 'parquet' (auto-detect if None) | |
| jit_compile | No | Enable JIT compilation for performance | |
| output_data | Yes | Output column mapping: {'columns': ['signal'], 'name': 'y', 'unit': 'volt'} OR {'columns': ['y1', 'y2'], 'name': 'y', 'unit': 'volt'} | |
| function_name | Yes | Function name that computes the model output | |
| optimizer_type | No | Optimizer: 'nlopt' (best default), 'scipy' (simple), 'nevergrad' (gradient-free) | nlopt |
| function_source | Yes | JAX function source code. MUST use jnp operations: jnp.exp, jnp.sin, etc. | |
| optimizer_config | No | Optimizer config: {'use_gradient': True, 'tol': 1e-6, 'max_function_eval': 1000000} | |
| cost_function_type | No | Cost function: 'mse' (default), 'mae', 'huber (with delta=1.0)', 'relative_mse' | mse |