Flory-Huggins

Flory-Huggins.py

import scipy.optimize
import numpy as np
import matplotlib.pyplot as plt


def F(
    phi: np.ndarray,
    uaa: float, ubb: float, uab: float, na: int, nb: int,
) -> np.ndarray:
    """
    Helmholtz-free energy for a Flory-Huggins-like model

    :param phi: (2xAxB) array of phi independent coordinates
    :return: AxB array of function values
    """
    coef_shape = (2,) + (1,)*(len(phi.shape) - 1)
    nanb = np.array((na, nb)).reshape(coef_shape)
    uaub = np.array((uaa, ubb)).reshape(coef_shape)
    tot = phi.sum(axis=0)

    entropy = (
        (phi / nanb * np.log(phi)).sum(axis=0)
        + (1 - tot)*np.log(1 - tot)
    )
    energy = (
        phi**2 * (-0.5 * uaub)
    ).sum(axis=0) - uab * phi.prod(axis=0)
    return entropy + energy


def gradF(
    phi: np.ndarray,
    uaa: float, ubb: float, uab: float, na: int, nb: int,
) -> np.ndarray:
    """
    First-order Jacobian of F in phi.
    :param phi: (2xAxB) array of phi independent coordinates
    :return: (2xAxB) array of gradient vectors
    """
    coef_shape = (2,) + (1,)*(len(phi.shape) - 1)
    nanb = np.array((na, nb)).reshape(coef_shape)
    uaub = np.array((uaa, ubb)).reshape(coef_shape)
    return (
        (np.log(phi) + 1) / nanb
        - np.log(1 - phi.sum(axis=0))
        - phi * uaub
        - phi[::-1] * uab
        - 1
    )


def hessianF(
    phi: np.ndarray,
    uaa: float, ubb: float, uab: float, na: int, nb: int,
) -> np.ndarray:
    """
    Hessian of F in phi.
    :param phi: (2xAxB) array of phi independent coordinates
    :return: (2x2xAxB) array of Hessian matrices
    """
    coef_shape = (2,) + (1,)*(len(phi.shape) - 1)
    nanb = np.array((na, nb)).reshape(coef_shape)
    uaub = np.array((uaa, ubb)).reshape(coef_shape)

    diag = 1/phi/nanb - uaub
    antidiag = -uab
    hessian = np.empty(shape=(2, 2, *phi.shape[1:]), dtype=phi.dtype)
    hessian[[0, 1], [0, 1], ...] = diag
    hessian[[0, 1], [1, 0], ...] = antidiag
    hessian += 1/(1 - phi.sum(axis=0))

    return hessian


def test_grad(params: tuple[float, ...]) -> None:
    """
    Sanity test to ensure that analytic gradients are correct.
    """
    xk = np.array((0.2, 0.3))
    assert scipy.optimize.check_grad(F, gradF, xk, *params) < 1e-6

    estimated = scipy.optimize.approx_fprime(xk, gradF, 1e-9, *params)
    analytic = hessianF(xk, *params)
    assert np.allclose(estimated, analytic, rtol=0, atol=1e-6)


def mu(
    phi: np.ndarray,
    uaa: float, ubb: float, uab: float, na: int, nb: int,
) -> np.ndarray:
    return gradF(phi, uaa, ubb, uab, na, nb)


def Pi(phi: np.ndarray, *params) -> np.ndarray:
    mu_ab = mu(phi, *params)
    return (mu_ab * phi).sum(axis=0) - F(phi, *params)


def plot_isodims(
    phi: np.ndarray,
    mua: np.ndarray,
    mub: np.ndarray,
    pi: np.ndarray,
) -> plt.Figure:
    """
    Contour plot of individual mua/b/pi series, all over phi_a and phi_b
    """
    fig: plt.Figure
    axes: tuple[plt.Axes, ...]
    fig, axes = plt.subplots(nrows=1, ncols=3)
    fig.suptitle('Individual series')

    (ax_mua, ax_mub, ax_pi) = axes
    for ax in axes:
        ax.set_xlabel('phi a')
        ax.grid(True)
    ax_mua.set_ylabel('phi b')
    ax_mub.set_yticklabels(())
    ax_pi.set_yticklabels(())

    ax_mua.contour(*phi, mua)
    ax_mua.set_title('mu a')
    ax_mub.contour(*phi, mub)
    ax_mub.set_title('mu b')
    ax_pi.contour(*phi, pi)
    ax_pi.set_title('pi')

    return fig


def plot_example_proximity(
    phi: np.ndarray,
    mua: np.ndarray,
    mub: np.ndarray,
    pi: np.ndarray,
    params: tuple[float, ...],
    phi_pimin: np.ndarray,
    phi_endpoints: np.ndarray,
    zapprox: np.ndarray,
) -> plt.Figure:
    """
    Plot, for the initial-fit origin point (1) and its pair (2), all three isocurves
    Include the Hessian determinant estimation (zapprox)
    """
    phi0 = phi_endpoints[:, 0]  # origin point (1)
    phi1 = phi_endpoints[:, 1]  # next point (2)
    mua0, mub0 = mu(phi0, *params)   # mu values at origin
    pi0 = Pi(phi0, *params)  # pi value at origin

    # Array indices where mu closely matches that at the origin
    mua_close_y, mua_close_x = (np.abs(mua - mua0) < 1e-3).nonzero()
    mub_close_y, mub_close_x = (np.abs(mub - mub0) < 1e-3).nonzero()

    # Error from the pi value seen at the origin
    pi_close = pi - pi0

    fig: plt.Figure
    ax: plt.Axes
    fig, ax = plt.subplots()
    ax.set_title('Isocurves from initial fit origin to next,'
                 '\nwith Hessian determinant')

    ax.scatter(*phi0, c='black', marker='+', s=100, label='origin (1)')
    ax.scatter(*phi1, c='pink', marker='+', s=100, label='next (2)')
    ax.scatter(*phi_pimin, c='blue', marker='+', s=100, label='pi_min')
    ax.plot(phi[0, 0, mua_close_x], phi[0, 0, mua_close_y], label='mua')
    ax.plot(phi[1, mub_close_x, 0], phi[1, mub_close_y, 0], label='mub')
    ax.contour(*phi, pi_close, levels=[0], colors='green')  # Pi isocurve
    ax.contour(*phi, zapprox, levels=[0], colors='red')  # Hessian determinant
    ax.legend()

    return fig


def initial_pair_estimate(
    phi_min: float, phi_max: float,
    phi_pimin: np.ndarray,
    params: tuple[float, ...],
) -> np.ndarray:
    """
    Perform an initial fit to find any pair on the target curve

    :param phi_min: Minimum search space bound of phi in both axes
    :param phi_max: Maximum search space bound of phi in both axes
    :param phi_pimin: Phi coordinates of minimum of pi; the search space is centred here
    :param params: uaa, ubb, uab, na, nb
    :return: endpoint coordinates in phi (2x2)
    """
    phia_pimin, phib_pimin = phi_pimin

    # As in original problem construction, the first point must be below-left of the second point;
    # and the distance must be greater than 0 to avoid degeneracy (superimposed pair).
    nondegenerate = scipy.optimize.LinearConstraint(
        A=[
            [-1, 1,  0, 0],
            [ 0, 0, -1, 1],
        ],
        lb=0.05,
    )

    def least_sq_difference(phi: np.ndarray) -> float:
        """
        From given phi endpoints, calculate exact (not estimated) values for mu and pi, and return
        their least-squared distance as the objective cost
        """
        phi = phi.reshape((2, 2))
        mua, mub = mu(phi, *params)
        pi = Pi(phi, *params)
        return (
            (mua[0] - mua[1])**2 +
            (mub[0] - mub[1])**2 +
            10*(pi[0] - pi[1])**2
        )

    result = scipy.optimize.minimize(
        fun=least_sq_difference,
        x0=(
            phia_pimin,   phia_pimin*3/2,
            phib_pimin/2, phib_pimin,
        ),
        # The origin is below-left of the minimum location of pi
        bounds=scipy.optimize.Bounds(
            lb=[phi_min,    phi_min,    phi_min,    phi_min],
            ub=[phia_pimin, phi_max,    phib_pimin, phi_max],
        ),
        constraints=nondegenerate,
        tol=1e-12,
    )
    if not result.success:
        raise ValueError(result.message)

    return result.x.reshape((2, 2))


def characterise_initial(
    params: tuple[float, ...],
    phi_min: float = 0,
    phi_max: float = 1,
) -> tuple[
    np.ndarray,  # mesh-like phi coordinate space, 2xAxB
    np.ndarray,  # mua, AxB
    np.ndarray,  # mub, AxB
    np.ndarray,  # pi, AxB
    np.ndarray,  # Hessian determinant estimate, AxB
]:
    phi_a = phi_b = np.linspace(phi_min, phi_max, 500)
    phi = np.stack(np.meshgrid(phi_a, phi_b))
    mua, mub = mu(phi, *params)
    pi = Pi(phi, *params)

    hess = hessianF(phi, *params)
    zapprox = np.linalg.det(hess.T)

    return phi, mua, mub, pi, zapprox


def estimate_pimin(phi: np.ndarray, pi: np.ndarray) -> np.ndarray:
    """
    Start referenced from the (estimated) minimum of pi, in the middle of the region of interest
    The Hessian estimate runs through this point.
    """
    ijmin = np.unravel_index(pi.argmin(), pi.shape)
    phi_pimin = phi[:, ijmin[0], ijmin[1]]
    print(f'Pi minimum point of {pi[ijmin]:.5f} at {phi_pimin}')
    return phi_pimin


def main() -> None:
    uaa, ubb, uab = 0, 0, 4.36
    na, nb = 10, 6
    phi_min, phi_max = 1e-2, 0.4
    params = (uaa, ubb, uab, na, nb)

    test_grad(params)

    phi, mua, mub, pi, zapprox = characterise_initial(params, phi_min, phi_max)
    phi_pimin = estimate_pimin(phi, pi)
    phi_endpoints = initial_pair_estimate(phi_min, phi_max, phi_pimin, params)

    plot_isodims(phi, mua, mub, pi)
    plot_example_proximity(
        phi, mua, mub, pi, params, phi_pimin,
        phi_endpoints=phi_endpoints, zapprox=zapprox,
    )
    plt.show()


if __name__ == '__main__':
    main()