polytope.py

# This file is part of BurnMan - a thermoelastic and thermodynamic toolkit
# for the Earth and Planetary Sciences
# Copyright (C) 2012 - 2021 by the BurnMan team, released under the GNU
# GPL v2 or later.

from __future__ import absolute_import

import numpy as np
from sympy import Matrix
from scipy.linalg import block_diag

import logging
import importlib
from ..classes.polytope import MaterialPolytope, independent_row_indices
from ..classes.solution import Solution
from ..classes.composite import Composite
from .solution import transform_solution_to_new_basis

try:
    cp = importlib.import_module("cvxpy")
except ImportError as err:
    print(
        f"Warning: {err}. " "For full functionality of BurnMan, please install cvxpy."
    )


def solution_polytope_from_charge_balance(
    charges, charge_total, return_fractions=False
):
    """
    Creates a polytope object from a list of the charges for each species on
    each site and the total charge for all site-species.

    :param charges: 2D list containing the total charge for species j on site i,
        including the site multiplicity. So, for example,
        a solution with the site formula [Mg,Fe]3[Mg,Al,Si]2Si3O12 would
        have the following list: [[6., 6.], [4., 6., 8.]].
    :type charges: 2D list of floats
    :param charge_total: The total charge for all site-species per formula unit.
        The example given above would have charge_total = 12.
    :type charge_total: float
    :param return_fractions: Determines whether the created polytope object returns its
        attributes (such as endmember occupancies) as fractions or as floats.
        Default is False.
    :type return_fractions: bool

    :returns: A polytope object corresponding to the parameters provided.
    :rtype: :class:`burnman.polytope.MaterialPolytope` object
    """
    n_sites = len(charges)
    all_charges = np.concatenate(charges)
    n_site_elements = len(all_charges)
    equalities = np.empty((n_sites + 1, n_site_elements + 1))
    equalities[:-1, 0] = -1
    i = 0
    for i_site, site_charges in enumerate(charges):
        equalities[i_site, 1:] = [
            1 if (j >= i and j < i + len(site_charges)) else 0
            for j in range(n_site_elements)
        ]
        i += len(site_charges)

    equalities[-1, 0] = -charge_total
    equalities[-1, 1:] = all_charges

    pos_constraints = np.concatenate(
        (np.zeros((len(equalities[0]) - 1, 1)), np.identity(len(equalities[0]) - 1)),
        axis=1,
    )
    return MaterialPolytope(
        equalities, pos_constraints, return_fractions=return_fractions
    )


def solution_polytope_from_endmember_occupancies(
    endmember_occupancies, return_fractions=False
):
    """
    Creates a polytope object from a list of independent endmember occupancies.

    :param endmember_occupancies: 2D list containing the
        site-species occupancies j for endmember i.
        So, for example, a solution with independent endmembers
        [Mg]3[Al]2Si3O12, [Mg]3[Mg0.5Si0.5]2Si3O12, [Fe]3[Al]2Si3O12
        might have the following array:
        [[1., 0., 1., 0., 0.],
        [1., 0., 0., 0.5, 0.5],
        [0., 1., 1., 0., 0.]],
        where the order of site-species is Mg_A, Fe_A, Al_B, Mg_B, Si_B.
    :type endmember_occupancies: 2D numpy array

    :param return_fractions: Determines whether the created polytope object
        returns its attributes (such as endmember occupancies) as fractions
        or as floats.
    :type return_fractions: bool

    :returns: A polytope object corresponding to the parameters provided.
    :rtype: :class:`burnman.polytope.MaterialPolytope` object
    """
    n_sites = sum(endmember_occupancies[0])
    n_occs = endmember_occupancies.shape[1]

    nullspace = np.array(Matrix(endmember_occupancies).nullspace(), dtype=float)

    equalities = np.zeros((len(nullspace) + 1, n_occs + 1))
    equalities[0, 0] = -n_sites
    equalities[0, 1:] = 1

    if len(nullspace) > 0:
        try:
            equalities[1:, 1:] = nullspace
        except ValueError:
            equalities[1:, 1:] = nullspace[:, :, 0]

    pos_constraints = np.concatenate(
        (np.zeros((len(equalities[0]) - 1, 1)), np.identity(len(equalities[0]) - 1)),
        axis=1,
    )

    return MaterialPolytope(
        equalities,
        pos_constraints,
        return_fractions=return_fractions,
        independent_endmember_occupancies=endmember_occupancies,
    )


def composite_polytope_at_constrained_composition(
    composite, composition, return_fractions=False
):
    """
    Creates a polytope object from a Composite object and a composition.
    This polytope describes the complete set of valid composite
    endmember amounts that satisfy the compositional constraints.

    :param composite: A composite containing one or more Solution and Mineral objects.
    :type composite: :class:`burnman.Composite` object

    :param composition: A dictionary containing the amounts of each element.
    :type composition: dict

    :param return_fractions: Determines whether the created polytope object returns its
        attributes (such as endmember occupancies) as fractions or as floats.
    :type return_fractions: bool


    :returns: A polytope object corresponding to the parameters provided.
    :rtype: :class:`burnman.polytope.MaterialPolytope` object
    """
    c_array = np.empty((composite.n_elements, 1))
    c_array[:, 0] = [
        -composition[e] if e in composition else 0.0 for e in composite.elements
    ]

    equalities = np.concatenate((c_array, composite.stoichiometric_array.T), axis=1)

    eoccs = []
    for i, ph in enumerate(composite.phases):
        if isinstance(ph, Solution):
            eoccs.append(ph.solution_model.endmember_occupancies.T)
        else:
            eoccs.append(np.ones((1, 1)))

    eoccs = block_diag(*eoccs)
    inequalities = np.concatenate((np.zeros((len(eoccs), 1)), eoccs), axis=1)

    return MaterialPolytope(
        equalities, inequalities, number_type="float", return_fractions=return_fractions
    )


def simplify_composite_with_composition(composite, composition):
    """
    Takes a composite and a composition, and returns the simplest composite
    object that spans the solution space at the given composition.

    For example, if the composition is given as {'Mg': 2., 'Si': 1.5, 'O': 5.},
    and the composite is given as a mix of Mg,Fe olivine and pyroxene
    solutions, this function will return a composite that only contains
    the Mg-bearing endmembers.

    :param composite: The initial Composite object.
    :type composite: :class:`burnman.Composite` object

    :param composition: A dictionary containing the amounts of each element.
    :type composition: dict

    :returns: The simplified Composite object
    :rtype: :class:`burnman.Composite` object
    """
    polytope = composite_polytope_at_constrained_composition(
        composite, composition, return_fractions=True
    )

    composite_changed = False
    new_phases = []
    mbr_amounts = polytope.endmember_occupancies
    i = 0
    for i_ph, n_mbrs in enumerate(composite.endmembers_per_phase):
        ph = composite.phases[i_ph]

        amounts = mbr_amounts[:, i : i + n_mbrs].astype(float)
        i += n_mbrs

        rank = np.linalg.matrix_rank(amounts, tol=1.0e-8)

        if rank < n_mbrs:
            if isinstance(ph, Solution) and rank > 0:
                if len(amounts) > 1:
                    c_mean = np.mean(amounts, axis=0)
                else:
                    c_mean = amounts[0]

                poly = solution_polytope_from_endmember_occupancies(
                    ph.solution_model.endmember_occupancies
                )
                dmbrs = poly.endmembers_as_independent_endmember_amounts

                x = cp.Variable(dmbrs.shape[0])
                objective = cp.Minimize(cp.sum_squares(x))
                constraints = [dmbrs.T @ x == c_mean, x >= 0]

                prob = cp.Problem(objective, constraints)
                prob.solve()

                mbr_indices = np.argsort(x.value)[::-1]
                ind_indices = [i for i in mbr_indices if x.value[i] > 1.0e-6]
                new_basis = dmbrs[ind_indices]

                # And now reduce the new basis if necessary
                new_basis = new_basis[independent_row_indices(new_basis)]

                if len(new_basis) < ph.n_endmembers:
                    logging.info(
                        f"Phase {i_ph} ({ph.name}) is "
                        "rank-deficient ({rank} < {n_mbrs}). "
                        "The transformed solution is described "
                        f"using {len(new_basis)} endmembers."
                    )

                    composite_changed = True
                    soln = transform_solution_to_new_basis(ph, new_basis)
                    new_phases.append(soln)
                else:
                    logging.info(
                        "This solution is rank-deficient "
                        f"({rank} < {n_mbrs}), "
                        "but its composition requires all "
                        "independent endmembers."
                    )
            else:
                composite_changed = True
                logging.info(
                    f"Phase {i_ph} ({ph.name}) removed from " "composite (rank = 0)."
                )
        else:
            new_phases.append(ph)

    if composite_changed:
        return Composite(new_phases)
    else:
        return composite