from __future__ import annotations
from typing import Any
import numpy as np
[docs]def square(params, func, r, z):
"""
calculate ||z - func(x, y, *params)||^2
where x and y are determine by z.shape
"""
z_guess = func(r, *params)
return np.mean((z - z_guess)**2)
[docs]def masked_square(params, func, r, z, mask):
z_guess = func(r, *params)
return np.mean((z[mask] - z_guess[mask])**2)
[docs]def diagonal_gaussian(r, *params):
a, b = params[-2:]
ndim = len(params[:-2])//2
mu = params[:ndim]
sg = params[ndim:-2]
z_value = np.array([(x0 - mu0)/sg0 for x0, mu0, sg0 in zip(r, mu, sg)])
return a * np.exp(-np.sum(z_value**2, axis=0) / 2) + b
[docs]class Gaussian:
[docs] def mu_inrange(self, low, high):
return np.logical_and(low<=self.mu, self.mu<=high).all()
[docs] def sg_inrange(self, low, high):
sg_ = np.abs(self.sg)
return np.logical_and(low<=sg_, sg_<=high).all()
[docs]class DiagonalGaussian(Gaussian):
def __init__(self, params=None):
if params is None:
self.mu = self.sg = self.a = self.b = None
else:
self.mu, self.sg, self.a, self.b = params
@property
def mu(self):
return self._mu
@mu.setter
def mu(self, value):
if value is None:
self._mu = None
else:
self._mu = np.asarray(value)
@property
def sg(self):
return self._sg
@sg.setter
def sg(self, value):
if value is None:
self._sg = None
else:
self._sg = np.asarray(value)
@property
def params(self):
"""Get flattened parameters."""
return tuple(self.mu) + tuple(self.sg) + (self.a, self.b)
@params.setter
def params(self, params:tuple):
self.a, self.b = params[-2:]
self.mu = params[:self.ndim]
self.sg = params[self.ndim:-2]
[docs] def asdict(self) -> dict[str, Any]:
return {
"mu": self.mu.tolist(),
"sigma": self.sg.tolist(),
"A": self.a,
"B": self.b,
}
@property
def ndim(self):
return self.mu.size
[docs] def fit(self, data: np.ndarray, method: str = "Powell", mask: np.ndarray | None = None):
from scipy.optimize import minimize
if self.mu is None or self.sg is None or self.a is None or self.b is None:
if mask is not None:
data = data.copy()
data[mask] = -np.inf
self._estimate_params(data)
else:
self._estimate_params(data)
r = np.indices(data.shape)
if mask is None:
result = minimize(
square,
self.params,
args=(diagonal_gaussian, r, data),
method=method
)
else:
if mask.shape != data.shape:
raise ValueError(
f"Shape mismatch between data {data.shape!r} and mask {mask.shape!r}."
)
result = minimize(
masked_square,
self.params,
args=(diagonal_gaussian, r, data, ~mask),
method=method
)
self.params = result.x
return result
[docs] def rescale(self, scale: float):
self.mu *= scale
self.sg *= scale
return None
[docs] def shift(self, dxdy):
self.mu += np.array(dxdy)
return None
[docs] def generate(self, shape:tuple) -> np.ndarray:
r = np.indices(shape, dtype=np.float32)
return diagonal_gaussian(r, *self.params)
def _estimate_params(self, data:np.ndarray):
pass
[docs]class GaussianParticle(DiagonalGaussian):
def __init__(self, params=None, initial_sg=1):
super().__init__(params)
self.initial_sg = initial_sg
def _estimate_params(self, data:np.ndarray):
# n-dim argmax
self.mu = np.array(np.unravel_index(np.argmax(data), data.shape), dtype="float32")
self.sg = np.full(data.ndim, self.initial_sg, dtype="float32")
self.b, p95 = np.percentile(data, [5, 95])
self.a = p95 - self.b
return None
[docs]class GaussianBackground(DiagonalGaussian):
def _estimate_params(self, data: np.ndarray):
# n-dim argmax
self.mu = np.array(np.unravel_index(np.argmax(data), data.shape), dtype="float32")
self.sg = np.array(data.shape, dtype="float32")
self.b, p95 = np.percentile(data, [5, 95])
self.a = p95 - self.b
return None