# Beam-Infer

**Beam-Infer** is a Python library for beam parameter inference from raster-scanned images. It provides three high-level wrappers for parameter estimation using optimization and MCMC methods, with full control over configuration while maintaining a simple, stable API.

> **Note:** This package is based on the original work by **Joel Henriksson** (Lund University, LTH) from his master's thesis *"Inverse Problems in Proton Beam Imaging at ESS: Analysis and Numerical Methods"*. The core calculation base and algorithms were implemented by Joel Henriksson. This package integrates that work into ESS systems.
>
> **Original Repository:** [ess-beam-imaging-inverse-problems](https://github.com/juhelh/ess-beam-imaging-inverse-problems)

**Numerical equivalence** to the legacy base implementation (forward model, loss, optimizer pipeline, Poisson MCMC) is tested and documented; see [docs/EQUIVALENCE.md](docs/EQUIVALENCE.md). This package focuses on the calculation layer only; EPICS integration and I/O are handled elsewhere.

## Features

- **Complete Pipeline**: `estimate_with_optimizer` handles preprocessing, grid construction, initialization, and optimization
- **Multistart Optimization**: `run_multistart_minimization` for robust parameter recovery
- **MCMC Sampling**: `run_mcmc_inference` for Bayesian parameter estimation with uncertainty quantification
- **Flexible Configuration**: Rich configuration via dataclasses for all inference methods
- **JAX-based**: Fast, GPU-accelerated computations using JAX
- **Reproducible**: Optional random seeds for deterministic multistart runs

## Installation

### From Source

```bash
git clone https://gitlab.esss.lu.se/hugovalim/beam-infer.git
cd beam-infer
pip install -e .
```

### Development Installation

```bash
pip install -e ".[dev]"
```

## Quick Start

### Full Pipeline (Recommended)

The `estimate_with_optimizer` function provides a complete pipeline from raw camera image to estimated parameters:

```python
import numpy as np
import beam_infer

# Raw camera image (2D numpy array)
raw_image = np.array(...)  # Your image data

# Metadata (known quantities)
meta = beam_infer.Metadata(
    pulse_duration_ms=0.05,  # Pulse duration in milliseconds
    fx=39.55,                 # Horizontal scan frequency (kHz) - optional, fixes if provided
    fy=29.05,                 # Vertical scan frequency (kHz) - optional, fixes if provided
)

# Configuration
config = beam_infer.OptimizerConfig(
    max_iterations=100,       # Maximum L-BFGS-B iterations per run
    multistart={"n_starts": 5, "type": "uniform"},  # Multistart configuration
    regularization={"params": {"sigx": 0.1, "sigy": 0.1}},  # Optional regularization
    random_seed=42,           # Optional: for reproducible multistart runs
    verbose=False,            # Print iteration details
)

# Run estimation
result = beam_infer.estimate_with_optimizer(raw_image, meta=meta, config=config)

# Access results
print(f"Estimated parameters: {result.estimated_parameters}")
print(f"Final loss: {result.objective_value:.6e}")
print(f"Success: {result.success}")
print(f"Crop bounds: {result.crop_bounds}")
```

### Multistart Minimization

For more control over the optimization process:

```python
import numpy as np
import jax.numpy as jnp
import beam_infer

# Preprocessed image and grids (you handle preprocessing)
observed_image = np.array(...)  # Already preprocessed
X, Y = ...  # 2D coordinate grids matching image shape
t_vals = ...  # 1D time array

# Initial guess (10 parameters)
k0 = jnp.array([60.0, 20.0, 13.5, 5.05, 0.0, 0.0, 39.55, 29.05, 0.0, 0.0])

# Bounds (optional)
lower = jnp.array([0.0, 0.0, 2.0, 2.0, -20.0, -20.0, 20.0, 20.0, -np.pi, -np.pi])
upper = jnp.array([100.0, 50.0, 20.0, 20.0, 20.0, 20.0, 50.0, 50.0, np.pi, np.pi])

# Configuration
config = beam_infer.MinimizationConfig(
    initial_guess=k0,
    bounds=(lower, upper),
    max_iterations=100,
    verbose=False,
)

# Run multistart optimization
result = beam_infer.run_multistart_minimization(
    observed_image,
    X, Y, t_vals,
    config,
    num_starts=10,
    random_seed=42,  # For reproducibility
)

# Access best result
best = result.best_result
print(f"Best parameters: {best.estimated_parameters}")
print(f"Best loss: {best.objective_value:.6e}")

# Access all results (sorted by loss)
for i, r in enumerate(result.results):
    print(f"Run {i}: loss = {r.objective_value:.6e}")
```

### MCMC Inference

For Bayesian parameter estimation with uncertainty quantification:

```python
import numpy as np
import beam_infer

# Observed image (counts or normalized intensity)
observed_image = np.array(...)  # 2D image
X, Y = ...  # 2D coordinate grids
t_vals = ...  # 1D time array

# Bounds
lower = np.array([0.0, 0.0, 2.0, 2.0, -20.0, -20.0, 20.0, 20.0, -np.pi, -np.pi])
upper = np.array([100.0, 50.0, 20.0, 20.0, 20.0, 20.0, 50.0, 50.0, np.pi, np.pi])

# MCMC configuration
config = beam_infer.MCMCConfig(
    lower_bounds=lower,
    upper_bounds=upper,
    num_samples=4000,      # Number of samples to collect
    burn_in=1000,          # Burn-in period
    thin=2,                # Thinning (keep every Nth sample)
    prior_type="uniform",  # or "gaussian"
    seed=42,               # Random seed
)

# Run MCMC
result = beam_infer.run_mcmc_inference(observed_image, X, Y, t_vals, config)

# Access results
print(f"Posterior mean: {result.posterior_mean}")
print(f"Posterior std: {result.posterior_std}")
print(f"Acceptance rate: {result.acceptance_rate:.2%}")
print(f"Samples shape: {result.samples.shape}")  # (num_samples, 10)
```

## Supported Public API

The library exposes three main entrypoints plus parallel tempering helpers:

1. **`estimate_with_optimizer`** — Complete pipeline (preprocessing + optimization)
2. **`run_multistart_minimization`** — Multistart L-BFGS-B optimization
3. **`run_mcmc_inference`** — Metropolis-Hastings MCMC sampling
4. **Parallel tempering** — `build_beta_ladder`, `sample_from_prior`, `run_parallel_tempering` for replica-exchange MCMC (see [API Reference](docs/API.md))

Each entrypoint has corresponding configuration and result dataclasses. See the [API Reference](docs/API.md) for complete documentation.

## Parameters

The beam model uses 10 parameters (in order):

| Index | Name | Description | Units |
|-------|------|-------------|-------|
| 0 | `Ax` | Horizontal raster amplitude | mm |
| 1 | `Ay` | Vertical raster amplitude | mm |
| 2 | `sigx` | Horizontal beam width (std dev) | mm |
| 3 | `sigy` | Vertical beam width (std dev) | mm |
| 4 | `cx` | Horizontal center offset | mm |
| 5 | `cy` | Vertical center offset | mm |
| 6 | `fx` | Horizontal scan frequency | kHz |
| 7 | `fy` | Vertical scan frequency | kHz |
| 8 | `phix` | Horizontal phase offset | radians |
| 9 | `phiy` | Vertical phase offset | radians |

## Configuration Guide

### OptimizerConfig

Controls the full pipeline (`estimate_with_optimizer`):

```python
config = beam_infer.OptimizerConfig(
    max_iterations=100,              # Max L-BFGS-B iterations per run
    multistart={"n_starts": 5, "type": "uniform"},  # Multistart config
    regularization={"params": {"sigx": 0.1}},  # Regularization strengths
    random_seed=42,                  # Optional: for reproducible multistart
    verbose=False,                    # Print iteration details
    debug=False,                      # Return intermediate results
)
```

**Multistart options:**
- `"type": "uniform"` - Uniform random sampling within bounds
- `"type": "gaussian"` - Gaussian sampling around initial guess

**Regularization:**
- Dictionary mapping parameter names to penalty strengths
- Penalizes deviation from initial guess: `lambda * (param - param0)^2`

### MinimizationConfig

Controls single/multistart optimization:

```python
config = beam_infer.MinimizationConfig(
    initial_guess=k0,                # Initial parameter vector (10-D)
    bounds=(lower, upper),           # Optional: parameter bounds
    max_iterations=100,              # Max L-BFGS-B iterations
    verbose=False,                   # Print iteration details
    regularization={"sigx": 0.1},    # Optional: regularization strengths
)
```

### MCMCConfig

Controls MCMC sampling:

```python
config = beam_infer.MCMCConfig(
    lower_bounds=lower,              # Lower bounds (10-D array)
    upper_bounds=upper,              # Upper bounds (10-D array)
    num_samples=4000,                # Number of samples to collect
    burn_in=1000,                    # Burn-in period
    thin=2,                          # Thinning interval
    prior_type="uniform",            # "uniform" or "gaussian"
    prior_mean=None,                 # Required for Gaussian prior
    prior_std=None,                  # Required for Gaussian prior
    proposal_std=None,               # Optional: custom proposal std dev
    seed=42,                         # Random seed
    use_scan=True,                   # Use lax.scan (faster, default)
)
```

## When to Use Which Method?

- **`estimate_with_optimizer`**: Use when you have a raw camera image and want a complete pipeline. Handles preprocessing, grid construction, and optimization automatically.

- **`run_multistart_minimization`**: Use when you need more control over preprocessing/grids, or want to run optimization from multiple starting points for robustness.

- **`run_mcmc_inference`**: Use when you need uncertainty quantification, want to explore the posterior distribution, or need Bayesian inference.

## Examples

See the `examples/` directory for complete examples:

- `simple_estimation.py` ⭐ - **Simplest example using only public API** (recommended for beginners)
- `basic_simulation.py` - Basic beam image generation
- `parameter_variations.py` - Effects of different parameters
- `raster_pattern.py` - Visualize raster scanning patterns
- `inference_example.py` - Parameter recovery from synthetic data
- `real_image_estimation.py` - **Real snapshot: optimizer + MCMC with comparison**
- `comprehensive_example.py` - Complete workflow demonstration

Run examples:

```bash
# Simplest example (public API only)
python -m examples.simple_estimation

# Real image: optimizer vs MCMC (requires beam snapshot NPZ in examples/)
python -m examples.real_image_estimation

# Comprehensive example (recommended)
python -m examples.comprehensive_example
```

See [examples/README.md](examples/README.md) for detailed documentation on all examples.

## Requirements

- Python >= 3.10
- numpy >= 1.20.0
- jax >= 0.4.0
- jaxlib >= 0.4.0
- scipy >= 1.7.0

### GPU Support (Optional)

For GPU acceleration, install JAX with CUDA support:

```bash
# For CUDA 12.x
pip install --upgrade "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_release

# Verify GPU detection
python -c "import jax; print(jax.devices())"
```

## Development

### Running Tests

```bash
# After clone, fetch the reference implementation submodule (for equivalence tests)
git submodule update --init --recursive

pytest tests/
```

### Code Formatting

```bash
black beam_infer/ tests/ examples/
```

### Type Checking

```bash
mypy beam_infer/
```

## License

MIT License - see LICENSE file for details.
