arviz 0.23.3__py3-none-any.whl → 1.0.0rc0__py3-none-any.whl

arviz/data/converters.py

@@ -1,203 +0,0 @@
-"""High level conversion functions."""
-
-import numpy as np
-import xarray as xr
-import pandas as pd
-
-try:
-    from tree import is_nested
-except ImportError:
-    is_nested = lambda obj: False
-
-from .base import dict_to_dataset
-from .inference_data import InferenceData
-from .io_beanmachine import from_beanmachine
-from .io_cmdstan import from_cmdstan
-from .io_cmdstanpy import from_cmdstanpy
-from .io_emcee import from_emcee
-from .io_numpyro import from_numpyro
-from .io_pyro import from_pyro
-from .io_pystan import from_pystan
-
-
-# pylint: disable=too-many-return-statements
-def convert_to_inference_data(obj, *, group="posterior", coords=None, dims=None, **kwargs):
-    r"""Convert a supported object to an InferenceData object.
-
-    This function sends `obj` to the right conversion function. It is idempotent,
-    in that it will return arviz.InferenceData objects unchanged.
-
-    Parameters
-    ----------
-    obj : dict, str, np.ndarray, xr.Dataset, pystan fit
-        A supported object to convert to InferenceData:
-            | InferenceData: returns unchanged
-            | str: Attempts to load the cmdstan csv or netcdf dataset from disk
-            | pystan fit: Automatically extracts data
-            | cmdstanpy fit: Automatically extracts data
-            | cmdstan csv-list: Automatically extracts data
-            | emcee sampler: Automatically extracts data
-            | pyro MCMC: Automatically extracts data
-            | beanmachine MonteCarloSamples: Automatically extracts data
-            | xarray.Dataset: adds to InferenceData as only group
-            | xarray.DataArray: creates an xarray dataset as the only group, gives the
-                         array an arbitrary name, if name not set
-            | dict: creates an xarray dataset as the only group
-            | numpy array: creates an xarray dataset as the only group, gives the
-                         array an arbitrary name
-            | object with __array__: converts to numpy array, then creates an xarray dataset as
-                         the only group, gives the array an arbitrary name
-    group : str
-        If `obj` is a dict or numpy array, assigns the resulting xarray
-        dataset to this group. Default: "posterior".
-    coords : dict[str, iterable]
-        A dictionary containing the values that are used as index. The key
-        is the name of the dimension, the values are the index values.
-    dims : dict[str, List(str)]
-        A mapping from variables to a list of coordinate names for the variable
-    kwargs
-        Rest of the supported keyword arguments transferred to conversion function.
-
-    Returns
-    -------
-    InferenceData
-    """
-    kwargs[group] = obj
-    kwargs["coords"] = coords
-    kwargs["dims"] = dims
-
-    # Cases that convert to InferenceData
-    if isinstance(obj, InferenceData):
-        if coords is not None or dims is not None:
-            raise TypeError("Cannot use coords or dims arguments with InferenceData value.")
-        return obj
-    elif isinstance(obj, str):
-        if obj.endswith(".csv"):
-            if group == "sample_stats":
-                kwargs["posterior"] = kwargs.pop(group)
-            elif group == "sample_stats_prior":
-                kwargs["prior"] = kwargs.pop(group)
-            return from_cmdstan(**kwargs)
-        else:
-            if coords is not None or dims is not None:
-                raise TypeError(
-                    "Cannot use coords or dims arguments reading InferenceData from netcdf."
-                )
-            return InferenceData.from_netcdf(obj)
-    elif (
-        obj.__class__.__name__ in {"StanFit4Model", "CmdStanMCMC"}
-        or obj.__class__.__module__ == "stan.fit"
-    ):
-        if group == "sample_stats":
-            kwargs["posterior"] = kwargs.pop(group)
-        elif group == "sample_stats_prior":
-            kwargs["prior"] = kwargs.pop(group)
-        if obj.__class__.__name__ == "CmdStanMCMC":
-            return from_cmdstanpy(**kwargs)
-        else:  # pystan or pystan3
-            return from_pystan(**kwargs)
-    elif obj.__class__.__name__ == "EnsembleSampler":  # ugly, but doesn't make emcee a requirement
-        return from_emcee(sampler=kwargs.pop(group), **kwargs)
-    elif obj.__class__.__name__ == "MonteCarloSamples":
-        return from_beanmachine(sampler=kwargs.pop(group), **kwargs)
-    elif obj.__class__.__name__ == "MCMC" and obj.__class__.__module__.startswith("pyro"):
-        return from_pyro(posterior=kwargs.pop(group), **kwargs)
-    elif obj.__class__.__name__ == "MCMC" and obj.__class__.__module__.startswith("numpyro"):
-        return from_numpyro(posterior=kwargs.pop(group), **kwargs)
-
-    # Cases that convert to xarray
-    if isinstance(obj, xr.Dataset):
-        dataset = obj
-    elif isinstance(obj, xr.DataArray):
-        if obj.name is None:
-            obj.name = "x"
-        dataset = obj.to_dataset()
-    elif isinstance(obj, dict):
-        dataset = dict_to_dataset(obj, coords=coords, dims=dims)
-    elif is_nested(obj) and not isinstance(obj, (list, tuple)):
-        dataset = dict_to_dataset(obj, coords=coords, dims=dims)
-    elif isinstance(obj, np.ndarray):
-        dataset = dict_to_dataset({"x": obj}, coords=coords, dims=dims)
-    elif (
-        hasattr(obj, "__array__")
-        and callable(getattr(obj, "__array__"))
-        and (not isinstance(obj, pd.DataFrame))
-    ):
-        obj = obj.__array__()
-        dataset = dict_to_dataset({"x": obj}, coords=coords, dims=dims)
-    elif isinstance(obj, (list, tuple)) and isinstance(obj[0], str) and obj[0].endswith(".csv"):
-        if group == "sample_stats":
-            kwargs["posterior"] = kwargs.pop(group)
-        elif group == "sample_stats_prior":
-            kwargs["prior"] = kwargs.pop(group)
-        return from_cmdstan(**kwargs)
-    else:
-        allowable_types = (
-            "xarray dataarray",
-            "xarray dataset",
-            "dict",
-            "pytree (if 'dm-tree' is installed)",
-            "netcdf filename",
-            "numpy array",
-            "object with __array__",
-            "pystan fit",
-            "emcee fit",
-            "pyro mcmc fit",
-            "numpyro mcmc fit",
-            "cmdstan fit csv filename",
-            "cmdstanpy fit",
-        )
-        raise ValueError(
-            f'Can only convert {", ".join(allowable_types)} to InferenceData, '
-            f"not {obj.__class__.__name__}"
-        )
-
-    return InferenceData(**{group: dataset})
-
-
-def convert_to_dataset(obj, *, group="posterior", coords=None, dims=None):
-    """Convert a supported object to an xarray dataset.
-
-    This function is idempotent, in that it will return xarray.Dataset functions
-    unchanged. Raises `ValueError` if the desired group can not be extracted.
-
-    Note this goes through a DataInference object. See `convert_to_inference_data`
-    for more details. Raises ValueError if it can not work out the desired
-    conversion.
-
-    Parameters
-    ----------
-    obj : dict, str, np.ndarray, xr.Dataset, pystan fit
-        A supported object to convert to InferenceData:
-
-        - InferenceData: returns unchanged
-        - str: Attempts to load the netcdf dataset from disk
-        - pystan fit: Automatically extracts data
-        - xarray.Dataset: adds to InferenceData as only group
-        - xarray.DataArray: creates an xarray dataset as the only group, gives the
-          array an arbitrary name, if name not set
-        - dict: creates an xarray dataset as the only group
-        - numpy array: creates an xarray dataset as the only group, gives the
-          array an arbitrary name
-
-    group : str
-        If `obj` is a dict or numpy array, assigns the resulting xarray
-        dataset to this group.
-    coords : dict[str, iterable]
-        A dictionary containing the values that are used as index. The key
-        is the name of the dimension, the values are the index values.
-    dims : dict[str, List(str)]
-        A mapping from variables to a list of coordinate names for the variable
-
-    Returns
-    -------
-    xarray.Dataset
-    """
-    inference_data = convert_to_inference_data(obj, group=group, coords=coords, dims=dims)
-    dataset = getattr(inference_data, group, None)
-    if dataset is None:
-        raise ValueError(
-            "Can not extract {group} from {obj}! See {filename} for other "
-            "conversion utilities.".format(group=group, obj=obj, filename=__file__)
-        )
-    return dataset

arviz/data/datasets.py

@@ -1,161 +0,0 @@
-"""Base IO code for all datasets. Heavily influenced by scikit-learn's implementation."""
-
-import hashlib
-import itertools
-import json
-import os
-import shutil
-from collections import namedtuple
-from urllib.request import urlretrieve
-
-from ..rcparams import rcParams
-from .io_netcdf import from_netcdf
-
-LocalFileMetadata = namedtuple("LocalFileMetadata", ["name", "filename", "description"])
-
-RemoteFileMetadata = namedtuple(
-    "RemoteFileMetadata", ["name", "filename", "url", "checksum", "description"]
-)
-_EXAMPLE_DATA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "example_data")
-_LOCAL_DATA_DIR = os.path.join(_EXAMPLE_DATA_DIR, "data")
-
-with open(os.path.join(_EXAMPLE_DATA_DIR, "data_local.json"), "r", encoding="utf-8") as f:
-    LOCAL_DATASETS = {
-        entry["name"]: LocalFileMetadata(
-            name=entry["name"],
-            filename=os.path.join(_LOCAL_DATA_DIR, entry["filename"]),
-            description=entry["description"],
-        )
-        for entry in json.load(f)
-    }
-
-with open(os.path.join(_EXAMPLE_DATA_DIR, "data_remote.json"), "r", encoding="utf-8") as f:
-    REMOTE_DATASETS = {entry["name"]: RemoteFileMetadata(**entry) for entry in json.load(f)}
-
-
-def get_data_home(data_home=None):
-    """Return the path of the arviz data dir.
-
-    This folder is used by some dataset loaders to avoid downloading the
-    data several times.
-
-    By default the data dir is set to a folder named 'arviz_data' in the
-    user home folder.
-
-    Alternatively, it can be set by the 'ARVIZ_DATA' environment
-    variable or programmatically by giving an explicit folder path. The '~'
-    symbol is expanded to the user home folder.
-
-    If the folder does not already exist, it is automatically created.
-
-    Parameters
-    ----------
-    data_home : str | None
-        The path to arviz data dir.
-    """
-    if data_home is None:
-        data_home = os.environ.get("ARVIZ_DATA", os.path.join("~", "arviz_data"))
-    data_home = os.path.expanduser(data_home)
-    if not os.path.exists(data_home):
-        os.makedirs(data_home)
-    return data_home
-
-
-def clear_data_home(data_home=None):
-    """Delete all the content of the data home cache.
-
-    Parameters
-    ----------
-    data_home : str | None
-        The path to arviz data dir.
-    """
-    data_home = get_data_home(data_home)
-    shutil.rmtree(data_home)
-
-
-def _sha256(path):
-    """Calculate the sha256 hash of the file at path."""
-    sha256hash = hashlib.sha256()
-    chunk_size = 8192
-    with open(path, "rb") as buff:
-        while True:
-            buffer = buff.read(chunk_size)
-            if not buffer:
-                break
-            sha256hash.update(buffer)
-    return sha256hash.hexdigest()
-
-
-def load_arviz_data(dataset=None, data_home=None, **kwargs):
-    """Load a local or remote pre-made dataset.
-
-    Run with no parameters to get a list of all available models.
-
-    The directory to save to can also be set with the environment
-    variable `ARVIZ_HOME`. The checksum of the dataset is checked against a
-    hardcoded value to watch for data corruption.
-
-    Run `az.clear_data_home` to clear the data directory.
-
-    Parameters
-    ----------
-    dataset : str
-        Name of dataset to load.
-    data_home : str, optional
-        Where to save remote datasets
-    **kwargs : dict, optional
-        Keyword arguments passed to :func:`arviz.from_netcdf`.
-
-    Returns
-    -------
-    xarray.Dataset
-
-    """
-    if dataset in LOCAL_DATASETS:
-        resource = LOCAL_DATASETS[dataset]
-        return from_netcdf(resource.filename, **kwargs)
-
-    elif dataset in REMOTE_DATASETS:
-        remote = REMOTE_DATASETS[dataset]
-        home_dir = get_data_home(data_home=data_home)
-        file_path = os.path.join(home_dir, remote.filename)
-
-        if not os.path.exists(file_path):
-            http_type = rcParams["data.http_protocol"]
-
-            # Replaces http type. Redundant if http_type is http, useful if http_type is https
-            url = remote.url.replace("http", http_type)
-            urlretrieve(url, file_path)
-
-        checksum = _sha256(file_path)
-        if remote.checksum != checksum:
-            raise IOError(
-                f"{file_path} has an SHA256 checksum ({checksum}) differing from expected "
-                "({remote.checksum}), file may be corrupted. "
-                "Run `arviz.clear_data_home()` and try again, or please open an issue."
-            )
-        return from_netcdf(file_path, **kwargs)
-    else:
-        if dataset is None:
-            return dict(itertools.chain(LOCAL_DATASETS.items(), REMOTE_DATASETS.items()))
-        else:
-            raise ValueError(
-                "Dataset {} not found! The following are available:\n{}".format(
-                    dataset, list_datasets()
-                )
-            )
-
-
-def list_datasets():
-    """Get a string representation of all available datasets with descriptions."""
-    lines = []
-    for name, resource in itertools.chain(LOCAL_DATASETS.items(), REMOTE_DATASETS.items()):
-        if isinstance(resource, LocalFileMetadata):
-            location = f"local: {resource.filename}"
-        elif isinstance(resource, RemoteFileMetadata):
-            location = f"remote: {resource.url}"
-        else:
-            location = "unknown"
-        lines.append(f"{name}\n{'=' * len(name)}\n{resource.description}\n\n{location}")
-
-    return f"\n\n{10 * '-'}\n\n".join(lines)