compositecalc/geometry/orientation.py

"""Orientation averaging for non-spherical inclusions.

For non-spherical inclusions the effective property depends on how the
inclusion's symmetry axis is oriented relative to the measurement direction.
This module provides wrappers that average a *directional* model over
common orientation distributions.

All averaging functions accept a *directional_eff* callable::

    P_eff_i = directional_eff(L_i)

where ``L_i`` is the depolarization factor along principal axis *i* and
``P_eff_i`` is the effective property computed by an aligned-inclusion EMT
model (e.g. Mori-Tanaka, Hatta-Taya, Nan) along that axis.

Orientation convention
----------------------
For a uniaxial inclusion (prolate or oblate spheroid) the depolarization
tensor is diagonal with two distinct eigenvalues:

- ``L_sym`` — factor along the inclusion's **symmetry axis**.
  For a prolate spheroid L_sym < 1/3; for an oblate spheroid L_sym > 1/3.
- ``L_trans`` — factor **transverse** to the symmetry axis
  (the other two principal values are equal by uniaxial symmetry).

When the symmetry axis makes a polar angle θ with the measurement direction,
the effective depolarization factor "seen" along the measurement direction is

    L(θ) = L_sym cos²θ + L_trans sin²θ

This rotated value is used in :func:`orient_avg_custom`.

References
----------
Nan, C.-W. et al. (1997). Effective thermal conductivity of particulate
    composites with interfacial thermal resistance. *J. Appl. Phys.*
    81(10), 6692. — eq. (A4)–(A6) for the random orientation average.
Hatta, H. & Taya, M. (1986). Thermal conductivity of coated filler
    composites. *J. Appl. Phys.* 59(6), 1851. — eq. (16).
Benveniste, Y. (1987). A new approach to the application of Mori-Tanaka's
    theory in composite materials. *Mech. Mater.* 6, 147.
"""

from __future__ import annotations

from typing import Callable, Union

import numpy as np

__all__ = [
    "orient_avg_3d",
    "orient_avg_2d",
    "orient_avg_custom",
]

# Type alias for the return value of directional_eff
_Result = Union[float, np.ndarray]


def orient_avg_3d(
    directional_eff: Callable[[float], _Result],
    L_a: float,
    L_b: float,
    L_c: float,
) -> _Result:
    """Isotropic 3D orientation average: arithmetic mean over three principal axes.

    For inclusions whose symmetry axes are distributed uniformly over all
    solid angles, the effective scalar property is the arithmetic mean of
    the aligned-inclusion property computed along each principal axis:

        P_eff = (P_a + P_b + P_c) / 3,   P_i = directional_eff(L_i)

    This is the standard Mori-Tanaka / Nan orientation-averaging
    approximation for randomly oriented inclusions.

    Parameters
    ----------
    directional_eff : callable
        Function ``f(L) -> P_eff`` that computes the effective property
        along one principal axis given its depolarization factor *L*.
        May return a Python float or a NumPy array (e.g. frequency-
        dependent complex permittivity).
    L_a, L_b, L_c : float
        Depolarization factors along the three principal axes.
        Must satisfy ``L_a + L_b + L_c = 1``.

    Returns
    -------
    P_eff : float or np.ndarray
        Isotropically averaged effective property, same type as the
        return value of *directional_eff*.

    References
    ----------
    Nan et al. (1997) eq. (A6); Benveniste (1987) eq. (6).

    Examples
    --------
    >>> f = lambda L: 1.0 - L          # simple linear model
    >>> orient_avg_3d(f, 1/3, 1/3, 1/3)  # sphere: average = f(1/3)
    0.6666666666666667
    >>> # Prolate needle (L_a=0, L_b=L_c=0.5): (1 + 0.5 + 0.5) / 3 = 2/3
    >>> round(orient_avg_3d(f, 0.0, 0.5, 0.5), 10)
    0.6666666667
    """
    return (directional_eff(L_a) + directional_eff(L_b) + directional_eff(L_c)) / 3.0


def orient_avg_2d(
    directional_eff: Callable[[float], _Result],
    L_a: float,
    L_b: float,
    L_c: float,
) -> tuple[_Result, _Result]:
    """In-plane random, out-of-plane aligned orientation average.

    Models thin-film or paper-like composites where inclusion symmetry
    axes lie randomly in the (*a*, *b*) plane but are fixed in the *c*
    direction:

        P_in  = (directional_eff(L_a) + directional_eff(L_b)) / 2
        P_out =  directional_eff(L_c)

    The returned pair gives the effective property measured in-plane and
    out-of-plane, respectively.

    Parameters
    ----------
    directional_eff : callable
        Function ``f(L) -> P_eff`` as for :func:`orient_avg_3d`.
    L_a, L_b : float
        Depolarization factors in the two in-plane directions.
    L_c : float
        Depolarization factor along the out-of-plane direction.

    Returns
    -------
    P_in_plane : float or np.ndarray
        Effective property for in-plane measurement.
    P_out_plane : float or np.ndarray
        Effective property for out-of-plane measurement (aligned value).

    Examples
    --------
    >>> f = lambda L: 1.0 - L
    >>> P_in, P_out = orient_avg_2d(f, 1/3, 1/3, 1/3)
    >>> P_in == P_out          # sphere: isotropic → in-plane = out-of-plane
    True
    >>> # Platelet in xy-plane: L_a = L_b = 0.1 (in-plane), L_c = 0.8 (thin axis)
    >>> P_in, P_out = orient_avg_2d(f, 0.1, 0.1, 0.8)
    >>> P_in > P_out           # lower L_a,b → higher f → P_in > P_out
    True
    """
    P_in = (directional_eff(L_a) + directional_eff(L_b)) / 2.0
    P_out = directional_eff(L_c)
    return P_in, P_out


def orient_avg_custom(
    directional_eff: Callable[[float], _Result],
    L_sym: float,
    L_trans: float,
    thetas: np.ndarray | list[float],
    weights: np.ndarray | list[float],
) -> _Result:
    """Weighted average over a discrete orientation distribution function (ODF).

    For uniaxial inclusions whose symmetry axes point at polar angles
    *thetas* from the measurement direction, the effective depolarization
    factor at each orientation is

        L(θ) = L_sym cos²θ + L_trans sin²θ

    The orientation-averaged effective property is then

        P_eff = Σᵢ weights[i] · directional_eff(L(θᵢ))

    This allows arbitrary discrete ODFs (e.g. Gaussian, bimodal, partial
    alignment) as a weighted sum over sampled orientations.

    Parameters
    ----------
    directional_eff : callable
        Function ``f(L) -> P_eff`` as for :func:`orient_avg_3d`.
    L_sym : float
        Depolarization factor along the inclusion's symmetry axis.
        For prolate spheroids L_sym < 1/3; for oblate L_sym > 1/3.
    L_trans : float
        Depolarization factor transverse to the symmetry axis.
    thetas : array-like of float
        Polar angles [rad] of the inclusion symmetry axes relative to the
        measurement direction.  Shape (N,).  Each entry in [0, π].
    weights : array-like of float
        Probability weights for each angle.  Shape (N,).  Must sum to 1
        within 1 × 10⁻⁸.

    Returns
    -------
    P_eff : float or np.ndarray
        Orientation-averaged effective property.

    Raises
    ------
    ValueError
        If *weights* do not sum to 1 within 1 × 10⁻⁸.

    References
    ----------
    Benveniste (1987) eq. (8) for the weighted orientation average.

    Examples
    --------
    >>> import numpy as np
    >>> f = lambda L: 1.0 - L
    >>> # Perfectly aligned along symmetry axis (θ = 0)
    >>> orient_avg_custom(f, L_sym=0.0, L_trans=0.5, thetas=[0.0], weights=[1.0])
    1.0
    >>> # Perfectly transverse (θ = π/2)
    >>> orient_avg_custom(f, L_sym=0.0, L_trans=0.5, thetas=[np.pi/2], weights=[1.0])
    0.5
    """
    thetas = np.asarray(thetas, dtype=float)
    weights = np.asarray(weights, dtype=float)

    weight_sum = float(weights.sum())
    if abs(weight_sum - 1.0) > 1e-8:
        raise ValueError(
            f"weights must sum to 1.0; got {weight_sum:.10g}."
        )

    L_eff = L_sym * np.cos(thetas) ** 2 + L_trans * np.sin(thetas) ** 2

    # Accumulate weighted sum; works for scalar and array-valued directional_eff.
    vals = np.array(
        [float(w) * directional_eff(float(L)) for L, w in zip(L_eff, weights)]
    )
    return np.sum(vals, axis=0)