API reference

PyCO2SYS

Marine carbonate system calculations in Python.

CO2System

Bases: UserDict

Source code in PyCO2SYS/engine.py

class CO2System(UserDict):
    def __init__(self, pd_index=None, xr_dims=None, xr_shape=None, **kwargs):
        super().__init__()
        opts = {k: v for k, v in kwargs.items() if k.startswith("opt_")}
        data = {k: v for k, v in kwargs.items() if k not in opts}
        # Get icase
        core_known = np.array([v in data for v in parameters_core])
        icase_all = np.arange(1, len(parameters_core) + 1)
        icase = icase_all[core_known]
        assert len(icase) < 3, "A maximum of 2 known core parameters can be provided."
        if len(icase) == 0:
            icase = np.array(0)
        elif len(icase) == 2:
            icase = icase[0] * 100 + icase[1]
        self.icase = icase.item()
        self.opts = opts_default.copy()
        # Assign opts
        for k, v in opts.items():
            if k in get_funcs_opts:
                assert np.isscalar(v)
                assert v in get_funcs_opts[k].keys(), (
                    "{} is not allowed for {}!".format(v, k)
                )
            else:
                warnings.warn(
                    "'{}' is not recognised".format(k)
                    + " - it will not be used in any calculations."
                )
        self.opts.update(opts)
        # Deal with tricky special cases
        if self.icase != 207:
            self.opts.pop("opt_HCO3_root")
        if self.icase not in [0, 4, 5, 8, 9]:
            self.opts.pop("opt_fCO2_temperature")
        # Assemble graphs and computation functions
        self.graph, self.funcs, self.data = self._assemble(self.icase, data)
        for k in self.data:
            if k not in self.graph.nodes:
                warnings.warn(
                    "'{}' is not recognised".format(k)
                    + " - it will not be used in any calculations."
                )
        self.grads = {}
        self.uncertainty = {}
        self.requested = set()  # keep track of all parameters that have been requested
        self.pd_index = pd_index
        if xr_dims is not None:
            assert xr_shape is not None
            assert len(xr_dims) == len(xr_shape)
        else:
            assert xr_shape is None
        self.xr_dims = xr_dims
        self.xr_shape = xr_shape

    def __getitem__(self, key):
        # When the user requests a dict key that hasn't been solved for yet, then solve
        # and provide the requested parameter
        self.solve(parameters=key)
        if isinstance(key, list):
            # If the user provides a list of keys to solve for, return all of them as
            # a dict
            return {k: self.data[k] for k in key}
        else:
            # If a single key is requested, return the corresponding value(s) directly
            return self.data[key]

    def __getattr__(self, attr):
        # This allows solved parameter values to be accessed with dot notation, purely
        # for convenience.
        # So, when the user tries to access something with dot notation...
        try:
            # ... then if it's an attribute, return it (this is the standard behaviour).
            return object.__getattribute__(self, attr)
        except AttributeError:
            # But if it's not an attribute...
            try:
                # ... return the corresponding parameter value, if it's already been
                # solved for...
                return self.data[attr]
            except KeyError:
                # ... but it if hasn't been solved for, throw an error.  The user needs
                # to use the normal dict notation (or solve method) to solve for it.
                raise AttributeError(attr)

    def __setitem__(self, key, value):
        # Don't allow the user to assign new key-value pairs to the dict
        raise RuntimeError("Item assignment is not supported.")

    def _assemble(self, icase, data):
        # Deal with tricky special cases
        if icase == 207:
            graph_opts = get_graph_opts()
        else:
            graph_opts = get_graph_opts(exclude="opt_HCO3_root")
        # Assemble graph and functions
        graph = nx.compose(graph_fixed, graph_core[icase])
        funcs = get_funcs.copy()
        funcs.update(get_funcs_core[icase])
        for opt, v in self.opts.items():
            graph = nx.compose(graph, graph_opts[opt][v])
            funcs.update(get_funcs_opts[opt][v])
        # Assign default values
        data = data.copy()
        for k, v in values_default.items():
            if k not in data:
                data[k] = v
                graph.add_node(k)
        # Save arguments
        for k, v in data.items():
            if v is not None:
                # state 1 means that the value was provided as an argument
                nx.set_node_attributes(graph, {k: 1}, name="state")
        self.nodes_original = list(k for k, v in data.items() if v is not None)
        return graph, funcs, data

    def solve(self, parameters=None, store_steps=1):
        """Calculate parameter(s) and store them internally.

        Parameters
        ----------
        parameters : str or list of str, optional
            Which parameter(s) to calculate and store, by default None, in which case
            all possible parameters are calculated and stored internally.
        store_steps : int, optional
            Whether/which non-requested parameters calculated during intermediate
            calculation steps should be stored, by default 1.  The options are
                0 - store only the specifically requested parameters,
                1 - store the most used set of intermediate parameters, or
                2 - store the complete set of parameters.
        """
        # Parse user-provided parameters (if there are any)
        if parameters is None:
            # If no parameters are provided, then we solve for everything possible
            parameters = list(self.graph.nodes)
        elif isinstance(parameters, str):
            # Allow user to provide a string if only one parameter is desired
            parameters = [parameters]
        parameters = set(parameters)  # get rid of duplicates
        self.requested |= parameters
        self_data = self.data.copy()  # what was already known before this solve
        # Remove known nodes from a copy of self.graph, so that ancestors of known nodes
        # are not unnecessarily recomputed
        graph_unknown = self.graph.copy()
        graph_unknown.remove_nodes_from([k for k in self_data if k not in parameters])
        # Add intermediate parameters that we need to know in order to calculate the
        # requested parameters
        parameters_all = parameters.copy()
        for p in parameters:
            parameters_all = parameters_all | nx.ancestors(graph_unknown, p)
        # Convert the set of parameters into a list, exclude already-known ones, and
        # organise the list into the order required for calculations
        parameters_all = [
            p
            for p in nx.topological_sort(self.graph)
            if p in parameters_all and p not in self_data
        ]
        store_parameters = []
        for p in parameters_all:
            priors = self.graph.pred[p]
            if len(priors) == 0 or all([r in self_data for r in priors]):
                self_data[p] = self.funcs[p](
                    *[
                        self_data[r]
                        for r in self.funcs[p].__code__.co_varnames[
                            : self.funcs[p].__code__.co_argcount
                        ]
                    ]
                )
                store_here = (
                    #  If store_steps is 0, store only requested parameters
                    (store_steps == 0 and p in parameters)
                    | (
                        # If store_steps is 1, store all but the equilibrium constants
                        # on the seawater scale, at 1 atm and their pressure-correction
                        # factors, and a few selected others
                        store_steps == 1
                        and not p.startswith("factor_k_")
                        and not (p.startswith("k_") and p.endswith("_sws"))
                        and not p.endswith("_1atm")
                        and p not in ["sws_to_opt", "opt_to_free", "ionic_strength"]
                    )
                    |  # If store_steps is 2, store everything
                    (store_steps == 2)
                )
                if store_here:
                    store_parameters.append(p)
                    # state = 2 means that the value was calculated internally
                    nx.set_node_attributes(self.graph, {p: 2}, name="state")
                    for f in self.funcs[p].__code__.co_varnames[
                        : self.funcs[p].__code__.co_argcount
                    ]:
                        nx.set_edge_attributes(self.graph, {(f, p): 2}, name="state")
        # Get rid of jax overhead on results
        self_data = {k: v for k, v in self_data.items() if k in store_parameters}
        _remove_jax_overhead(self_data)
        self.data.update(self_data)

    def to_pandas(self, parameters=None, store_steps=1):
        """Return parameters as a pandas `Series` or `DataFrame`.  All parameters should
        be scalar or one-dimensional vectors of the same size.

        Parameters
        ----------
        parameters : str or list of str, optional
            The parameter(s) to return.  These are solved for if not already available.
            If `None`, then all parameters that have already been solved for are
            returned.
        store_steps : int, optional
            See `solve`.

        Returns
        -------
        pd.Series or pd.DataFrame
            The parameter(s) as a `pd.Series` (if `parameters` is a `str`) or as a
            `pd.DataFrame` (if `parameters` is a `list`) with the original pandas index
            passed into the `CO2System` as `data`.  If `data` was not a `pd.DataFrame`
            then the default index will be used.
        """
        try:
            import pandas as pd

            if parameters is None:
                parameters = self.keys()
            self.solve(parameters=parameters, store_steps=store_steps)
            if isinstance(parameters, str):
                return pd.Series(data=self[parameters], index=self.pd_index)
            else:
                return pd.DataFrame(
                    {
                        p: pd.Series(
                            data=self[p] * np.ones(self.pd_index.shape),
                            index=self.pd_index,
                        )
                        for p in parameters
                    }
                )
        except ImportError:
            warnings.warn("pandas could not be imported.")

    def _get_xr_ndims(self, parameter):
        ndims = []
        if not np.isscalar(self[parameter]):
            for i, vs in enumerate(self[parameter].shape):
                if vs == self.xr_shape[i]:
                    ndims.append(self.xr_dims[i])
        return ndims

    def to_xarray(self, parameters=None, store_steps=1):
        """Return parameters as an xarray `DataArray` or `Dataset`.

        Parameters
        ----------
        parameters : str or list of str, optional
            The parameter(s) to return.  These are solved for if not already available.
            If `None`, then all parameters that have already been solved for are
            returned.
        store_steps : int, optional
            See `solve`.

        Returns
        -------
        xr.DataArray or xr.Dataset
            The parameter(s) as a `xr.DataArray` (if `parameters` is a `str`) or as a
            `xr.Dataset` (if `parameters` is a `list`) with the original xarray
            dimensions passed into the `CO2System` as `data`.  If `data` was not an
            `xr.Dataset` then this function will not work.
        """
        assert self.xr_dims is not None and self.xr_shape is not None, (
            "`data` was not provided as an `xr.Dataset` "
            + "when creating this `CO2System`."
        )
        try:
            import xarray as xr

            if parameters is None:
                parameters = self.keys()
            self.solve(parameters=parameters, store_steps=store_steps)
            if isinstance(parameters, str):
                ndims = self._get_xr_ndims(parameters)
                return xr.DataArray(np.squeeze(self[parameters]), dims=ndims)
            else:
                return xr.Dataset(
                    {
                        p: xr.DataArray(np.squeeze(self[p]), dims=self._get_xr_ndims(p))
                        for p in parameters
                    }
                )
        except ImportError:
            warnings.warn("xarray could not be imported.")

    def _get_expUps(
        self,
        method_fCO2,
        temperature,
        bh_upsilon=None,
        opt_which_fCO2_insitu=1,
    ):
        if method_fCO2 in [1, 2, 3, 4]:
            self.solve("gas_constant")
        match method_fCO2:
            case 1:
                self.solve("fCO2", store_steps=0)
                fCO2 = self.fCO2
                assert opt_which_fCO2_insitu in [1, 2]
                if opt_which_fCO2_insitu == 2:
                    # If the output conditions are the environmental ones, then we need
                    # to provide an estimate of output fCO2 in order to use the bh
                    # parameterisation; we get this using the method_fCO2=2 approach:
                    fCO2 = fCO2 * upsilon.expUps_TOG93_H24(
                        self.data["temperature"],
                        temperature,
                        self.data["gas_constant"],
                    )
                return upsilon.expUps_parameterised_H24(
                    self.data["temperature"],
                    temperature,
                    self.data["salinity"],
                    fCO2,
                    self.data["gas_constant"],
                    opt_which_fCO2_insitu=opt_which_fCO2_insitu,
                )
            case 2:
                return upsilon.expUps_TOG93_H24(
                    self.data["temperature"],
                    temperature,
                    self.data["gas_constant"],
                )
            case 3:
                return upsilon.expUps_enthalpy_H24(
                    self.data["temperature"],
                    temperature,
                    self.data["gas_constant"],
                )
            case 4:
                assert bh_upsilon is not None, (
                    "A bh_upsilon value must be provided for method_fCO2=4."
                )
                return upsilon.expUps_Hoff_H24(
                    self.data["temperature"],
                    temperature,
                    self.data["gas_constant"],
                    bh_upsilon,
                )
            case 5:
                return upsilon.expUps_linear_TOG93(
                    self.data["temperature"],
                    temperature,
                )
            case 6:
                return upsilon.expUps_quadratic_TOG93(
                    self.data["temperature"],
                    temperature,
                )

    def adjust(
        self,
        temperature=None,
        pressure=None,
        store_steps=1,
        method_fCO2=1,
        opt_which_fCO2_insitu=1,
        bh_upsilon=None,
    ):
        """Adjust the system to a different temperature and/or pressure.

        Parameters
        ----------
        temperature : float, optional
            Temperature in °C to adjust to.  If `None`, temperature is not adjusted.
        pressure : float, optional
            Hydrostatic pressure in dbar to adjust to.  If `None`, pressure is not
            adjusted.
        store_steps : int, optional
            Whether/which non-requested parameters calculated during intermediate
            calculation steps should be stored.  The options are:

              - `0`: Store only the requested parameters.
              - `1`: Store the requested and most commonly used set of intermediate
              parameters (default).
              - `2`: Store the requested and complete set of intermediate parameters.
        method_fCO2 : int, optional
            If this is a single-parameter system, which method to use for the
            adjustment.  The options are:

              - `1`: parameterisation of H24 (default).
              - `2`: constant bh fitted to TOG93 dataset by H24.
              - `3`: constant theoretical bh of H24.
              - `4`: user-specified bh with the equations of H24.
              - `5`: linear fit of TOG93.
              - `6`: quadratic fit of TOG93.
        opt_which_fCO2_insitu : int, optional
            If this is a single-parameter system and `method_fCO2` is `1`, whether:
              - `1` the input condition (starting; default) or
              - `2` output condition (adjusted) temperature
            should be used to calculate $b_h$.
        bh_upsilon : float, optional
            If this is a single-parameter system and `method_fCO2` is `4`, then the
            value of $b_h$ in J/mol must be specified here.

        Returns
        -------
        CO2System
            A new `CO2System` with all values adjusted to the requested temperature(s)
            and/or pressure(s).
        """
        if self.icase > 100:
            # If we know (any) two MCS parameters, solve for alkalinity and DIC under
            # the "input" conditions
            self.solve(parameters=["alkalinity", "dic"], store_steps=store_steps)
            core = {k: self[k] for k in ["alkalinity", "dic"]}
            data = {}
        elif self.icase in [4, 5, 8, 9]:
            assert pressure is None, (
                "Cannot adjust pressure for a single-parameter system!"
            )
            # If we know only one of pCO2, fCO2, xCO2 or CO2(aq), first get fCO2 under
            # the "input" conditions
            self.solve(parameters="fCO2", store_steps=store_steps)
            core = {"fCO2": self.fCO2}
            # Then, convert this to the value at the new temperature using the requested
            # method
            assert method_fCO2 in range(1, 7), (
                "`method_fCO2` must be an integer from 1-6."
            )
            expUps = self._get_expUps(
                method_fCO2,
                temperature,
                bh_upsilon=bh_upsilon,
                opt_which_fCO2_insitu=opt_which_fCO2_insitu,
            )
            data = {"fCO2": core["fCO2"] * expUps}
        else:
            warnings.warn(
                "Single-parameter temperature adjustments are possible only "
                + "if the known parameter is one of pCO2, fCO2, xCO2 and CO2."
            )
        # Copy all parameters that are T/P independent into new data dict
        for k in condition_independent:
            if k in self.nodes_original:
                data[k] = self.data[k]
            elif k in core:
                data[k] = core[k]
        # Set temperature and/or pressure to adjusted value(s), unless they are None, in
        # which case, don't adjust
        if temperature is not None:
            data["temperature"] = temperature
        else:
            data["temperature"] = self.data["temperature"]
        if pressure is not None:
            data["pressure"] = pressure
        else:
            data["pressure"] = self.data["pressure"]
        sys = CO2System(
            **data,
            **self.opts,
            pd_index=self.pd_index,
            xr_dims=self.xr_dims,
            xr_shape=self.xr_shape,
        )
        sys.solve(parameters=self.data)
        return sys

    def _get_func_of(self, var_of):
        """Create a function to compute ``var_of`` directly from an input set of values.

        The created function has the signature

            value_of = get_value_of(**value)

        where the ``values`` are the originally user-defined values, obtained with
        either of the following:

            values_original = {k: sys.data[k] for k in sys.nodes_original}
            values_original = sys.get_values_original()
        """
        # We get a sub-graph of the node of interest and all its ancestors, excluding
        # originally fixed / user-defined values
        nodes_vo_all = nx.ancestors(self.graph, var_of)
        nodes_vo_all.add(var_of)
        # nodes_vo = onp.array([n for n in nodes_vo if n not in self.nodes_original])
        nodes_vo = [n for n in nodes_vo_all if n not in self.nodes_original]
        graph_vo = self.graph.subgraph(nodes_vo)

        def get_value_of(**data):
            data = data.copy()
            # This loops through the functions in the correct order determined above so
            # we end up calculating the value of interest, which is returned
            for n in nx.topological_sort(graph_vo):
                data.update(
                    {
                        n: self.funcs[n](
                            *[
                                data[v]
                                for v in self.funcs[n].__code__.co_varnames[
                                    : self.funcs[n].__code__.co_argcount
                                ]
                            ]
                        )
                    }
                )
            return data[var_of]

        # Generate docstring
        get_value_of.__doc__ = (
            "Calculate ``{}``.".format(var_of)
            + "\n\nParameters\n----------"
            + "\nvalues : dict"
            + "\n    Key-value pairs for the following parameters:"
        )
        for p in nodes_vo_all:
            if p in self.nodes_original:
                get_value_of.__doc__ += "\n        {}".format(p)
        get_value_of.__doc__ += "\n\nReturns\n-------"
        get_value_of.__doc__ += "\n{}".format(var_of)
        return get_value_of

    def _get_func_of_from_wrt(self, get_value_of, var_wrt):
        """Reorganise a function created with ``_get_func_of`` so that one of its kwargs
        is instead a positional arg (and which can thus be gradded).

        Parameters
        ----------
        get_value_of : func
            Function created with ``_get_func_of``.
        var_wrt : str
            Name of the value to use as a positional arg instead.

        Returns
        -------
        A function with the signature
            value_of = get_of_from_wrt(value_wrt, **other_values_original)
        """

        def get_value_of_from_wrt(value_wrt, **other_values_original):
            other_values_original = other_values_original.copy()
            other_values_original.update({var_wrt: value_wrt})
            return get_value_of(**other_values_original)

        return get_value_of_from_wrt

    def get_grad_func(self, var_of, var_wrt):
        get_value_of = self._get_func_of(var_of)
        get_value_of_from_wrt = self._get_func_of_from_wrt(get_value_of, var_wrt)
        return meta.egrad(get_value_of_from_wrt)

    def get_grad(self, var_of, var_wrt):
        """Compute the derivative of `var_of` with respect to `var_wrt` and store it in
        `sys.grads[var_of][var_wrt]`.  If there is already a value there, then that
        value is returned instead of recalculating.

        Parameters
        ----------
        var_of : str
            The name of the variable to get the derivative of.
        var_wrt : str
            The name of the variable to get the derivative with respect to.  This must
            be one of the fixed values provided when creating the `CO2System`, i.e.,
            listed in its `nodes_original` attribute.
        """
        assert var_wrt in self.nodes_original, (
            "`var_wrt` must be one of `sys.nodes_original!`"
        )
        try:  # see if we've already calculated this value
            d_of__d_wrt = self.grads[var_of][var_wrt]
        except KeyError:  # only do the calculations if there isn't already a value
            # We need to know the shape of the variable that we want the grad of, the
            # easy way to get this is just to solve for it (if that hasn't already been
            # done)
            if var_of not in self.data:
                self.solve(var_of)
            # Next, we extract the originally set values, which are fixed during the
            # differentiation
            values_original = self.get_values_original()
            other_values_original = values_original.copy()
            # We have to make sure the value we are differentiating with respect to has
            # the same shape as the value we want the differential of
            value_wrt = other_values_original.pop(var_wrt) * np.ones_like(
                self.data[var_of]
            )
            # Here we compute the gradient
            grad_func = self.get_grad_func(var_of, var_wrt)
            d_of__d_wrt = grad_func(value_wrt, **other_values_original)
            # Put the final value into self.grads, first creating a new sub-dict if
            # necessary
            if var_of not in self.grads:
                self.grads[var_of] = {}
            self.grads[var_of][var_wrt] = d_of__d_wrt
        return d_of__d_wrt

    def get_grads(self, vars_of, vars_wrt):
        """Compute the derivatives of `vars_of` with respect to `vars_wrt` and store
        them in `sys.grads[var_of][var_wrt]`.  If there are already values there, then
        those values are returned instead of recalculating.

        Parameters
        ----------
        vars_of : list
            The names of the variables to get the derivatives of.
        vars_wrt : list
            The names of the variables to get the derivatives with respect to.  These
            must all be one of the fixed values provided when creating the `CO2System`,
            i.e., listed in its `nodes_original` attribute.
        """
        if isinstance(vars_of, str):
            vars_of = [vars_of]
        if isinstance(vars_wrt, str):
            vars_wrt = [vars_wrt]
        for var_of, var_wrt in itertools.product(vars_of, vars_wrt):
            self.get_grad(var_of, var_wrt)

    def get_values_original(self):
        return {k: self.data[k] for k in self.nodes_original}

    def propagate(self, uncertainty_in, uncertainty_from):
        """Propagate independent uncertainties through the calculations.  Covariances
        are not accounted for.

        New entries are added in the `uncertainty` attribute, for example:

            co2s = CO2System(dic=2100, alkalinity=2300)
            co2s.propagate("pH", {"dic": 2, "alkalinity": 2})
            co2s.uncertainty["pH"]["total"]  # total uncertainty in pH
            co2s.uncertainty["pH"]["dic"]  # component of ^ due to DIC uncertainty

        Parameters
        ----------
        uncertainty_in : list
            The parameters to calculate the uncertainty in.
        uncertainty_from : dict
            The parameters to propagate the uncertainty from (keys) and their
            uncertainties (values).
        """
        self.solve(uncertainty_in)
        if isinstance(uncertainty_in, str):
            uncertainty_in = [uncertainty_in]
        for var_in in uncertainty_in:
            # This should always be reset to zero and all values wiped, even if it
            # already exists (so you don't end up with old uncertainty_from components
            # from a previous calculation which are no longer part of the total)
            self.uncertainty[var_in] = {"total": np.zeros_like(self.data[var_in])}
            u_total = self.uncertainty[var_in]["total"]
            for var_from, u_from in uncertainty_from.items():
                is_pk = var_from.startswith("pk_")
                if is_pk:
                    # If the uncertainty is given in terms of a pK value, we do the
                    # calculations as if it were a K value, and convert at the end
                    var_from = var_from[1:]
                is_fractional = var_from.endswith("__f")
                if is_fractional:
                    # If the uncertainty is fractional, multiply through by this
                    var_from = var_from[:-3]
                    u_from = self.data[var_from] * u_from
                if var_from in self.nodes_original:
                    self.get_grad(var_in, var_from)
                    u_part = np.abs(self.grads[var_in][var_from] * u_from)
                else:
                    # If the uncertainty is from some internally calculated value, then
                    # we need to make a second CO2System where that value is one of the
                    # known inputs, and get the grad from that
                    self.solve(var_from)
                    data = self.get_values_original()
                    data.update({var_from: self.data[var_from]})
                    sys = CO2System(**data, **self.opts)
                    sys.get_grad(var_in, var_from)
                    u_part = np.abs(sys.grads[var_in][var_from] * u_from)
                # Add the p back and convert value, if necessary
                if is_pk:
                    u_part = u_part * np.log(10) * np.abs(sys.data[var_from])
                    var_from = "p" + var_from
                if is_fractional:
                    var_from += "__f"
                self.uncertainty[var_in][var_from] = u_part
                u_total = u_total + u_part**2
            self.uncertainty[var_in]["total"] = np.sqrt(u_total)

    def plot_graph(
        self,
        ax=None,
        exclude_nodes=None,
        prog_graphviz="neato",
        show_tsp=True,
        show_unknown=True,
        show_isolated=True,
        skip_nodes=None,
    ):
        """Draw a graph showing the relationships between the different parameters.

        Parameters
        ----------
        ax : matplotlib axes, optional
            The axes on which to plot.  If `None`, a new figure and axes are created.
        exclude_nodes : list of str, optional
            List of nodes to exclude from the plot, by default `None`.  Nodes in this
            list are not shown, nor are connections to them or through them.
        prog_graphviz : str, optional
            Name of Graphviz layout program, by default "neato".
        show_tsp : bool, optional
            Whether to show temperature, salinity and pressure nodes, by default
            `False`.
        show_unknown : bool, optional
            Whether to show nodes for parameters that have not (yet) been calculated,
            by default `True`.
        show_isolated : bool, optional
            Whether to show nodes for parameters that are not connected to the graph,
            by default `True`.
        skip_nodes : bool, optional
            List of nodes to skip from the plot, by default `None`.  Nodes in this list
            are not shown, but the connections between their predecessors and children
            are still drawn.

        Returns
        -------
        matplotlib axes
            The axes on which the graph is plotted.
        """
        from matplotlib import pyplot as plt

        if ax is None:
            ax = plt.subplots(dpi=300, figsize=(8, 7))[1]
        self_graph = self.graph.copy()
        node_states = nx.get_node_attributes(self_graph, "state", default=0)
        edge_states = nx.get_edge_attributes(self_graph, "state", default=0)
        if not show_tsp:
            self_graph.remove_nodes_from(["pressure", "salinity", "temperature"])
        if not show_unknown:
            self_graph.remove_nodes_from([n for n, s in node_states.items() if s == 0])
        if not show_isolated:
            self_graph.remove_nodes_from(
                [n for n, d in dict(self_graph.degree).items() if d == 0]
            )
        if exclude_nodes:
            # Excluding nodes just makes them disappear from the graph without caring
            # about what they were connected to
            if isinstance(exclude_nodes, str):
                exclude_nodes = [exclude_nodes]
            self_graph.remove_nodes_from(exclude_nodes)
        if skip_nodes:
            # Skipping nodes removes them but then shows their predecessors as being
            # directly connected to their children
            if isinstance(skip_nodes, str):
                skip_nodes = [skip_nodes]
            for n in skip_nodes:
                for p, s in itertools.product(
                    self_graph.predecessors(n), self_graph.successors(n)
                ):
                    self_graph.add_edge(p, s)
                    if edge_states[(p, n)] + edge_states[(n, s)] == 4:
                        new_state = {(p, s): 2}
                    else:
                        new_state = {(p, s): 0}
                    nx.set_edge_attributes(self_graph, new_state, name="state")
                    edge_states.update(new_state)
                self_graph.remove_node(n)
        state_colours = {0: "xkcd:grey", 1: "xkcd:grass", 2: "xkcd:azure"}
        node_colour = [state_colours[node_states[n]] for n in nx.nodes(self_graph)]
        edge_colour = [state_colours[edge_states[e]] for e in nx.edges(self_graph)]
        pos = nx.nx_agraph.graphviz_layout(self.graph, prog=prog_graphviz)
        node_labels = {k: k for k in self_graph.nodes}
        for k, v in set_node_labels.items():
            if k in node_labels:
                node_labels[k] = v
        nx.draw_networkx(
            self_graph,
            ax=ax,
            clip_on=False,
            with_labels=True,
            node_color=node_colour,
            edge_color=edge_colour,
            pos=pos,
            labels=node_labels,
        )
        return ax

adjust(temperature=None, pressure=None, store_steps=1, method_fCO2=1, opt_which_fCO2_insitu=1, bh_upsilon=None)

Adjust the system to a different temperature and/or pressure.

Parameters:

Name	Type	Description	Default
temperature	float	Temperature in °C to adjust to. If None, temperature is not adjusted.	None
pressure	float	Hydrostatic pressure in dbar to adjust to. If None, pressure is not adjusted.	None
store_steps	int	Whether/which non-requested parameters calculated during intermediate calculation steps should be stored. The options are: 0: Store only the requested parameters. 1: Store the requested and most commonly used set of intermediate parameters (default). 2: Store the requested and complete set of intermediate parameters.	1
method_fCO2	int	If this is a single-parameter system, which method to use for the adjustment. The options are: 1: parameterisation of H24 (default). 2: constant bh fitted to TOG93 dataset by H24. 3: constant theoretical bh of H24. 4: user-specified bh with the equations of H24. 5: linear fit of TOG93. 6: quadratic fit of TOG93.	1
opt_which_fCO2_insitu	int	If this is a single-parameter system and method_fCO2 is 1, whether: - 1 the input condition (starting; default) or - 2 output condition (adjusted) temperature should be used to calculate $b_h$.	1
bh_upsilon	float	If this is a single-parameter system and method_fCO2 is 4, then the value of $b_h$ in J/mol must be specified here.	None

Returns:

Type	Description
CO2System	A new CO2System with all values adjusted to the requested temperature(s) and/or pressure(s).

Source code in PyCO2SYS/engine.py

def adjust(
    self,
    temperature=None,
    pressure=None,
    store_steps=1,
    method_fCO2=1,
    opt_which_fCO2_insitu=1,
    bh_upsilon=None,
):
    """Adjust the system to a different temperature and/or pressure.

    Parameters
    ----------
    temperature : float, optional
        Temperature in °C to adjust to.  If `None`, temperature is not adjusted.
    pressure : float, optional
        Hydrostatic pressure in dbar to adjust to.  If `None`, pressure is not
        adjusted.
    store_steps : int, optional
        Whether/which non-requested parameters calculated during intermediate
        calculation steps should be stored.  The options are:

          - `0`: Store only the requested parameters.
          - `1`: Store the requested and most commonly used set of intermediate
          parameters (default).
          - `2`: Store the requested and complete set of intermediate parameters.
    method_fCO2 : int, optional
        If this is a single-parameter system, which method to use for the
        adjustment.  The options are:

          - `1`: parameterisation of H24 (default).
          - `2`: constant bh fitted to TOG93 dataset by H24.
          - `3`: constant theoretical bh of H24.
          - `4`: user-specified bh with the equations of H24.
          - `5`: linear fit of TOG93.
          - `6`: quadratic fit of TOG93.
    opt_which_fCO2_insitu : int, optional
        If this is a single-parameter system and `method_fCO2` is `1`, whether:
          - `1` the input condition (starting; default) or
          - `2` output condition (adjusted) temperature
        should be used to calculate $b_h$.
    bh_upsilon : float, optional
        If this is a single-parameter system and `method_fCO2` is `4`, then the
        value of $b_h$ in J/mol must be specified here.

    Returns
    -------
    CO2System
        A new `CO2System` with all values adjusted to the requested temperature(s)
        and/or pressure(s).
    """
    if self.icase > 100:
        # If we know (any) two MCS parameters, solve for alkalinity and DIC under
        # the "input" conditions
        self.solve(parameters=["alkalinity", "dic"], store_steps=store_steps)
        core = {k: self[k] for k in ["alkalinity", "dic"]}
        data = {}
    elif self.icase in [4, 5, 8, 9]:
        assert pressure is None, (
            "Cannot adjust pressure for a single-parameter system!"
        )
        # If we know only one of pCO2, fCO2, xCO2 or CO2(aq), first get fCO2 under
        # the "input" conditions
        self.solve(parameters="fCO2", store_steps=store_steps)
        core = {"fCO2": self.fCO2}
        # Then, convert this to the value at the new temperature using the requested
        # method
        assert method_fCO2 in range(1, 7), (
            "`method_fCO2` must be an integer from 1-6."
        )
        expUps = self._get_expUps(
            method_fCO2,
            temperature,
            bh_upsilon=bh_upsilon,
            opt_which_fCO2_insitu=opt_which_fCO2_insitu,
        )
        data = {"fCO2": core["fCO2"] * expUps}
    else:
        warnings.warn(
            "Single-parameter temperature adjustments are possible only "
            + "if the known parameter is one of pCO2, fCO2, xCO2 and CO2."
        )
    # Copy all parameters that are T/P independent into new data dict
    for k in condition_independent:
        if k in self.nodes_original:
            data[k] = self.data[k]
        elif k in core:
            data[k] = core[k]
    # Set temperature and/or pressure to adjusted value(s), unless they are None, in
    # which case, don't adjust
    if temperature is not None:
        data["temperature"] = temperature
    else:
        data["temperature"] = self.data["temperature"]
    if pressure is not None:
        data["pressure"] = pressure
    else:
        data["pressure"] = self.data["pressure"]
    sys = CO2System(
        **data,
        **self.opts,
        pd_index=self.pd_index,
        xr_dims=self.xr_dims,
        xr_shape=self.xr_shape,
    )
    sys.solve(parameters=self.data)
    return sys

get_grad(var_of, var_wrt)

Compute the derivative of var_of with respect to var_wrt and store it in sys.grads[var_of][var_wrt]. If there is already a value there, then that value is returned instead of recalculating.

Parameters:

Name	Type	Description	Default
var_of	str	The name of the variable to get the derivative of.	required
var_wrt	str	The name of the variable to get the derivative with respect to. This must be one of the fixed values provided when creating the CO2System, i.e., listed in its nodes_original attribute.	required

Source code in PyCO2SYS/engine.py

def get_grad(self, var_of, var_wrt):
    """Compute the derivative of `var_of` with respect to `var_wrt` and store it in
    `sys.grads[var_of][var_wrt]`.  If there is already a value there, then that
    value is returned instead of recalculating.

    Parameters
    ----------
    var_of : str
        The name of the variable to get the derivative of.
    var_wrt : str
        The name of the variable to get the derivative with respect to.  This must
        be one of the fixed values provided when creating the `CO2System`, i.e.,
        listed in its `nodes_original` attribute.
    """
    assert var_wrt in self.nodes_original, (
        "`var_wrt` must be one of `sys.nodes_original!`"
    )
    try:  # see if we've already calculated this value
        d_of__d_wrt = self.grads[var_of][var_wrt]
    except KeyError:  # only do the calculations if there isn't already a value
        # We need to know the shape of the variable that we want the grad of, the
        # easy way to get this is just to solve for it (if that hasn't already been
        # done)
        if var_of not in self.data:
            self.solve(var_of)
        # Next, we extract the originally set values, which are fixed during the
        # differentiation
        values_original = self.get_values_original()
        other_values_original = values_original.copy()
        # We have to make sure the value we are differentiating with respect to has
        # the same shape as the value we want the differential of
        value_wrt = other_values_original.pop(var_wrt) * np.ones_like(
            self.data[var_of]
        )
        # Here we compute the gradient
        grad_func = self.get_grad_func(var_of, var_wrt)
        d_of__d_wrt = grad_func(value_wrt, **other_values_original)
        # Put the final value into self.grads, first creating a new sub-dict if
        # necessary
        if var_of not in self.grads:
            self.grads[var_of] = {}
        self.grads[var_of][var_wrt] = d_of__d_wrt
    return d_of__d_wrt

get_grads(vars_of, vars_wrt)

Compute the derivatives of vars_of with respect to vars_wrt and store them in sys.grads[var_of][var_wrt]. If there are already values there, then those values are returned instead of recalculating.

Parameters:

Name	Type	Description	Default
vars_of	list	The names of the variables to get the derivatives of.	required
vars_wrt	list	The names of the variables to get the derivatives with respect to. These must all be one of the fixed values provided when creating the CO2System, i.e., listed in its nodes_original attribute.	required

Source code in PyCO2SYS/engine.py

def get_grads(self, vars_of, vars_wrt):
    """Compute the derivatives of `vars_of` with respect to `vars_wrt` and store
    them in `sys.grads[var_of][var_wrt]`.  If there are already values there, then
    those values are returned instead of recalculating.

    Parameters
    ----------
    vars_of : list
        The names of the variables to get the derivatives of.
    vars_wrt : list
        The names of the variables to get the derivatives with respect to.  These
        must all be one of the fixed values provided when creating the `CO2System`,
        i.e., listed in its `nodes_original` attribute.
    """
    if isinstance(vars_of, str):
        vars_of = [vars_of]
    if isinstance(vars_wrt, str):
        vars_wrt = [vars_wrt]
    for var_of, var_wrt in itertools.product(vars_of, vars_wrt):
        self.get_grad(var_of, var_wrt)

plot_graph(ax=None, exclude_nodes=None, prog_graphviz='neato', show_tsp=True, show_unknown=True, show_isolated=True, skip_nodes=None)

Draw a graph showing the relationships between the different parameters.

Parameters:

Name	Type	Description	Default
ax	matplotlib axes	The axes on which to plot. If None, a new figure and axes are created.	None
exclude_nodes	list of str	List of nodes to exclude from the plot, by default None. Nodes in this list are not shown, nor are connections to them or through them.	None
prog_graphviz	str	Name of Graphviz layout program, by default "neato".	'neato'
show_tsp	bool	Whether to show temperature, salinity and pressure nodes, by default False.	True
show_unknown	bool	Whether to show nodes for parameters that have not (yet) been calculated, by default True.	True
show_isolated	bool	Whether to show nodes for parameters that are not connected to the graph, by default True.	True
skip_nodes	bool	List of nodes to skip from the plot, by default None. Nodes in this list are not shown, but the connections between their predecessors and children are still drawn.	None

Returns:

Type	Description
matplotlib axes	The axes on which the graph is plotted.

Source code in PyCO2SYS/engine.py

def plot_graph(
    self,
    ax=None,
    exclude_nodes=None,
    prog_graphviz="neato",
    show_tsp=True,
    show_unknown=True,
    show_isolated=True,
    skip_nodes=None,
):
    """Draw a graph showing the relationships between the different parameters.

    Parameters
    ----------
    ax : matplotlib axes, optional
        The axes on which to plot.  If `None`, a new figure and axes are created.
    exclude_nodes : list of str, optional
        List of nodes to exclude from the plot, by default `None`.  Nodes in this
        list are not shown, nor are connections to them or through them.
    prog_graphviz : str, optional
        Name of Graphviz layout program, by default "neato".
    show_tsp : bool, optional
        Whether to show temperature, salinity and pressure nodes, by default
        `False`.
    show_unknown : bool, optional
        Whether to show nodes for parameters that have not (yet) been calculated,
        by default `True`.
    show_isolated : bool, optional
        Whether to show nodes for parameters that are not connected to the graph,
        by default `True`.
    skip_nodes : bool, optional
        List of nodes to skip from the plot, by default `None`.  Nodes in this list
        are not shown, but the connections between their predecessors and children
        are still drawn.

    Returns
    -------
    matplotlib axes
        The axes on which the graph is plotted.
    """
    from matplotlib import pyplot as plt

    if ax is None:
        ax = plt.subplots(dpi=300, figsize=(8, 7))[1]
    self_graph = self.graph.copy()
    node_states = nx.get_node_attributes(self_graph, "state", default=0)
    edge_states = nx.get_edge_attributes(self_graph, "state", default=0)
    if not show_tsp:
        self_graph.remove_nodes_from(["pressure", "salinity", "temperature"])
    if not show_unknown:
        self_graph.remove_nodes_from([n for n, s in node_states.items() if s == 0])
    if not show_isolated:
        self_graph.remove_nodes_from(
            [n for n, d in dict(self_graph.degree).items() if d == 0]
        )
    if exclude_nodes:
        # Excluding nodes just makes them disappear from the graph without caring
        # about what they were connected to
        if isinstance(exclude_nodes, str):
            exclude_nodes = [exclude_nodes]
        self_graph.remove_nodes_from(exclude_nodes)
    if skip_nodes:
        # Skipping nodes removes them but then shows their predecessors as being
        # directly connected to their children
        if isinstance(skip_nodes, str):
            skip_nodes = [skip_nodes]
        for n in skip_nodes:
            for p, s in itertools.product(
                self_graph.predecessors(n), self_graph.successors(n)
            ):
                self_graph.add_edge(p, s)
                if edge_states[(p, n)] + edge_states[(n, s)] == 4:
                    new_state = {(p, s): 2}
                else:
                    new_state = {(p, s): 0}
                nx.set_edge_attributes(self_graph, new_state, name="state")
                edge_states.update(new_state)
            self_graph.remove_node(n)
    state_colours = {0: "xkcd:grey", 1: "xkcd:grass", 2: "xkcd:azure"}
    node_colour = [state_colours[node_states[n]] for n in nx.nodes(self_graph)]
    edge_colour = [state_colours[edge_states[e]] for e in nx.edges(self_graph)]
    pos = nx.nx_agraph.graphviz_layout(self.graph, prog=prog_graphviz)
    node_labels = {k: k for k in self_graph.nodes}
    for k, v in set_node_labels.items():
        if k in node_labels:
            node_labels[k] = v
    nx.draw_networkx(
        self_graph,
        ax=ax,
        clip_on=False,
        with_labels=True,
        node_color=node_colour,
        edge_color=edge_colour,
        pos=pos,
        labels=node_labels,
    )
    return ax

propagate(uncertainty_in, uncertainty_from)

Propagate independent uncertainties through the calculations. Covariances are not accounted for.

New entries are added in the uncertainty attribute, for example:

co2s = CO2System(dic=2100, alkalinity=2300)
co2s.propagate("pH", {"dic": 2, "alkalinity": 2})
co2s.uncertainty["pH"]["total"]  # total uncertainty in pH
co2s.uncertainty["pH"]["dic"]  # component of ^ due to DIC uncertainty

Parameters:

Name	Type	Description	Default
uncertainty_in	list	The parameters to calculate the uncertainty in.	required
uncertainty_from	dict	The parameters to propagate the uncertainty from (keys) and their uncertainties (values).	required

Source code in PyCO2SYS/engine.py

def propagate(self, uncertainty_in, uncertainty_from):
    """Propagate independent uncertainties through the calculations.  Covariances
    are not accounted for.

    New entries are added in the `uncertainty` attribute, for example:

        co2s = CO2System(dic=2100, alkalinity=2300)
        co2s.propagate("pH", {"dic": 2, "alkalinity": 2})
        co2s.uncertainty["pH"]["total"]  # total uncertainty in pH
        co2s.uncertainty["pH"]["dic"]  # component of ^ due to DIC uncertainty

    Parameters
    ----------
    uncertainty_in : list
        The parameters to calculate the uncertainty in.
    uncertainty_from : dict
        The parameters to propagate the uncertainty from (keys) and their
        uncertainties (values).
    """
    self.solve(uncertainty_in)
    if isinstance(uncertainty_in, str):
        uncertainty_in = [uncertainty_in]
    for var_in in uncertainty_in:
        # This should always be reset to zero and all values wiped, even if it
        # already exists (so you don't end up with old uncertainty_from components
        # from a previous calculation which are no longer part of the total)
        self.uncertainty[var_in] = {"total": np.zeros_like(self.data[var_in])}
        u_total = self.uncertainty[var_in]["total"]
        for var_from, u_from in uncertainty_from.items():
            is_pk = var_from.startswith("pk_")
            if is_pk:
                # If the uncertainty is given in terms of a pK value, we do the
                # calculations as if it were a K value, and convert at the end
                var_from = var_from[1:]
            is_fractional = var_from.endswith("__f")
            if is_fractional:
                # If the uncertainty is fractional, multiply through by this
                var_from = var_from[:-3]
                u_from = self.data[var_from] * u_from
            if var_from in self.nodes_original:
                self.get_grad(var_in, var_from)
                u_part = np.abs(self.grads[var_in][var_from] * u_from)
            else:
                # If the uncertainty is from some internally calculated value, then
                # we need to make a second CO2System where that value is one of the
                # known inputs, and get the grad from that
                self.solve(var_from)
                data = self.get_values_original()
                data.update({var_from: self.data[var_from]})
                sys = CO2System(**data, **self.opts)
                sys.get_grad(var_in, var_from)
                u_part = np.abs(sys.grads[var_in][var_from] * u_from)
            # Add the p back and convert value, if necessary
            if is_pk:
                u_part = u_part * np.log(10) * np.abs(sys.data[var_from])
                var_from = "p" + var_from
            if is_fractional:
                var_from += "__f"
            self.uncertainty[var_in][var_from] = u_part
            u_total = u_total + u_part**2
        self.uncertainty[var_in]["total"] = np.sqrt(u_total)

solve(parameters=None, store_steps=1)

Calculate parameter(s) and store them internally.

Parameters:

Name	Type	Description	Default
parameters	str or list of str	Which parameter(s) to calculate and store, by default None, in which case all possible parameters are calculated and stored internally.	None
store_steps	int	Whether/which non-requested parameters calculated during intermediate calculation steps should be stored, by default 1. The options are 0 - store only the specifically requested parameters, 1 - store the most used set of intermediate parameters, or 2 - store the complete set of parameters.	1

Source code in PyCO2SYS/engine.py

def solve(self, parameters=None, store_steps=1):
    """Calculate parameter(s) and store them internally.

    Parameters
    ----------
    parameters : str or list of str, optional
        Which parameter(s) to calculate and store, by default None, in which case
        all possible parameters are calculated and stored internally.
    store_steps : int, optional
        Whether/which non-requested parameters calculated during intermediate
        calculation steps should be stored, by default 1.  The options are
            0 - store only the specifically requested parameters,
            1 - store the most used set of intermediate parameters, or
            2 - store the complete set of parameters.
    """
    # Parse user-provided parameters (if there are any)
    if parameters is None:
        # If no parameters are provided, then we solve for everything possible
        parameters = list(self.graph.nodes)
    elif isinstance(parameters, str):
        # Allow user to provide a string if only one parameter is desired
        parameters = [parameters]
    parameters = set(parameters)  # get rid of duplicates
    self.requested |= parameters
    self_data = self.data.copy()  # what was already known before this solve
    # Remove known nodes from a copy of self.graph, so that ancestors of known nodes
    # are not unnecessarily recomputed
    graph_unknown = self.graph.copy()
    graph_unknown.remove_nodes_from([k for k in self_data if k not in parameters])
    # Add intermediate parameters that we need to know in order to calculate the
    # requested parameters
    parameters_all = parameters.copy()
    for p in parameters:
        parameters_all = parameters_all | nx.ancestors(graph_unknown, p)
    # Convert the set of parameters into a list, exclude already-known ones, and
    # organise the list into the order required for calculations
    parameters_all = [
        p
        for p in nx.topological_sort(self.graph)
        if p in parameters_all and p not in self_data
    ]
    store_parameters = []
    for p in parameters_all:
        priors = self.graph.pred[p]
        if len(priors) == 0 or all([r in self_data for r in priors]):
            self_data[p] = self.funcs[p](
                *[
                    self_data[r]
                    for r in self.funcs[p].__code__.co_varnames[
                        : self.funcs[p].__code__.co_argcount
                    ]
                ]
            )
            store_here = (
                #  If store_steps is 0, store only requested parameters
                (store_steps == 0 and p in parameters)
                | (
                    # If store_steps is 1, store all but the equilibrium constants
                    # on the seawater scale, at 1 atm and their pressure-correction
                    # factors, and a few selected others
                    store_steps == 1
                    and not p.startswith("factor_k_")
                    and not (p.startswith("k_") and p.endswith("_sws"))
                    and not p.endswith("_1atm")
                    and p not in ["sws_to_opt", "opt_to_free", "ionic_strength"]
                )
                |  # If store_steps is 2, store everything
                (store_steps == 2)
            )
            if store_here:
                store_parameters.append(p)
                # state = 2 means that the value was calculated internally
                nx.set_node_attributes(self.graph, {p: 2}, name="state")
                for f in self.funcs[p].__code__.co_varnames[
                    : self.funcs[p].__code__.co_argcount
                ]:
                    nx.set_edge_attributes(self.graph, {(f, p): 2}, name="state")
    # Get rid of jax overhead on results
    self_data = {k: v for k, v in self_data.items() if k in store_parameters}
    _remove_jax_overhead(self_data)
    self.data.update(self_data)

to_pandas(parameters=None, store_steps=1)

Return parameters as a pandas Series or DataFrame. All parameters should be scalar or one-dimensional vectors of the same size.

Parameters:

Name	Type	Description	Default
parameters	str or list of str	The parameter(s) to return. These are solved for if not already available. If None, then all parameters that have already been solved for are returned.	None
store_steps	int	See solve.	1

Returns:

Type	Description
Series or DataFrame	The parameter(s) as a pd.Series (if parameters is a str) or as a pd.DataFrame (if parameters is a list) with the original pandas index passed into the CO2System as data. If data was not a pd.DataFrame then the default index will be used.

Source code in PyCO2SYS/engine.py

def to_pandas(self, parameters=None, store_steps=1):
    """Return parameters as a pandas `Series` or `DataFrame`.  All parameters should
    be scalar or one-dimensional vectors of the same size.

    Parameters
    ----------
    parameters : str or list of str, optional
        The parameter(s) to return.  These are solved for if not already available.
        If `None`, then all parameters that have already been solved for are
        returned.
    store_steps : int, optional
        See `solve`.

    Returns
    -------
    pd.Series or pd.DataFrame
        The parameter(s) as a `pd.Series` (if `parameters` is a `str`) or as a
        `pd.DataFrame` (if `parameters` is a `list`) with the original pandas index
        passed into the `CO2System` as `data`.  If `data` was not a `pd.DataFrame`
        then the default index will be used.
    """
    try:
        import pandas as pd

        if parameters is None:
            parameters = self.keys()
        self.solve(parameters=parameters, store_steps=store_steps)
        if isinstance(parameters, str):
            return pd.Series(data=self[parameters], index=self.pd_index)
        else:
            return pd.DataFrame(
                {
                    p: pd.Series(
                        data=self[p] * np.ones(self.pd_index.shape),
                        index=self.pd_index,
                    )
                    for p in parameters
                }
            )
    except ImportError:
        warnings.warn("pandas could not be imported.")

to_xarray(parameters=None, store_steps=1)

Return parameters as an xarray DataArray or Dataset.

Parameters:

Name	Type	Description	Default
parameters	str or list of str	The parameter(s) to return. These are solved for if not already available. If None, then all parameters that have already been solved for are returned.	None
store_steps	int	See solve.	1

Returns:

Type	Description
DataArray or Dataset	The parameter(s) as a xr.DataArray (if parameters is a str) or as a xr.Dataset (if parameters is a list) with the original xarray dimensions passed into the CO2System as data. If data was not an xr.Dataset then this function will not work.

Source code in PyCO2SYS/engine.py

def to_xarray(self, parameters=None, store_steps=1):
    """Return parameters as an xarray `DataArray` or `Dataset`.

    Parameters
    ----------
    parameters : str or list of str, optional
        The parameter(s) to return.  These are solved for if not already available.
        If `None`, then all parameters that have already been solved for are
        returned.
    store_steps : int, optional
        See `solve`.

    Returns
    -------
    xr.DataArray or xr.Dataset
        The parameter(s) as a `xr.DataArray` (if `parameters` is a `str`) or as a
        `xr.Dataset` (if `parameters` is a `list`) with the original xarray
        dimensions passed into the `CO2System` as `data`.  If `data` was not an
        `xr.Dataset` then this function will not work.
    """
    assert self.xr_dims is not None and self.xr_shape is not None, (
        "`data` was not provided as an `xr.Dataset` "
        + "when creating this `CO2System`."
    )
    try:
        import xarray as xr

        if parameters is None:
            parameters = self.keys()
        self.solve(parameters=parameters, store_steps=store_steps)
        if isinstance(parameters, str):
            ndims = self._get_xr_ndims(parameters)
            return xr.DataArray(np.squeeze(self[parameters]), dims=ndims)
        else:
            return xr.Dataset(
                {
                    p: xr.DataArray(np.squeeze(self[p]), dims=self._get_xr_ndims(p))
                    for p in parameters
                }
            )
    except ImportError:
        warnings.warn("xarray could not be imported.")