beam-infer 1.0.0__tar.gz

beam_infer-1.0.0/LICENSE

@@ -0,0 +1,24 @@
+MIT License
+
+Copyright (c) 2024 Joel Henriksson, Hugo Valim
+
+Original work by Joel Henriksson (Lund University, LTH)
+Package integration by Hugo Valim (ESS)
+
+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.

beam_infer-1.0.0/PKG-INFO

@@ -0,0 +1,364 @@
+Metadata-Version: 2.4
+Name: beam-infer
+Version: 1.0.0
+Summary: A Python package for beam simulation and parameter inference using raster scanning models
+Author-email: Joel Henriksson <carljoelhenriksson@gmail.com>, Hugo Valim <hugo.valim@ess.eu>
+License: MIT
+Project-URL: Original_Work, https://github.com/juhelh/ess-beam-imaging-inverse-problems
+Keywords: beam,simulation,raster,mcmc,inference,optimization,fitting
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.20.0
+Requires-Dist: jax>=0.4.0
+Requires-Dist: jaxlib>=0.4.0
+Requires-Dist: scipy>=1.7.0
+Provides-Extra: gpu
+Requires-Dist: jax[cuda12]>=0.4.0; extra == "gpu"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
+Requires-Dist: matplotlib>=3.5.0; extra == "dev"
+Requires-Dist: tqdm>=4.60.0; extra == "dev"
+Requires-Dist: black>=22.0.0; extra == "dev"
+Requires-Dist: flake8>=5.0.0; extra == "dev"
+Requires-Dist: mypy>=1.0.0; extra == "dev"
+Dynamic: license-file
+
+# 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.