Add Gradient Capability

espei/error_functions/activity_error.py

@@ -19,6 +19,7 @@
 from pycalphad.core.utils import filter_phases, unpack_species
 from scipy.stats import norm
 from pycalphad import Workspace
+from pycalphad.property_framework import JanssonDerivative
 
 from espei.core_utils import ravel_conditions
 from espei.error_functions.residual_base import ResidualFunction, residual_function_registry
@@ -29,7 +30,7 @@
 _log = logging.getLogger(__name__)
 
 
-def target_chempots_from_activity(component, target_activity, temperatures, wks_ref):
+def target_chempots_from_activity(component, parameters, target_activity, temperatures, wks_ref):
     """
     Return an array of experimental chemical potentials for the component
 
@@ -49,15 +50,21 @@ def target_chempots_from_activity(component, target_activity, temperatures, wks_
     numpy.ndarray
         Array of experimental chemical potentials
     """
-    # acr_i = exp((mu_i - mu_i^{ref})/(RT))
-    # so mu_i = R*T*ln(acr_i) + mu_i^{ref}
+
     ref_chempot = wks_ref.get(v.MU(component))
     exp_chem_pots = v.R * temperatures * np.log(target_activity) + ref_chempot
-    return exp_chem_pots
+
+    gradient_props = [JanssonDerivative(v.MU(component), key) for key in parameters]
+    gradients = wks_ref.get(*gradient_props)
+    if type(gradients) is list:
+        ref_grads = [float(element) for element in gradients]
+    else:
+        ref_grads = gradients
+    return exp_chem_pots, ref_grads
 
 
 # TODO: roll this function into ActivityResidual
-def calculate_activity_residuals(dbf, comps, phases, datasets, parameters=None, phase_models=None, callables=None, data_weight=1.0) -> Tuple[List[float], List[float]]:
+def calculate_activity_residuals(dbf, comps, phases, datasets, parameters=None, phase_models=None, callables=None, data_weight=1.0) -> Tuple[List[float], List[float], List[float]]:
     """
     Notes
     -----
@@ -75,12 +82,23 @@ def calculate_activity_residuals(dbf, comps, phases, datasets, parameters=None,
     if parameters is None:
         parameters = {}
 
+    params_keys = []
+
+    # This mutates the global pycalphad namespace
+    for key in parameters.keys():
+        if not hasattr(v, key):
+            setattr(v, key, v.IndependentPotential(key))
+        params_keys.append(getattr(v, key))
+        # Mutates argument to function
+        dbf.symbols.pop(key,None)
+
     activity_datasets = datasets.search(
         (tinydb.where('output').test(lambda x: 'ACR' in x)) &
         (tinydb.where('components').test(lambda x: set(x).issubset(comps))))
 
     residuals = []
     weights = []
+    gradients = []
     for ds in activity_datasets:
         acr_component = ds['output'].split('_')[1]  # the component of interest
         # calculate the reference state equilibrium
@@ -90,7 +108,9 @@ def calculate_activity_residuals(dbf, comps, phases, datasets, parameters=None,
         data_comps = ds['components']
         data_phases = filter_phases(dbf, unpack_species(dbf, data_comps), candidate_phases=phases)
         ref_conditions = {_map_coord_to_variable(coord): val for coord, val in ref['conditions'].items()}
-        wks_ref = Workspace(database=dbf, components=data_comps, phases=ref['phases'], conditions=ref_conditions, parameters=parameters)
+        # removed parameter assignment from wks_ref
+        ref_conditions.update(parameters)
+        wks_ref = Workspace(database=dbf, components=data_comps, phases=ref['phases'], conditions=ref_conditions)
 
         # calculate current chemical potentials
         # get the conditions
@@ -102,6 +122,7 @@ def calculate_activity_residuals(dbf, comps, phases, datasets, parameters=None,
         # we will ravel each composition individually, since they all must have the same shape
         dataset_computed_chempots = []
         dataset_weights = []
+        dataset_gradients = []
         for comp_name, comp_x in conds_list:
             P, T, X = ravel_conditions(ds['values'], ds['conditions']['P'], ds['conditions']['T'], comp_x)
             conditions[v.P] = P
@@ -113,22 +134,34 @@ def calculate_activity_residuals(dbf, comps, phases, datasets, parameters=None,
         # assume now that the ravelled conditions all have the same size
         conditions_list = [{c: conditions[c][i] for c in conditions.keys()} for i in range(len(conditions[v.T]))]
         for conds in conditions_list:
-            wks_sample = Workspace(database=dbf, components=data_comps, phases=data_phases, conditions=conds, parameters=parameters)
+            conds.update(parameters)
+            wks_sample = Workspace(database=dbf, components=data_comps, phases=data_phases, conditions=conds)
             dataset_computed_chempots.append(wks_sample.get(v.MU(acr_component)))
             dataset_weights.append(std_dev / data_weight / ds.get("weight", 1.0))
+            gradient_props = [JanssonDerivative(v.MU(acr_component), key) for key in parameters]
+            grads = wks_sample.get(*gradient_props)
+            if type(grads) is list:
+                sample_grads = [float(element) for element in grads]
+            else:
+                sample_grads = grads
+            dataset_gradients.append(sample_grads)
 
         # calculate target chempots
         dataset_activities = np.array(ds['values']).flatten()
-        dataset_target_chempots = target_chempots_from_activity(acr_component, dataset_activities, conditions[v.T], wks_ref)
+        dataset_target_chempots, ref_grads = target_chempots_from_activity(acr_component, parameters, dataset_activities, conditions[v.T], wks_ref)
         dataset_residuals = (np.asarray(dataset_computed_chempots) - np.asarray(dataset_target_chempots, dtype=float)).tolist()
+        adjusted_gradient = []
+        for element in dataset_gradients:
+            adjusted_gradient.append((np.asarray(element) - np.asarray(ref_grads)).tolist())
         _log.debug('Data: %s, chemical potential difference: %s, reference: %s', dataset_activities, dataset_residuals, ds["reference"])
         residuals.extend(dataset_residuals)
         weights.extend(dataset_weights)
-    return residuals, weights
+        gradients.append(adjusted_gradient)
+    return residuals, weights, gradients
 
 
 # TODO: roll this function into ActivityResidual
-def calculate_activity_error(dbf, comps, phases, datasets, parameters=None, phase_models=None, callables=None, data_weight=1.0) -> float:
+def calculate_activity_error(dbf, comps, phases, datasets, parameters=None, phase_models=None, callables=None, data_weight=1.0) -> Tuple[float, List[float]]:
     """
     Return the sum of square error from activity data
 
@@ -160,14 +193,23 @@ def calculate_activity_error(dbf, comps, phases, datasets, parameters=None, phas
 
 
     """
-    residuals, weights = calculate_activity_residuals(dbf, comps, phases, datasets, parameters=parameters, phase_models=phase_models, callables=callables, data_weight=data_weight)
+    residuals, weights, gradients = calculate_activity_residuals(dbf, comps, phases, datasets, parameters=parameters, phase_models=phase_models, callables=callables, data_weight=data_weight)
     likelihood = np.sum(norm(0, scale=weights).logpdf(residuals))
+    if len(gradients) == 0:
+        likelihood_grads = []
+    else: 
+        gradients = np.concatenate(gradients)
+        derivative = -np.array(residuals)*np.array(gradients).T/np.array(weights)**2
+        if derivative.ndim == 1:
+            likelihood_grads = np.sum(derivative, axis=0)
+        else:
+            likelihood_grads = np.sum(derivative, axis=1)
     if np.isnan(likelihood):
         # TODO: revisit this case and evaluate whether it is resonable for NaN
         # to show up here. When this comment was written, the test
         # test_subsystem_activity_probability would trigger a NaN.
-        return -np.inf
-    return likelihood
+        return -np.inf, np.zeros(len(parameters))
+    return likelihood, likelihood_grads
 
 
 # TODO: the __init__ method should pre-compute Model and PhaseRecord objects
@@ -215,13 +257,14 @@ def __init__(
 
     def get_residuals(self, parameters: npt.ArrayLike) -> Tuple[List[float], List[float]]:
         parameters = {param_name: param for param_name, param in zip(self._symbols_to_fit, parameters.tolist())}
-        residuals, weights = calculate_activity_residuals(parameters=parameters, **self._activity_likelihood_kwargs)
+        residuals, weights, grads = calculate_activity_residuals(parameters=parameters, **self._activity_likelihood_kwargs)
         return residuals, weights
 
-    def get_likelihood(self, parameters: npt.NDArray) -> float:
+    def get_likelihood(self, parameters: npt.NDArray) -> Tuple[float, List[float]]:
         parameters = {param_name: param for param_name, param in zip(self._symbols_to_fit, parameters.tolist())}
-        likelihood = calculate_activity_error(parameters=parameters, **self._activity_likelihood_kwargs)
-        return likelihood
+        likelihood, gradients = calculate_activity_error(parameters=parameters, **self._activity_likelihood_kwargs)
+        return likelihood, gradients
 
 
 residual_function_registry.register(ActivityResidual)
+

espei/error_functions/equilibrium_thermochemical_error.py

@@ -12,13 +12,13 @@
 from tinydb import where
 from scipy.stats import norm
 from pycalphad import Database, Model, ReferenceState, variables as v
-from pycalphad.core.utils import instantiate_models, filter_phases, extract_parameters, unpack_species, unpack_condition
+from pycalphad.core.utils import instantiate_models, filter_phases, unpack_species, unpack_condition
 from pycalphad.codegen.phase_record_factory import PhaseRecordFactory
 from pycalphad import Workspace, as_property
+from pycalphad.property_framework import JanssonDerivative
 
 from espei.error_functions.residual_base import ResidualFunction, residual_function_registry
 from espei.phase_models import PhaseModelSpecification
-from espei.shadow_functions import update_phase_record_parameters
 from espei.typing import SymbolName
 from espei.utils import PickleableTinyDB, database_symbols_to_fit
 
@@ -73,13 +73,21 @@ def build_eqpropdata(data: tinydb.database.Document,
         'SM':   0.2,  # J/K-mol
         'CPM':  0.2,  # J/K-mol
     }
+
+    params_keys = []
 
-    params_keys, _ = extract_parameters(parameters)
+    # This mutates the global pycalphad namespace
+    for key in parameters.keys():
+        if not hasattr(v, key):
+            setattr(v, key, v.IndependentPotential(key))
+        params_keys.append(getattr(v, key))
+        # Mutates argument to function
+        dbf.symbols.pop(key,None)
 
     data_comps = list(set(data['components']).union({'VA'}))
     species = sorted(unpack_species(dbf, data_comps), key=str)
     data_phases = filter_phases(dbf, species, candidate_phases=data['phases'])
-    models = instantiate_models(dbf, species, data_phases, model=model, parameters=parameters)
+    models = instantiate_models(dbf, species, data_phases, model=model)
     output = data['output']
     property_output = output.split('_')[0]  # property without _FORM, _MIX, etc.
     samples = np.array(data['values']).flatten()
@@ -99,14 +107,7 @@ def build_eqpropdata(data: tinydb.database.Document,
     pot_conds = OrderedDict([(getattr(v, key), unpack_condition(data['conditions'][key])) for key in sorted(data['conditions'].keys()) if not key.startswith('X_')])
     comp_conds = OrderedDict([(v.X(key[2:]), unpack_condition(data['conditions'][key])) for key in sorted(data['conditions'].keys()) if key.startswith('X_')])
 
-    phase_record_factory = PhaseRecordFactory(dbf, species, {**pot_conds, **comp_conds}, models, parameters=parameters)
-
-    # Now we need to unravel the composition conditions
-    # (from Dict[v.X, Sequence[float]] to Sequence[Dict[v.X, float]]), since the
-    # composition conditions are only broadcast against the potentials, not
-    # each other. Each individual composition needs to be computed
-    # independently, since broadcasting over composition cannot be turned off
-    # in pycalphad.
+    phase_record_factory = PhaseRecordFactory(dbf, species, {**pot_conds, **comp_conds, **parameters}, models)
     rav_comp_conds = [OrderedDict(zip(comp_conds.keys(), pt_comps)) for pt_comps in zip(*comp_conds.values())]
 
     # Build weights, should be the same size as the values
@@ -173,7 +174,7 @@ def get_equilibrium_thermochemical_data(dbf: Database, comps: Sequence[str],
 def calc_prop_differences(eqpropdata: EqPropData,
                           parameters: np.ndarray,
                           approximate_equilibrium: Optional[bool] = False,
-                          ) -> Tuple[np.ndarray, np.ndarray]:
+                          ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     """
     Calculate differences between the expected and calculated values for a property
 
@@ -196,42 +197,45 @@ def calc_prop_differences(eqpropdata: EqPropData,
         * weights for this dataset
 
     """
+
     dbf = eqpropdata.dbf
     species = eqpropdata.species
     phases = eqpropdata.phases
     pot_conds = eqpropdata.potential_conds
     models = eqpropdata.models
     phase_record_factory = eqpropdata.phase_record_factory
-    update_phase_record_parameters(phase_record_factory, parameters)
     params_dict = OrderedDict(zip(map(str, eqpropdata.params_keys), parameters))
     output = as_property(eqpropdata.output)
     weights = np.array(eqpropdata.weight, dtype=np.float64)
     samples = np.array(eqpropdata.samples, dtype=np.float64)
-    wks = Workspace(database=dbf, components=species, phases=phases, models=models, phase_record_factory=phase_record_factory, parameters=params_dict)
+    wks = Workspace(database=dbf, components=species, phases=phases, models=models, phase_record_factory=phase_record_factory)
 
     calculated_data = []
+    gradient_data = []
     for comp_conds in eqpropdata.composition_conds:
-        cond_dict = OrderedDict(**pot_conds, **comp_conds)
+        cond_dict = OrderedDict(**pot_conds, **comp_conds, **params_dict)
         wks.conditions = cond_dict
-        wks.parameters = params_dict  # these reset models and phase_record_factory through depends_on -> lose Model.shift_reference_state, etc.
         wks.models = models
-        wks.phase_record_factory = phase_record_factory
         vals = wks.get(output)
         calculated_data.extend(np.atleast_1d(vals).tolist())
+        gradient_props = [JanssonDerivative(output, key) for key in params_dict]
+        gradients = wks.get(*gradient_props)
+        gradient_data.append(gradients)
 
     calculated_data = np.array(calculated_data, dtype=np.float64)
+    gradient_data = np.array(gradient_data, dtype=np.float64)
 
     assert calculated_data.shape == samples.shape, f"Calculated data shape {calculated_data.shape} does not match samples shape {samples.shape}"
     assert calculated_data.shape == weights.shape, f"Calculated data shape {calculated_data.shape} does not match weights shape {weights.shape}"
     differences = calculated_data - samples
     _log.debug('Output: %s differences: %s, weights: %s, reference: %s', output, differences, weights, eqpropdata.reference)
-    return differences, weights
+    return differences, weights, gradient_data
 
 
 def calculate_equilibrium_thermochemical_probability(eq_thermochemical_data: Sequence[EqPropData],
                                                      parameters: np.ndarray,
                                                      approximate_equilibrium: Optional[bool] = False,
-                                                     ) -> float:
+                                                     ) -> Tuple[float, List[float]]:
     """
     Calculate the total equilibrium thermochemical probability for all EqPropData
 
@@ -252,23 +256,31 @@ def calculate_equilibrium_thermochemical_probability(eq_thermochemical_data: Seq
 
     """
     if len(eq_thermochemical_data) == 0:
-        return 0.0
+        return 0.0, np.zeros(len(parameters))
 
     differences = []
     weights = []
+    gradients = []
     for eqpropdata in eq_thermochemical_data:
-        diffs, wts = calc_prop_differences(eqpropdata, parameters, approximate_equilibrium)
+        diffs, wts, grads = calc_prop_differences(eqpropdata, parameters, approximate_equilibrium)
         if np.any(np.isinf(diffs) | np.isnan(diffs)):
             # NaN or infinity are assumed calculation failures. If we are
             # calculating log-probability, just bail out and return -infinity.
-            return -np.inf
+            return -np.inf, np.zeros(len(parameters))
         differences.append(diffs)
         weights.append(wts)
+        gradients.append(grads)
 
     differences = np.concatenate(differences, axis=0)
     weights = np.concatenate(weights, axis=0)
+    gradients = np.concatenate(gradients, axis=0)
     probs = norm(loc=0.0, scale=weights).logpdf(differences)
-    return np.sum(probs)
+    grad_probs = -differences*gradients.T/weights**2
+    if grad_probs.ndim == 1:
+        grad_probs_sum = np.sum(grad_probs, axis=0)
+    else:
+        grad_probs_sum = np.sum(grad_probs, axis=1)
+    return np.sum(probs), grad_probs_sum
 
 
 class EquilibriumPropertyResidual(ResidualFunction):
@@ -304,14 +316,14 @@ def get_residuals(self, parameters: npt.ArrayLike) -> Tuple[List[float], List[fl
         residuals = []
         weights = []
         for data in self.property_data:
-            dataset_residuals, dataset_weights = calc_prop_differences(data, parameters)
+            dataset_residuals, dataset_weights, dataset_grads = calc_prop_differences(data, parameters)
             residuals.extend(dataset_residuals.tolist())
             weights.extend(dataset_weights.tolist())
         return residuals, weights
 
-    def get_likelihood(self, parameters) -> float:
-        likelihood = calculate_equilibrium_thermochemical_probability(self.property_data, parameters)
-        return likelihood
+    def get_likelihood(self, parameters) -> Tuple[float, List[float]]:
+        likelihood, gradients = calculate_equilibrium_thermochemical_probability(self.property_data, parameters)
+        return likelihood, gradients
 
 
-residual_function_registry.register(EquilibriumPropertyResidual)
+residual_function_registry.register(EquilibriumPropertyResidual)