Merge pull request #10 from nury12n/workspace_compat

Pycalphad workspace compatibility. Some slight refactoring was done to take advantage of the Workspace variables over equilibrium calls.

.github/workflows/pytest.yaml

@@ -11,7 +11,7 @@ jobs:
       fail-fast: false
       max-parallel: 100
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
         pycalphad_develop_version: [true, false]
 
     steps:

kawin/diffusion/Homogenization.py

@@ -84,21 +84,10 @@ def _newEqCalc(self, x, T):
         '''
         Calculates equilibrium and returns a CompositionSet
         '''
-        eq = self.therm.getEq(x, T, 0, self.phases)
-        state_variables = np.array([0, 1, 101325, T], dtype=np.float64)
-        stable_phases = eq.Phase.values.ravel()
-        phase_amounts = eq.NP.values.ravel()
-        comp = []
-        for p in stable_phases:
-            if p != '':
-                idx = np.where(stable_phases == p)[0]
-                cs, misc = self.therm._createCompositionSet(eq, state_variables, p, phase_amounts, idx)
-                comp.append(cs)
-
-        if len(comp) == 0:
-            comp = None
-
-        return self.therm.getLocalEq(x, T, 0, self.phases, comp)
+        wks = self.therm.getEq(x, T, 0, self.phases)
+        chemical_potentials = np.squeeze(wks.eq.MU)
+        composition_sets = wks.get_composition_sets()
+        return chemical_potentials, composition_sets
 
     def updateCompSets(self, xarray):
         '''
@@ -127,18 +116,14 @@ def updateCompSets(self, xarray):
             if self.cache:
                 hashValue = self._getHash(xarray[:,i], self.T[i])
                 if hashValue not in self.hashTable:
-                    result, comp = self._newEqCalc(xarray[:,i], self.T[i])
-                    #result, comp = self.therm.getLocalEq(xarray[:,i], self.T, 0, self.phases, self.compSets[i])
-                    self.hashTable[hashValue] = (result, comp, None)
+                    chemical_potentials, comp = self._newEqCalc(xarray[:,i], self.T[i])
+                    self.hashTable[hashValue] = (chemical_potentials, comp, None)
                 else:
-                    result, comp, _ = self.hashTable[hashValue]
-                results, self.compSets[i] = copy.copy(result), copy.copy(comp)
+                    chemical_potentials, comp, _ = self.hashTable[hashValue]
+                chemical_potentials, self.compSets[i] = copy.copy(chemical_potentials), copy.copy(comp)
             else:
-                if self.compSets[i] is None:
-                    results, self.compSets[i] = self._newEqCalc(xarray[:,i], self.T[i])
-                else:
-                    results, self.compSets[i] = self.therm.getLocalEq(xarray[:,i], self.T[i], 0, self.phases, self.compSets[i])
-            self.mu[:,i] = results.chemical_potentials[self.unsortIndices]
+                chemical_potentials, self.compSets[i] = self._newEqCalc(xarray[:,i], self.T[i])
+            self.mu[:,i] = chemical_potentials[self.unsortIndices]
             cs_phases = [cs.phase_record.phase_name for cs in self.compSets[i]]
             for p in range(len(cs_phases)):
                 parray[self._getPhaseIndex(cs_phases[p]), i] = self.compSets[i][p].NP

kawin/tests/test_diffusion.py

@@ -309,18 +309,18 @@ def test_homogenization_dxdt():
     dt = m.getDt(dxdt)
 
     #Index 5
-    ind5, vals5 = 5, np.array([-1.592463e-9, 1.211067e-9])
+    ind5, vals5 = 5, np.array([-1.581478e-9, 1.212876e-9])
 
     #Index 10
-    ind10, vals10 = 10, np.array([-9.751858e-10, 1.702190e-9])
+    ind10, vals10 = 10, np.array([-9.722631e-10, 1.703447e-9])
 
     #Index 15
-    ind15, vals15 = 15, np.array([-4.728854e-10, 8.590127e-10])
+    ind15, vals15 = 15, np.array([-4.720562e-10, 8.600518e-10])
 
     assert_allclose(dxdt[0][:,ind5], vals5, atol=0, rtol=1e-3)
     assert_allclose(dxdt[0][:,ind10], vals10, atol=0, rtol=1e-3)
     assert_allclose(dxdt[0][:,ind15], vals15, atol=0, rtol=1e-3)
-    assert_allclose(dt, 61865.352193, rtol=1e-3)
+    assert_allclose(dt, 62271.050081, rtol=1e-3)
 
 
 

kawin/thermo/BinTherm.py

@@ -1,6 +1,6 @@
 from kawin.thermo.Thermodynamics import GeneralThermodynamics
 import numpy as np
-from pycalphad import equilibrium, calculate, variables as v
+from pycalphad import Workspace, equilibrium, calculate, variables as v
 from pycalphad.core.composition_set import CompositionSet
 from kawin.thermo.FreeEnergyHessian import dMudX
 
@@ -160,50 +160,46 @@ def _interfacialCompositionFromEq(self, T, gExtra = 0, precPhase = None):
 
         #Compute equilibrium at guess composition
         cond = {v.X(self.elements[1]): self._guessComposition[precPhase], v.T: T, v.P: 101325, v.GE: gExtra}
-        eq = equilibrium(self.db, self.elements, [self.phases[0], precPhase], cond, model=self.models,
-                        phase_records={self.phases[0]: self.phase_records[self.phases[0]], precPhase: self.phase_records[precPhase]},
-                        calc_opts = {'pdens': self.pDens})
-
-        xParentArray = np.zeros(len(gExtra))
-        xPrecArray = np.zeros(len(gExtra))
-        for g in range(len(gExtra)):
-            eqG = eq.where(eq.GE == gExtra[g], drop=True)
-            gm = eqG.GM.values.ravel()
-            for i in range(len(gm)):
-                eqSub = eqG.where(eqG.GM == gm[i], drop=True)
-
-                ph = eqSub.Phase.values.ravel()
-                ph = ph[ph != '']
-
-                #Check if matrix and precipitate phase are stable, and check if there's no miscibility gaps
+        phases, sub_models = self._setupSubModels(precPhase)
+        wks = Workspace(self.db, self.elements, phases, cond, models=sub_models,
+                        phase_record_factory = self.phase_records, calc_opts={'pdens': self.pDens})
+
+        coord_keys = list(wks.eq.coords.keys())
+        ge_var_idx = coord_keys.index('GE')
+
+        xMatrixArray = -1*np.ones(len(gExtra), dtype=np.float64)
+        xPrecipArray = -1*np.ones(len(gExtra), dtype=np.float64)
+        gIndex = 0
+        # cs_idx is a 5 len tuple of (GE, N, P, T, X)
+        # where it iterates by X first, then by GE
+        for cs_idx, cs_list in wks.enumerate_composition_sets():
+            # If the GE index is greater than the current record index, then
+            # update the current index. This occurs if we did not find any
+            # two-phase regions at a given value of GE and now we want to move
+            # on to the next value of GE
+            if cs_idx[ge_var_idx] > gIndex:
+                gIndex = cs_idx[ge_var_idx]
+
+            # Only check for two-phase region if the cs list corresponds
+            # to the current index of GE that we're looking at
+            # If the current index of GE is greater than what's in
+            # cs_idx, then we want to iterate the composition sets until we
+            # catch up to gIndex
+            if cs_idx[ge_var_idx] == gIndex:
+                ph = [cs.phase_record.phase_name for cs in cs_list]
                 if len(ph) == 2 and self.phases[0] in ph and precPhase in ph:
-                    #Get indices for each phase
-                    eqPa = eqSub.where(eqSub.Phase == self.phases[0], drop=True)
-                    eqPr = eqSub.where(eqSub.Phase == precPhase, drop=True)
-
-                    cParent = eqPa.X.values.ravel()
-                    cPrec = eqPr.X.values.ravel()
-
-                    #Get composition of element, use element index of 1 is the parent index is first alphabetically
-                    if self.reverse:
-                        xParent = cParent[0]
-                        xPrec = cPrec[0]
-                    else:
-                        xParent = cParent[1]
-                        xPrec = cPrec[1]
-
-                    xParentArray[g] = xParent
-                    xPrecArray[g] = xPrec
-                    break
-            if xParentArray[g] == 0:
-                xParentArray[g] = -1
-                xPrecArray[g] = -1
+                    cs_matrix = [cs for cs in cs_list if cs.phase_record.phase_name == self.phases[0]][0]
+                    cs_precip = [cs for cs in cs_list if cs.phase_record.phase_name == precPhase][0]
+
+                    c_idx = 0 if self.reverse else 1
+                    xMatrixArray[gIndex] = cs_matrix.X[c_idx]
+                    xPrecipArray[gIndex] = cs_precip.X[c_idx]
+                    gIndex += 1
 
         if len(gExtra) == 1:
-            return xParentArray[0], xPrecArray[0]
+            return xMatrixArray[0], xPrecipArray[0]
         else:
-            return xParentArray, xPrecArray
-
+            return xMatrixArray, xPrecipArray
 
     def _interfacialCompositionFromCurvature(self, T, gExtra = 0, precPhase = None):
         '''
@@ -236,83 +232,54 @@ def _interfacialCompositionFromCurvature(self, T, gExtra = 0, precPhase = None):
             gExtra = np.array([gExtra])
 
         #Compute equilibrium at guess composition
-        cond = {v.X(self.elements[1]): self._guessComposition[precPhase], v.T: T, v.P: 101325, v.GE: self.gOffset}
-        eq = equilibrium(self.db, self.elements, [self.phases[0], precPhase], cond, model=self.models,
-                        phase_records={self.phases[0]: self.phase_records[self.phases[0]], precPhase: self.phase_records[precPhase]},
-                        calc_opts = {'pdens': self.pDens})
-
-        gm = eq.GM.values.ravel()
-        for g in gm:
-            eqSub = eq.where(eq.GM == g, drop=True)
-
-            ph = eqSub.Phase.values.ravel()
-            ph = ph[ph != '']
-
-            #Check if matrix and precipitate phase are stable, and check if there's no miscibility gaps
+        cond = {v.X(self.elements[1]): self._guessComposition[precPhase], v.T: T, v.P: 101325, v.GE: self.gOffset, v.N: 1}
+        phases, sub_models = self._setupSubModels(precPhase)
+        wks = Workspace(self.db, self.elements, phases, cond, models=sub_models,
+                        phase_record_factory = self.phase_records, calc_opts={'pdens': self.pDens})
+
+        chemical_potentials = np.squeeze(wks.eq.MU)
+
+        idx = 0
+        for _, cs_list in wks.enumerate_composition_sets():
+            ph = [cs.phase_record.phase_name for cs in cs_list]
             if len(ph) == 2 and self.phases[0] in ph and precPhase in ph:
-                #Cast values in state_variables to double for updating composition sets
-                state_variables = np.array([cond[v.GE], cond[v.N], cond[v.P], cond[v.T]], dtype=np.float64)
-                stable_phases = eqSub.Phase.values.ravel()
-                phase_amounts = eqSub.NP.values.ravel()
-                matrix_idx = np.where(stable_phases == self.phases[0])[0]
-                precip_idx = np.where(stable_phases == precPhase)[0]
-
-                cs_matrix = CompositionSet(self.phase_records[self.phases[0]])
-                if len(matrix_idx) > 1:
-                    matrix_idx = [matrix_idx[np.argmax(phase_amounts[matrix_idx])]]
-                cs_matrix.update(eqSub.Y.isel(vertex=matrix_idx).values.ravel()[:cs_matrix.phase_record.phase_dof],
-                                    phase_amounts[matrix_idx], state_variables)
-                cs_precip = CompositionSet(self.phase_records[precPhase])
-                if len(precip_idx) > 1:
-                    precip_idx = [precip_idx[np.argmax(phase_amounts[precip_idx])]]
-                cs_precip.update(eqSub.Y.isel(vertex=precip_idx).values.ravel()[:cs_precip.phase_record.phase_dof],
-                                    phase_amounts[precip_idx], state_variables)
-
-                chemical_potentials = eqSub.MU.values.ravel()
-                cPrec = eqSub.isel(vertex=precip_idx).X.values.ravel()
-                cParent = eqSub.isel(vertex=matrix_idx).X.values.ravel()
-
-                dMudxParent = dMudX(chemical_potentials, cs_matrix, self.elements[0])
-                dMudxPrec = dMudX(chemical_potentials, cs_precip, self.elements[0])
-
-                #Get composition of element, use element index of 1 is the parent index is first alphabetically
-                if self.reverse:
-                    xParentEq = cParent[0]
-                    xPrecEq = cPrec[0]
-                else:
-                    xParentEq = cParent[1]
-                    xPrecEq = cPrec[1]
+                cs_matrix = [cs for cs in cs_list if cs.phase_record.phase_name == self.phases[0]][0]
+                cs_precip = [cs for cs in cs_list if cs.phase_record.phase_name == precPhase][0]
+                chem_pot = chemical_potentials[idx]
 
-                #dmudx are scalars here
-                dMudxParent = dMudxParent[0,0]
-                dMudxPrec = dMudxPrec[0,0]
+                dMudxMatrix = dMudX(chem_pot, cs_matrix, self.elements[0])[0,0]
+                dMudxPrecip = dMudX(chem_pot, cs_precip, self.elements[0])[0,0]
 
-                if dMudxParent != 0:
-                    xParent = gExtra / dMudxParent / (xPrecEq - xParentEq) + xParentEq
+                c_idx = 0 if self.reverse else 1
+                xMatrixEq = cs_matrix.X[c_idx]
+                xPrecipEq = cs_precip.X[c_idx]
+
+                if dMudxMatrix != 0:
+                    xMatrix = gExtra / dMudxMatrix / (xPrecipEq - xMatrixEq) + xMatrixEq
                 else:
-                    xParent = xParentEq*np.ones(len(gExtra))
+                    xMatrix = xMatrixEq*np.ones(len(gExtra))
 
-                if dMudxPrec != 0:
-                    xPrec = dMudxParent * (xParent - xParentEq) / dMudxPrec + xPrecEq
+                if dMudxPrecip != 0:
+                    xPrecip = dMudxMatrix * (xMatrix - xMatrixEq) / dMudxPrecip + xPrecipEq
                 else:
-                    xPrec = xPrecEq*np.ones(len(gExtra))
+                    xPrecip = xPrecipEq*np.ones(len(gExtra))
 
-                xParent[xParent < 0] = 0
-                xParent[xParent > 1] = 1
-                xPrec[xPrec < 0] = 0
-                xPrec[xPrec > 1] = 1
+                xMatrix[xMatrix < 0] = 0
+                xMatrix[xMatrix > 1] = 1
+                xPrecip[xPrecip < 0] = 0
+                xPrecip[xPrecip > 1] = 1
 
                 if len(gExtra) == 1:
-                    return xParent[0], xPrec[0]
+                    return xMatrix[0], xPrecip[0]
                 else:
-                    return xParent, xPrec
+                    return xMatrix, xPrecip
+            idx += 1
 
         if len(gExtra) == 1:
             return -1, -1
         else:
             return -1*np.ones(len(gExtra)), -1*np.ones(len(gExtra))
 
-
     def plotPhases(self, ax, T, gExtra = 0, plotGibbsOffset = False, *args, **kwargs):
         '''
         Plots sampled points from the parent and precipitate phase
@@ -331,12 +298,14 @@ def plotPhases(self, ax, T, gExtra = 0, plotGibbsOffset = False, *args, **kwargs
                 with no extra Gibbs free energy contributions
             Defualts to False
         '''
-        points = calculate(self.db, self.elements, self.phases[0], P=101325, T=T, GE=0, model=self.models, phase_records=self.phase_records, output='GM')
+        phases, sub_models = self._setupSubModels([self.phases[0]])
+        points = calculate(self.db, self.elements, phases[0], P=101325, T=T, GE=0, model=sub_models, phase_records=self.phase_records, output='GM')
         ax.scatter(points.X.sel(component=self.elements[1]), points.GM / 1000, label=self.phases[0], *args, **kwargs)
 
         #Add gExtra to precipitate phase
         for i in range(1, len(self.phases)):
-            points = calculate(self.db, self.elements, self.phases[i], P=101325, T=T, GE=0, model=self.models, phase_records=self.phase_records, output='GM')
+            phases, sub_models = self._setupSubModels([self.phases[i]])
+            points = calculate(self.db, self.elements, phases[0], P=101325, T=T, GE=0, model=sub_models, phase_records=self.phase_records, output='GM')
             ax.scatter(points.X.sel(component=self.elements[1]), (points.GM + gExtra) / 1000, label=self.phases[i], *args, **kwargs)
 
             #Plot non-offset precipitate phase

kawin/thermo/FreeEnergyHessian.py

@@ -1,8 +1,4 @@
 import numpy as np
-from pycalphad import Model, variables as v
-from pycalphad.codegen.callables import build_callables
-
-setattr(v, 'GE', v.StateVariable('GE'))
 
 def hessian(chemical_potentials, composition_set):
     '''

kawin/thermo/LocalEquilibrium.py

@@ -1,6 +1,5 @@
 from pycalphad.core.solver import Solver
 from pycalphad.core.composition_set import CompositionSet
-from pycalphad.codegen.callables import build_phase_records
 from pycalphad import calculate, variables as v
 import numpy as np
 
@@ -28,10 +27,13 @@ def local_equilibrium(dbf, comps, phases, conds, models, phase_records, composit
     '''
     # Broadcasting conditions not supported
     cur_conds = {str(k): float(v) for k, v in conds.items()}
-    if 'GE' in cur_conds:
-        state_variables = np.array([cur_conds['GE'], cur_conds['N'], cur_conds['P'], cur_conds['T']], dtype=np.float64)
+    cur_conds = {k: v for k, v in conds.items()}
+
+    # State variables are in alphabetical order
+    if v.GE in cur_conds:
+        state_variables = np.array([cur_conds[v.GE], cur_conds[v.N], cur_conds[v.P], cur_conds[v.T]], dtype=np.float64)
     else:
-        state_variables = np.array([0, cur_conds['N'], cur_conds['P'], cur_conds['T']], dtype=np.float64)
+        state_variables = np.array([0, cur_conds[v.N], cur_conds[v.P], cur_conds[v.T]], dtype=np.float64)
     if composition_sets is None:
         # Note: filter_phases() not called, so all specified phases must be valid
         composition_sets = []
@@ -43,7 +45,7 @@ def local_equilibrium(dbf, comps, phases, conds, models, phase_records, composit
         for phase in phases:
             # arbitrary guess
             phase_amt = 1./len(phases)
-            calc_p = calculate(dbf, comps, phase, T=cur_conds['T'], P=cur_conds['P'], N=cur_conds['N'], GE=cur_conds['GE'],
+            calc_p = calculate(dbf, comps, phase, T=cur_conds[v.T], P=cur_conds[v.P], N=cur_conds[v.N], GE=cur_conds[v.GE],
                                pdens=10, model=models, phase_records=phase_records)
             idx_p = np.argmin(calc_p.GM.values.squeeze())
             compset = CompositionSet(phase_records[phase])