simulate_equilibrium_solidification should do order/disorder checks automatically

[bocklund]

Currently there are no checks for order/disorder phase name checking in simulate_equilibrium_solidification, so a order/disorder modeled phase that is in a disordered configuration will still show as the ordered phase name.

The main challenge in making this change that the order/disorder check assumes a pycalphad LightDataset, rather than xarray Dataset.

[bocklund]

I wrote some code to do this for another project, but it needs to be integrated still. rename_disordered_phases could be used for both equilibrium and Scheil solidification simulations, but it would require using the xarray Dataset objects from pycalphad, i.e. calls to equilibrium in simulate_scheil_solidification cannot use the to_xarray=False performance optimization.

from dataclasses import dataclass
from typing import Sequence, List
import itertools
from collections import defaultdict
import numpy as np
import xarray as xr
from pycalphad.core.utils import unpack_components, generate_dof

@dataclass
class OrderingRecord:
    ordered_phase_name: str
    disordered_phase_name: str
    subl_dof: Sequence[int]  # number of degrees of freedom in each sublattice of the ordered phase
    symmetric_subl_idx: Sequence[Sequence[int]]  # List of sublattices (of the ordered phase) that are symmetric

    def is_disordered(self, site_fractions):
        # For each sublattice, create a `slice` object for slicing the site
        # fractions of that particular sublattice from the site fraction array
        subl_slices = []
        for subl_idx in range(len(self.subl_dof)):
            start_idx = np.sum(self.subl_dof[:subl_idx], dtype=np.int_)
            end_idx = start_idx + self.subl_dof[subl_idx]
            subl_slices.append(slice(start_idx, end_idx))

        # For each set of symmetrically equivalent sublattices
        for symm_subl in self.symmetric_subl_idx:
            # Check whether the site fractions of each pair of symmetrically
            # equivalent sublattices are ordered or disordered
            for idx1, idx2 in itertools.combinations(symm_subl, 2):
                # A phase is ordered if any pair of sublattices does not have
                # equal (within numerical tolerance) site fractions
                pair_is_ordered = np.any(~np.isclose(site_fractions[subl_slices[idx1]], site_fractions[subl_slices[idx2]]))
                if pair_is_ordered:
                    return False
        return True

def create_ordering_records(dbf, comps, phases):
    """Return a dictionary with the sublattice degrees of freedom and equivalent
    sublattices for order/disorder phases

    Parameters
    ----------
    dbf : pycalphad.Database
    comps : list[str]
        List of active components to consider
    phases : list[str]
        List of active phases to consider

    Returns
    -------
    List[OrderingRecord]

    Notes
    -----
    Phases which should be checked for ordered/disordered configurations are
    determined heuristically for this script.

    The heuristic for a phase satisfies the following:
    1. The phase is the ordered part of an order-disorder model
    2. The equivalent sublattices have all the same number of elements
    """
    species = unpack_components(dbf, comps)
    ordering_records = []
    for phase_name in phases:
        phase_obj = dbf.phases[phase_name]
        if phase_name == phase_obj.model_hints.get('ordered_phase', ''):
            # This phase is active and modeled with an order/disorder model.
            dof = generate_dof(dbf.phases[phase_name], species)[1]
            # Define the symmetrically equivalent sublattices as any sublattices
            # that have the same site ratio. Create a {site_ratio: [subl idx]} dict
            site_ratio_idxs = defaultdict(lambda: [])
            for subl_idx, site_ratio in enumerate(phase_obj.sublattices):
                site_ratio_idxs[site_ratio].append(subl_idx)
            equiv_sublattices = list(site_ratio_idxs.values())
            ordering_records.append(OrderingRecord(phase_name, phase_obj.model_hints['disordered_phase'], dof, equiv_sublattices))
    return ordering_records

def rename_disordered_phases(eq_result, ordering_records):
    """
    Modify an xarray Dataset to rename the ordered phase names to the disordered phase
    names if the equilibrium configuration is disordered

    Parameters
    ----------
    eq_result : xarray.Dataset
    order_disorder_dict : OrderingRecord
        Output from scheil.utils.order_disorder_dict

    Returns
    -------
    xrray.Dataset
        Dataset modified in-place

    Notes
    -----
    This function does _not_ change the site fractions array of the disordered
    configurations to match the site fractions matching the internal degrees of freedom
    of the disordered phase's constituents (although that should be possible).

    Examples
    --------
    >>> from pycalphad import Database, equilibrium, variables as v
    >>> dbf = Database('Al-Ni.tdb')
    >>> comps = ['AL', 'NI', 'VA']
    >>> phases = list(dbf.phases.keys())
    >>> eq_res = equilibrium(dbf, comps, phases, {v.P: 101325, v.T: (300, 2200, 10), v.N: 1, v.X('NI'): (0, 1, 0.01)})
    >>> ordering_records = create_ordering_records(dbf, comps, phases)
    >>> rename_disordered_phases(eq_res, ordering_records)
    """

    for ord_rec in ordering_records:
        # Array indices matching phase with ordered phase name
        mask = eq_result.Phase == ord_rec.ordered_phase_name
        # disordered_mask is a boolean mask which is True if the element listed
        # as an ordered phase is a disordered configuration. We want to
        # broadcast over all dimensions except for internal_dof (we need all
        # internal dof to justify whether the site fractions are disordered).
        # The OrderingRecord.is_disordered method is not vectorized and works
        # on one set of site fractions at a time, so we use `vectorize=True`
        disordered_mask = xr.apply_ufunc(ord_rec.is_disordered, eq_result.where(mask).Y, input_core_dims=[['internal_dof']], vectorize=True)
        # Finally, use xr.where to set the value of the phase name to the disordered phase everywhere the mask is true and use the existing value otherwise
        eq_result['Phase'] = xr.where(disordered_mask, ord_rec.disordered_phase_name, eq_result.Phase)
    return eq_result