Source code for icet.core.local_cluster_expansion

"""
This module provides the LocalClusterExpansion class.
"""

import json
from typing import Any

import numpy as np

from ase import Atoms

from icet.core.base_cluster_expansion import BaseClusterExpansion
from icet.core.local_cluster_space import LocalClusterSpace
from icet.input_output.repr_tools import (html_representation,
                                          text_representation)


def encode_numpy(value):
    """Converts numpy values to plain Python for JSON serialization.

    Metadata commonly comes from a fitting summary, for example the one of
    the trainstation CrossValidationEstimator, whose values carry numpy
    types.
    """
    if isinstance(value, np.integer):
        return int(value)
    if isinstance(value, np.floating):
        return float(value)
    if isinstance(value, np.ndarray):
        return value.tolist()
    raise TypeError(f'{type(value)} is not JSON serializable')




class LocalClusterExpansion(BaseClusterExpansion):
    """A local cluster expansion pairs a local cluster space with a set of
    parameters, where the latter is commonly obtained by fitting to locally
    defined energy differences such as kinetically resolved activation
    barriers.
    Instances of this class predict the local property for an event
    instance in a supercell.

    Parameters
    ----------
    local_cluster_space
        Local cluster space that defines the basis of the expansion.
    parameters
        Parameter vector, one element per element of the space.
    metadata
        User-defined metadata to be stored together with the expansion.

    Raises
    ------
    ValueError
        If :attr:`local_cluster_space` and :attr:`parameters` differ in
        length.
    """

    _space_name = 'local_cluster_space'

    def __init__(self,
                 local_cluster_space: LocalClusterSpace,
                 parameters: np.ndarray,
                 metadata: dict | None = None) -> None:
        self._local_cluster_space = local_cluster_space
        super().__init__(local_cluster_space, parameters, metadata)

    @property
    def _space(self) -> LocalClusterSpace:
        """The space that defines the basis of the expansion."""
        return self._local_cluster_space

    @property
    def local_cluster_space(self) -> LocalClusterSpace:
        """Local cluster space on which the expansion is based."""
        return self._local_cluster_space



    def predict(self, structure: Atoms, event_sites: tuple) -> float:
        """Returns the local property predicted for an event instance.

        Parameters
        ----------
        structure
            Supercell whose occupations define the environment.
        event_sites
            Indices of the anchor sites in the supercell.
        """
        vector = self._local_cluster_space.get_local_cluster_vector(
            structure, event_sites)
        return float(np.dot(vector, self.parameters))


    def _get_rows(self) -> list[tuple[str, Any]]:
        """Returns the label and value pairs that summarize this expansion,
        which are those of the space with the number of nonzero parameters
        added to the number of parameters."""
        nonzero = int(np.count_nonzero(self._parameters))
        rows = []
        for label, value in self._local_cluster_space._get_rows():
            if label == 'total number of parameters':
                label, value = (f'{label} (nonzero)', f'{value} ({nonzero})')
            rows.append((label, value))
        return rows

    def __str__(self) -> str:
        """ String representation. """
        return text_representation(title='Local Cluster Expansion',
                                   rows=self._get_rows(), width=60,
                                   label_width=36)

    def _repr_html_(self) -> str:
        """ HTML representation. Used, e.g., in jupyter notebooks. """
        return html_representation(title='Local Cluster Expansion',
                                   rows=self._get_rows())



    def write(self, filename: str) -> None:
        """Writes the local cluster expansion to a file.

        Metadata is serialized as JSON, with numpy values converted to
        plain numbers and lists.

        Parameters
        ----------
        filename
            Name of the file to which to write.
        """
        data = {'type': type(self).__name__,
                'local_cluster_space': self._local_cluster_space.to_dict(),
                'parameters': [float(p) for p in self._parameters],
                'metadata': self._metadata}
        # Serialization happens before the file is opened, so a metadata
        # value that is not serializable raises without truncating a
        # previously saved expansion.
        payload = json.dumps(data, ensure_ascii=False,
                             default=encode_numpy)
        with open(filename, 'w', encoding='utf-8') as handle:
            handle.write(payload)




    @classmethod
    def read(cls, filename: str) -> 'LocalClusterExpansion':
        """Reads a local cluster expansion from a file.

        Parameters
        ----------
        filename
            Name of the file from which to read.
        """
        with open(filename, encoding='utf-8') as handle:
            data = json.load(handle)
        if data.get('type') != cls.__name__:
            raise ValueError(f'{filename} does not contain a {cls.__name__}')
        expansion = cls(LocalClusterSpace.from_dict(data['local_cluster_space']),
                        np.array(data['parameters']))
        expansion._metadata = data['metadata']
        return expansion