hyperion.medium

def henyey_greenstein_scattering_angle(key, g=0.9):
    """Henyey-Greenstein scattering in one plane.

    Parameters
    ----------
    key : jax.random.PRNGKey
        Random key for sampling.
    g : float, optional
        Asymmetry parameter (default is 0.9).

    Returns
    -------
    float
        Scattering angle in radians.
    """
    eta = random.uniform(key)
    costheta = 1 / (2 * g) * (1 + g**2 - ((1 - g**2) / (1 + g * (2 * eta - 1))) ** 2)
    return jnp.arccos(costheta)

def liu_scattering_angle(key, g=0.95):
    """Simplified Liu scattering.

    Parameters
    ----------
    key : jax.random.PRNGKey
        Random key for sampling.
    g : float, optional
        Asymmetry parameter (default is 0.95).

    Returns
    -------
    float
        Scattering angle in radians.

    Notes
    -----
    See: https://arxiv.org/pdf/1301.5361.pdf
    """
    beta = (1 - g) / (1 + g)
    xi = random.uniform(key)
    costheta = 2 * xi**beta - 1
    return jnp.arccos(costheta)

def make_mixed_scattering_func(f1, f2, ratio):
    """Create a mixture model with two sampling functions.

    Parameters
    ----------
    f1 : callable
        Sampling function taking one argument (random key).
    f2 : callable
        Sampling function taking one argument (random key).
    ratio : float
        Fraction of samples drawn from ``f1``.

    Returns
    -------
    callable
        Mixture sampling function.
    """

    def _f(key):
        """Mixture sampler selecting between ``f1`` and ``f2`` based on ``ratio``.

        Parameters
        ----------
        key : jax.random.PRNGKey
            PRNG key for sampling.

        Returns
        -------
        float
            Sampled scattering angle.
        """
        k1, k2 = random.split(key)
        is_f1 = random.uniform(k1) < ratio

        return cond(is_f1, f1, f2, k2)

    return _f

def make_ref_index_func(salinity, temperature, pressure):
    """Create a function that returns refractive index as function of wavelength.

    Parameters
    ----------
    salinity : float
        Salinity in parts per thousand.
    temperature : float
        Temperature in C.
    pressure : float
        Pressure in bar.

    Returns
    -------
    callable
        Function mapping wavelength (nm) to refractive index.
    """
    n0 = 1.31405
    n1 = 1.45e-5
    n2 = 1.779e-4
    n3 = 1.05e-6
    n4 = 1.6e-8
    n5 = 2.02e-6
    n6 = 15.868
    n7 = 0.01155
    n8 = 0.00423
    n9 = 4382
    n10 = 1.1455e6

    a01 = (
        n0
        + (n2 - n3 * temperature + n4 * temperature * temperature) * salinity
        - n5 * temperature * temperature
        + n1 * pressure
    )
    a2 = n6 + n7 * salinity - n8 * temperature
    a3 = -n9
    a4 = n10

    def ref_index_func(wavelength):
        """Return refractive index for given wavelength(s).

        Parameters
        ----------
        wavelength : float or array-like
            Wavelength in nanometres.

        Returns
        -------
        float or array-like
            Refractive index corresponding to input wavelength(s).
        """

        x = 1 / wavelength
        return a01 + x * (a2 + x * (a3 + x * a4))

    return ref_index_func

def make_wl_dep_sca_len_func(vol_conc_small_part, vol_conc_large_part):
    """Create a function that calculates the scattering length based on particle concentrations.

    Copied from clsim.

    Parameters
    ----------
    vol_conc_small_part : float
        Volumetric concentration of small particles (ppm).
    vol_conc_large_part : float
        Volumetric concentration of large particles (ppm).

    Returns
    -------
    callable
        Function that maps wavelength (nm) to scattering length.
    """

    def sca_len(wavelength):
        """Compute scattering length as a function of wavelength (nm).

        Parameters
        ----------
        wavelength : float or array-like
            Wavelength in nanometres.

        Returns
        -------
        float or np.ndarray
            Scattering length in same units as the input.
        """
        ref_wlen = 550  # nm
        x = ref_wlen / wavelength

        sca_coeff = (
            0.0017 * jnp.power(x, 4.3)
            + 1.34 * vol_conc_small_part * jnp.power(x, 1.7)
            + 0.312 * vol_conc_large_part * jnp.power(x, 0.3)
        )

        return 1 / sca_coeff

    return sca_len

def rayleigh_scattering_angle(key):
    """Rayleigh scattering. Adapted from clsim.

    Parameters
    ----------
    key : jax.random.PRNGKey
        Random key for sampling.

    Returns
    -------
    float
        Scattering angle in radians.
    """
    b = 0.835
    p = 1.0 / 0.835

    q = (b + 3.0) * ((random.uniform(key)) - 0.5) / b
    d = q * q + p * p * p

    u1 = -q + jnp.sqrt(d)
    u = jnp.cbrt(jnp.abs(u1)) * jnp.sign(u1)

    v1 = -q - jnp.sqrt(d)
    v = jnp.cbrt(jnp.abs(v1)) * jnp.sign(v1)

    return jnp.arccos(jax.lax.clamp(-1.0, u + v, 1.0))