Source code for pybamm.expression_tree.input_parameter

#
# Parameter classes
#
from __future__ import annotations

import numbers

import casadi
import numpy as np
import numpy.typing as npt
import scipy.sparse

import pybamm
from pybamm.type_definitions import DomainType

DUMMY_INPUT_PARAMETER_VALUE = np.nan


class InputParameter(pybamm.Symbol):
    """
    A node in the expression tree representing an input parameter.

    This node's value can be set at the point of solving, allowing parameter estimation
    and control

    Parameters
    ----------
    name : str
        name of the node
    domain : iterable of str, or str
        list of domains over which the node is valid (empty list indicates the symbol
        is valid over all domains)
    expected_size : int
        The size of the input parameter expected, defaults to 1 (scalar input)
    """

    __slots__ = ("_expected_size",)

    def __init__(
        self,
        name: str,
        domain: DomainType = None,
        expected_size: int | None = None,
    ) -> None:
        # Expected size defaults to 1 if no domain else None (gets set later)
        if expected_size is None:
            if domain is None:
                expected_size = 1
            else:
                expected_size = None
        self._expected_size = expected_size
        super().__init__(name, domain=domain)

    @classmethod
    def _from_json(cls, snippet: dict):
        return cls(
            snippet["name"],
            domain=snippet["domain"],
            expected_size=snippet["expected_size"],
        )

    def create_copy(
        self,
        new_children=None,
        perform_simplifications=True,
    ) -> pybamm.InputParameter:
        """See :meth:`pybamm.Symbol.new_copy()`."""
        new_input_parameter = InputParameter(
            self.name, self._domains["primary"], expected_size=self._expected_size
        )
        return new_input_parameter

    def _evaluate_for_shape(self):
        """
        Returns the scalar 'NaN' to represent the shape of a parameter.
        See :meth:`pybamm.Symbol.evaluate_for_shape()`
        """
        if self._expected_size is None:
            return pybamm.evaluate_for_shape_using_domain(self._domains)
        elif self._expected_size == 1:
            return np.nan
        else:
            return np.nan * np.ones((self._expected_size, 1))

    def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols):
        """See :meth:`pybamm.Symbol._to_casadi()`."""
        return casadi.MX(self.evaluate(t, y, y_dot, inputs))

    def _jac(self, variable: pybamm.StateVector) -> pybamm.Matrix:
        """See :meth:`pybamm.Symbol._jac()`."""
        n_variable = variable.evaluation_array.count(True)
        nan_vector = self._evaluate_for_shape()
        if isinstance(nan_vector, numbers.Number):
            n_self = 1
        else:
            n_self = nan_vector.shape[0]
        zero_matrix = scipy.sparse.csr_matrix((n_self, n_variable))
        return pybamm.Matrix(zero_matrix)

    def _base_evaluate(
        self,
        t: float | None = None,
        y: npt.NDArray[np.float64] | None = None,
        y_dot: npt.NDArray[np.float64] | None = None,
        inputs: dict | str | None = None,
    ):
        # inputs should be a dictionary
        # convert 'None' to empty dictionary for more informative error
        if inputs is None:
            inputs = {}
        if not isinstance(inputs, dict):
            # if the special input "shape test" is passed, just return NaN
            if inputs == "shape test":
                return self.evaluate_for_shape()
            raise TypeError("inputs should be a dictionary")
        try:
            input_eval = inputs[self.name]
        # raise more informative error if can't find name in dict
        except KeyError as error:
            raise KeyError(f"Input parameter '{self.name}' not found") from error

        if isinstance(input_eval, numbers.Number):
            input_size = 1
            input_ndim = 0
        else:
            input_size = input_eval.shape[0]
            input_ndim = len(input_eval.shape)
        if input_size == self._expected_size:
            if input_ndim == 1:
                return input_eval[:, np.newaxis]
            else:
                return input_eval
        else:
            raise ValueError(
                f"Input parameter '{self.name}' was given an object of size '{input_size}'"
                + f" but was expecting an object of size '{self._expected_size}'."
            )

    def to_json(self):
        """
        Method to serialise an InputParameter object into JSON.
        """

        json_dict = {
            "name": self.name,
            "domain": self._domains["primary"],
            "expected_size": self._expected_size,
        }

        return json_dict