AMS-BP 0.0.2__py3-none-any.whl

AMS_BP/optics/psf/psf_engine.py

@@ -0,0 +1,215 @@
+from dataclasses import dataclass
+from functools import cache, cached_property, lru_cache
+from typing import Literal, Optional, Tuple
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+@dataclass(frozen=True)
+class PSFParameters:
+    """Parameters for PSF (Point Spread Function) generation.
+
+    Attributes:
+        wavelength: Light wavelength in nanometers
+        numerical_aperture: Numerical aperture of the optical system
+        pixel_size: Size of pixels in micrometers
+        z_step: Axial step size in micrometers
+        refractive_index: Refractive index of the medium (default: 1.0 for air)
+    """
+
+    wavelength: float
+    numerical_aperture: float
+    pixel_size: float
+    z_step: float
+    refractive_index: float = 1.0
+
+    # def __post_init__(self) -> None:
+    #     """Validate parameters after initialization."""
+    #     if any(
+    #         param <= 0
+    #         for param in (
+    #             self.wavelength,
+    #             self.numerical_aperture,
+    #             self.pixel_size,
+    #             self.z_step,
+    #             self.refractive_index,
+    #         )
+    #     ):
+    #         raise ValueError("All parameters must be positive numbers")
+    #     if self.numerical_aperture >= self.refractive_index:
+    #         raise ValueError("Numerical aperture must be less than refractive index")
+
+    @cached_property
+    def wavelength_um(self) -> float:
+        """Wavelength in micrometers."""
+        return self.wavelength / 1000.0
+
+
+class PSFEngine:
+    """Engine for generating various microscope Point Spread Functions.
+
+    This class implements calculations for both 2D and 3D Point Spread Functions
+    using Gaussian approximations.
+    """
+
+    def __init__(self, params: PSFParameters):
+        """Initialize PSF engine with given parameters."""
+        self.params = params
+        self._initialize_calculations()
+
+    def _initialize_calculations(self) -> None:
+        """Initialize commonly used calculations."""
+        self._sigma_xy = _calculate_sigma_xy(
+            self.params.wavelength_um, self.params.numerical_aperture
+        )
+        self._sigma_z = _calculate_sigma_z(
+            self.params.wavelength_um,
+            self.params.numerical_aperture,
+            self.params.refractive_index,
+        )
+        self._psf_size = calculate_psf_size(
+            sigma_xy=self._sigma_xy,
+            pixel_size=self.params.pixel_size,
+            sigma_z=self._sigma_z,
+        )
+        self._grid_xy = _generate_grid(self._psf_size, self.params.pixel_size)
+
+        # Pre-calculate normalized sigma values
+        self._norm_sigma_xy = self._sigma_xy / 2.355
+        self._norm_sigma_z = self._sigma_z / 2.355
+
+    @lru_cache(maxsize=128)
+    def psf_z(self, z_val: float) -> NDArray[np.float64]:
+        """Generate z=z_val Gaussian approximation of PSF.
+
+        Args:
+            z_val: Z-position in micrometers
+
+        Returns:
+            2D array containing the PSF at given z position
+        """
+        x, y = self._grid_xy
+
+        # Vectorized calculation
+        r_squared = (x / self._norm_sigma_xy) ** 2 + (y / self._norm_sigma_xy) ** 2
+        z_term = (z_val / self._norm_sigma_z) ** 2
+        return np.exp(-0.5 * (r_squared + z_term))
+
+    @lru_cache(maxsize=128)
+    def psf_z_xy0(self, z_val: float) -> float:
+        """Generate z=z_val Gaussian approximation of PSF with x=y=0.
+
+        Args:
+            z_val: Z-position in micrometers
+
+        Returns:
+            PSF value at x=y=0 and given z position
+        """
+        return np.exp(-0.5 * (z_val / self._norm_sigma_z) ** 2)
+
+    @cache
+    def _3d_normalization_A(
+        self, sigma_z: float, sigma_x: float, sigma_y: float
+    ) -> float:
+        return 1.0 / (((2.0 * np.pi) ** (3.0 / 2.0)) * sigma_x * sigma_y * sigma_z)
+
+    @cache
+    def _2d_normalization_A(self, sigma_x: float, sigma_y: float) -> float:
+        return 1.0 / ((2.0 * np.pi) * sigma_x * sigma_y)
+
+    @staticmethod
+    def normalize_psf(
+        psf: NDArray[np.float64], mode: Literal["sum", "max", "energy"] = "sum"
+    ) -> NDArray[np.float64]:
+        """Normalize PSF with different schemes.
+
+        Args:
+            psf: Input PSF array
+            mode: Normalization mode
+                - 'sum': Normalize so sum equals 1 (energy conservation)
+                - 'max': Normalize so maximum equals 1
+                - 'energy': Normalize so squared sum equals 1
+
+        Returns:
+            Normalized PSF array
+
+        Raises:
+            ValueError: If unknown normalization mode is specified
+        """
+        if not np.any(psf):  # Check if array is all zeros
+            return psf
+
+        normalizers = {
+            "sum": np.sum,
+            "max": np.max,
+            "energy": lambda x: np.sqrt(np.sum(x**2)),
+        }
+
+        try:
+            normalizer = normalizers[mode]
+            return psf / normalizer(psf)
+        except KeyError:
+            raise ValueError(
+                f"Unknown normalization mode: {mode}. Valid modes: {list(normalizers.keys())}"
+            )
+
+
+@cache
+def _calculate_sigma_xy(wavelength_um: float, numerical_aperture: float) -> float:
+    """Calculate lateral sigma value."""
+    return 0.61 * wavelength_um / numerical_aperture
+
+
+@cache
+def _calculate_sigma_z(
+    wavelength_um: float, numerical_aperture: float, refractive_index: float
+) -> float:
+    """Calculate axial sigma value."""
+    return 2.0 * wavelength_um * refractive_index / (numerical_aperture**2)
+
+
+@cache
+def _generate_grid(
+    size: Tuple[int, int], pixel_size: float
+) -> Tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Generate coordinate grids for PSF calculation.
+
+    Args:
+        size: Tuple of (height, width) for the grid
+
+    Returns:
+        Tuple of x and y coordinate arrays
+    """
+    y, x = np.ogrid[: size[0], : size[1]]
+    center_y, center_x = [(s - 1) / 2 for s in size]
+    y = (y - center_y) * pixel_size
+    x = (x - center_x) * pixel_size
+    return x, y
+
+
+@cache
+def calculate_psf_size(
+    sigma_xy: float, pixel_size: float, sigma_z: float, z_size: Optional[int] = None
+) -> Tuple[int, ...]:
+    """Calculate appropriate PSF size based on physical parameters.
+
+    Args:
+        z_size: Optional number of z-planes for 3D PSF
+
+    Returns:
+        Tuple of dimensions (z,y,x) or (y,x) for the PSF calculation
+    """
+    # Calculate radius to capture important features (2x Airy radius)
+    r_psf = 2 * sigma_xy
+
+    # Convert to pixels and ensure odd number
+    pixels_xy = int(np.ceil(r_psf / pixel_size))
+    pixels_xy += (pixels_xy + 1) % 2
+
+    if z_size is not None:
+        pixels_z = int(np.ceil(2 * sigma_z / z_size))
+        pixels_z += (pixels_z + 1) % 2
+        return (pixels_z, pixels_xy, pixels_xy)
+
+    return (pixels_xy, pixels_xy)

AMS_BP/photophysics/photon_physics.py

@@ -0,0 +1,181 @@
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Tuple
+
+import numpy as np
+
+from ..optics.camera.detectors import photon_noise
+from ..optics.camera.quantum_eff import QuantumEfficiency
+from ..optics.filters.filters import FilterSpectrum
+from ..optics.psf.psf_engine import PSFEngine
+from ..sample.flurophores.flurophore_schema import (
+    SpectralData,
+    WavelengthDependentProperty,
+)
+from ..utils.constants import H_C_COM
+
+
+@dataclass
+class AbsorptionBase:
+    excitation_spectrum: SpectralData  # wl in nm, relative intensity
+    intensity_incident: WavelengthDependentProperty  # wl in nm, intensity in W/um^2
+    absorb_cross_section_spectrum: (
+        WavelengthDependentProperty  # wl in nm, cross section in cm^2
+    )
+
+    def __post_init__(self):
+        self.flux_density_precursor_lambda = self._calc_flux_density_precursor()
+
+    def _calc_flux_density_precursor(self) -> WavelengthDependentProperty:
+        """Per wavelength of incoming light W/cm^2 (intensity), find the quantity W/cm over the excitation_spectrum provided"""
+        wavelengths = []
+        ex_flux_density_lambda = []
+        for i in range(len(self.intensity_incident.wavelengths)):
+            intensity = self.intensity_incident.values[i]
+            wavelength = self.intensity_incident.wavelengths[i]
+            ex_spectrum = self.excitation_spectrum.get_value(wavelength)
+            ex_flux_density_lambda.append(intensity * wavelength * ex_spectrum)
+            wavelengths.append(wavelength)
+        return WavelengthDependentProperty(
+            wavelengths=wavelengths, values=ex_flux_density_lambda
+        )
+
+
+@dataclass
+class AbsorptionPhysics(AbsorptionBase):
+    fluorescent_lifetime_inverse: float
+
+    def saturation_rate(self, rate: float, max_rate: float) -> float:
+        return np.clip(rate, 0, max_rate)
+
+    def absorbed_photon_rate(self) -> float:
+        """Calculate the rate of incident photons"""
+        if self.flux_density_precursor_lambda is None:
+            raise ValueError("Flux density not calculated")
+
+        photon_rate_lambda = 0  # adding up all the wavelength based intensity rates
+        for i in range(len(self.flux_density_precursor_lambda.wavelengths)):
+            cross_section = self.absorb_cross_section_spectrum.values[i]
+            int_inverse_seconds_i = (
+                cross_section
+                * self.flux_density_precursor_lambda.values[i]
+                * H_C_COM
+                * 1e-1
+            )
+
+            photon_rate_lambda += int_inverse_seconds_i
+        return self.saturation_rate(
+            photon_rate_lambda,
+            self.fluorescent_lifetime_inverse,
+        )  # 1/s, 10^-1 combined all conversion factors
+
+
+@dataclass
+class PhotoStateSwitchPhysics(AbsorptionBase):
+    quantum_yeild: float  # switching events per photon absorbed
+
+
+@dataclass
+class EmissionPhysics:
+    emission_spectrum: SpectralData  # wl in nm, normalied intensity
+    quantum_yield: WavelengthDependentProperty
+    transmission_filter: FilterSpectrum
+
+    def __post_init__(self):
+        # normalize emission spectrum
+        emission_spectrum_sum = sum(self.emission_spectrum.values)
+        self.emission_spectrum = SpectralData(
+            wavelengths=self.emission_spectrum.wavelengths,
+            intensities=[
+                val / emission_spectrum_sum for val in self.emission_spectrum.values
+            ],
+        )
+
+    def emission_photon_rate(
+        self,
+        total_absorbed_rate: float,  # 1/s
+    ) -> WavelengthDependentProperty:
+        """Calculate the rate of emitted photons (1/s)
+
+        Parameters:
+            total_absorbed_rate: float
+        """
+
+        wavelengths = []
+        emission_rate_lambda = []
+        for i in range(len(self.emission_spectrum.wavelengths)):
+            wavelengths.append(self.emission_spectrum.wavelengths[i])
+            emission_rate_lambda.append(
+                total_absorbed_rate
+                * self.quantum_yield.values[i]
+                * self.emission_spectrum.values[i]
+            )
+
+        return WavelengthDependentProperty(
+            wavelengths=wavelengths, values=emission_rate_lambda
+        )
+
+    def transmission_photon_rate(
+        self, emission_photon_rate_lambda: WavelengthDependentProperty
+    ) -> WavelengthDependentProperty:
+        """Calculate the rate of transmitted photons (1/s)
+
+        Parameters:
+            emission_photon_rate_lambda: WavelengthDependentProperty
+        """
+        wavelengths = []
+        transmission_rate_lambda = []
+        for i in range(len(emission_photon_rate_lambda.wavelengths)):
+            wavelengths.append(emission_photon_rate_lambda.wavelengths[i])
+            transmission_rate_lambda.append(
+                emission_photon_rate_lambda.values[i]
+                * self.transmission_filter.find_transmission(
+                    emission_photon_rate_lambda.wavelengths[i]
+                )
+            )
+
+        return WavelengthDependentProperty(
+            wavelengths=wavelengths, values=transmission_rate_lambda
+        )
+
+
+@dataclass
+class incident_photons:
+    transmission_photon_rate: WavelengthDependentProperty
+    quantumEff: QuantumEfficiency
+    psf: Callable[[float | int, Optional[float | int]], PSFEngine]
+    position: Tuple[float, float, float]
+
+    def __post_init__(self):
+        self.generator = []
+        for i in range(len(self.transmission_photon_rate.wavelengths)):
+            if self.transmission_photon_rate.values[i] > 0:
+                self.generator.append(
+                    self.psf(
+                        self.transmission_photon_rate.wavelengths[i], self.position[2]
+                    )
+                )
+            else:
+                self.generator.append(0)
+
+    def incident_photons_calc(self, dt: float) -> Tuple[float, List]:
+        photons = 0
+        psf_hold = []
+        for i in range(len(self.transmission_photon_rate.wavelengths)):
+            if self.transmission_photon_rate.values[i] > 0:
+                qe_lam = self.quantumEff.get_qe(
+                    self.transmission_photon_rate.wavelengths[i]
+                )
+                photons_n = self.transmission_photon_rate.values[i] * dt
+                photons += photons_n
+                psf_gen = (
+                    self.generator[i].normalize_psf(
+                        self.generator[i].psf_z(z_val=self.position[2]),
+                        mode="sum",
+                    )
+                    * self.generator[i].psf_z_xy0(z_val=self.position[2])
+                    * photons_n
+                )
+
+                psf_hold.append(photon_noise(psf_gen) * qe_lam)
+
+        return photons, psf_hold

AMS_BP/photophysics/state_kinetics.py

@@ -0,0 +1,146 @@
+from functools import cache
+from typing import Callable, Optional, Sequence, Tuple
+
+import numpy as np
+from pydantic import BaseModel
+
+from ..sample.flurophores.flurophore_schema import State, StateType
+from ..sample.sim_sampleplane import FluorescentObject
+
+
+class ErnoMsg(BaseModel):
+    success: bool
+    erno_time: Optional[float | None] = None
+    erno_end_state: Optional[State | None] = None
+
+
+class StateTransitionCalculator:
+    def __init__(
+        self,
+        flurophoreobj: FluorescentObject,
+        time_duration: int | float,
+        current_global_time: int,
+        laser_intensity_generator: Callable,
+    ) -> None:
+        self.flurophoreobj = flurophoreobj
+        self.time_duration = time_duration  # seconds
+        self.current_global_time = current_global_time  # ms (oversample motion time)
+        self.laser_intensity_generator = laser_intensity_generator
+        self.fluorescent_state_history = {}  # {fluorescent.state.name : [delta time (seconds), laser_intensites], ...}
+
+    def __call__(
+        self,
+    ) -> Tuple[dict, State, ErnoMsg]:
+        state, erno = self.MCMC()
+        return self.fluorescent_state_history, state, erno
+
+    def _initialize_state_hist(self, time_pos: int, time_laser: float) -> dict:
+        laser_intensities = self.laser_intensity_generator(
+            florPos=self.flurophoreobj.position_history[time_pos],
+            time=time_laser,
+        )
+        for i in self.flurophoreobj.fluorophore.states.values():
+            if i.state_type == StateType.FLUORESCENT:
+                self.fluorescent_state_history[i.name] = [0, laser_intensities]
+        return laser_intensities
+
+    def MCMC(self) -> Tuple[State, ErnoMsg]:
+        time = 0
+        transitions = self.flurophoreobj.state_history[self.current_global_time][2]
+        final_state_name = transitions[0].from_state
+        laser_intensities = self._initialize_state_hist(self.current_global_time, time)
+
+        while time < self.time_duration:
+            stateTransitionMatrixR = [
+                sum(
+                    state_transitions.rate()(laser["wavelength"], laser["intensity"])
+                    for laser in laser_intensities.values()
+                )
+                for state_transitions in transitions
+            ]  # 1/s
+            if not stateTransitionMatrixR:
+                break
+            if sum(stateTransitionMatrixR) == 0:
+                break
+
+            # print(final_state_name)
+            new_time, state_indx = ssa_step(
+                stateTransitionMatrixR
+            )  # seconds, index on transitions
+
+            state_name = transitions[state_indx].to_state
+            if new_time > self.time_duration:
+                erno_time = new_time - time
+                new_time = self.time_duration - time
+                erno = ErnoMsg(
+                    success=False,
+                    erno_time=erno_time,
+                    erno_end_state=self.flurophoreobj.fluorophore.states[state_name],
+                )
+                if (
+                    self.flurophoreobj.fluorophore.states[final_state_name].state_type
+                    == StateType.FLUORESCENT
+                ):
+                    # print("I glow inside")
+                    self.fluorescent_state_history[
+                        self.flurophoreobj.fluorophore.states[final_state_name].name
+                    ][0] += new_time
+                return self.flurophoreobj.fluorophore.states[final_state_name], erno
+
+            # print(new_time)
+
+            if (
+                self.flurophoreobj.fluorophore.states[final_state_name].state_type
+                == StateType.FLUORESCENT
+            ):
+                # print("I glow")
+                self.fluorescent_state_history[
+                    self.flurophoreobj.fluorophore.states[final_state_name].name
+                ][0] += new_time
+            final_state_name = state_name
+            transitions = self._find_transitions(state_name)
+            time += new_time
+        # find state
+        return self.flurophoreobj.fluorophore.states[final_state_name], ErnoMsg(
+            success=True
+        )
+
+    @cache
+    def _find_transitions(self, statename: str) -> list:
+        return [
+            stateTrans
+            for stateTrans in self.flurophoreobj.fluorophore.transitions.values()
+            if stateTrans.from_state == statename
+        ]
+
+
+def ssa_step(reaction_rates: Sequence[float | int]) -> tuple[float, int]:
+    """
+    Perform one step of the SSA simulation.
+
+    Parameters:
+    - reaction_rates: List of reaction rates [k1, k2, ...]
+
+    Returns:
+    - dt: Time step to advance
+    - next_event: Index of the next reaction (0-based)
+    """
+    # Calculate propensities
+    propensities = np.array(reaction_rates)
+    total_propensity = np.sum(propensities)
+
+    if total_propensity == 0:
+        raise ValueError("Total propensity is zero; no reactions can occur.")
+
+    # Draw two random numbers
+    r1, r2 = np.random.uniform(0, 1, size=2)
+
+    # Compute time step
+    dt = -np.log(r1) / total_propensity
+
+    # Determine the next reaction
+    cumulative_propensities = np.cumsum(propensities)
+    threshold = r2 * total_propensity
+    next_event = np.searchsorted(cumulative_propensities, threshold)
+
+    return dt, next_event

AMS_BP/probabilityfuncs/markov_chain.py

@@ -0,0 +1,143 @@
+from functools import cache
+
+import numpy as np
+from scipy.linalg import fractional_matrix_power
+
+
+def MCMC_state_selection(
+    initial_state_index: int,
+    transition_matrix: np.ndarray,
+    possible_states: np.ndarray,
+    n: int,
+) -> np.ndarray:
+    """
+    Markov Chain Monte Carlo (MCMC) state selection.
+
+    This function simulates state transitions using a Markov Chain Monte Carlo method.
+    It selects the next state based on the current state and a transition matrix over `n` iterations.
+    The probability in the transition matrix is the probability of switching to a new state in the "time" step from n-1 -> n.
+
+    Parameters:
+    -----------
+    initial_state_index : int
+        The index of the initial state in the possible states.
+    transition_matrix : np.ndarray
+        A square matrix representing the transition probabilities between states.
+    possible_states : np.ndarray
+        An array of possible states for the system.
+    n : int
+        The number of iterations to perform.
+
+    Returns:
+    --------
+    np.ndarray
+        An array of selected states at each iteration.
+    """
+    # initialize the state selection
+    state_selection = np.zeros(n)
+    # initialize the current state
+    current_state = possible_states[initial_state_index]
+    current_state_index = initial_state_index
+    # iterate through the number of iterations
+    for i in range(n):
+        # find the probability of switching to each state
+        state_probability = transition_matrix[current_state_index]
+        # find the next state
+        next_state_index = np.random.choice(
+            np.arange(len(possible_states)), p=state_probability
+        )
+        next_state = possible_states[next_state_index]
+        # update the current state
+        current_state = next_state
+        current_state_index = next_state_index
+        state_selection[i] = current_state
+    return state_selection
+
+
+def MCMC_state_selection_rate(
+    initial_state_index: int,
+    transition_matrix: np.ndarray,  # in rate, (1/s) s= seconds
+    possible_states: np.ndarray,
+    n: int,
+    time_unit: int,  # amount of time (ms) in one n; ms = milliseconds
+):
+    # convert transition_matrix to probability
+    # divide elementwise to convert 1/s -> 1/ms */1000
+    transition_matrix = transition_matrix * (1.0 / 1000.0)
+    # convert to prob
+    for i in range(len(transition_matrix)):
+        for j in range(len(transition_matrix[i])):
+            transition_matrix[i][j] = rate_to_probability(
+                transition_matrix[i][j], time_unit
+            )
+
+    assert np.sum(transition_matrix, axis=0) == 1.0
+
+    # apply "MCMC_state_selection
+    return MCMC_state_selection(
+        initial_state_index=initial_state_index,
+        transition_matrix=transition_matrix,
+        possible_states=possible_states,
+        n=n,
+    )
+
+
+# convert from rate (1/s) to probability (0-1)
+@cache
+def rate_to_probability(rate: float, dt: float) -> float:
+    """Convert from rate (1/s) to probability (0-1)
+
+    Parameters:
+    -----------
+    rate : float
+        The rate (1/s)
+    dt : float
+        The time step (s) for the probability calculation
+
+    Returns:
+    --------
+    float
+        The probability (0-1)
+    """
+    return 1 - np.exp(-rate * dt)
+
+
+# convert from probability (0-1) to rate (1/s)
+@cache
+def probability_to_rate(probability: float, dt: float) -> float:
+    """Convert from probability (0-1) to rate (1/s)
+
+    Parameters:
+    -----------
+    probability : float
+        The probability (0-1)
+    dt : float
+        The time step (s) for the probability calculation
+    """
+    return -np.log(1 - probability) / dt
+
+
+# fractional probability util
+def change_prob_time(
+    probability: np.ndarray | float, dt: float, dt_prime: float
+) -> np.ndarray:
+    """Change the probability defined for dt to dt'
+
+    Parameters:
+    -----------
+    probability : np.ndarray | float
+        The probability (0-1)
+    dt : float
+        The time step (s) for the probability calculation
+    dt_prime : float
+        The new time step (s) for the probability calculation
+
+    Returns:
+    --------
+    np.ndarray | float
+        The probability (0-1)
+    """
+    if isinstance(probability, np.ndarray):
+        return fractional_matrix_power(probability, dt_prime / dt)
+    else:
+        return probability ** (dt_prime / dt)