jelli 0.0.1__py3-none-any.whl

jelli/__init__.py

@@ -0,0 +1,5 @@
+from ._version import __version__
+from .core.global_likelihood import GlobalLikelihood
+from .core.observable_sector import ObservableSector
+from .core.measurement import Measurement
+from .core.custom_basis import CustomBasis

jelli/_version.py

@@ -0,0 +1 @@
+__version__ = '0.0.1'

jelli/core/custom_basis.py

@@ -0,0 +1,177 @@
+from typing import List, Dict, Union
+
+class CustomBasis:
+    '''
+    A class to represent a custom parameter basis.
+
+    Parameters
+    ----------
+    name : str
+        The name of the custom basis.
+    parameters : list or dict
+        The parameters in the custom basis. If a list is provided, all parameters are assumed to be real. If a dict is provided, the keys are the parameter names and the values are the types (e.g., 'R' for real, 'C' for complex).
+
+    Attributes
+    ----------
+    name : str
+        The name of the custom basis.
+    parameters : dict
+        The parameters in the custom basis. The keys are the parameter names and the values are the types (e.g., 'R' for real, 'C' for complex).
+
+    Methods
+    -------
+    get_all_names() -> List[str]
+        Get all custom basis names.
+    get(name: str) -> 'CustomBasis'
+        Get a custom basis by name.
+    get_all() -> List['CustomBasis']
+        Get all custom basis objects.
+    get_parameter_basis() -> List
+        Get the parameter basis.
+
+    Examples
+    --------
+    Initialize a custom basis with real parameters:
+
+    >>> custom_basis = CustomBasis('example_basis', ['param1', 'param2'])
+
+    Initialize a custom basis with mixed parameters:
+
+    >>> custom_basis = CustomBasis('example_basis', {'param1': 'R', 'param2': 'C'})
+
+    Get all custom basis names:
+
+    >>> CustomBasis.get_all_names()
+
+    Get a custom basis by name:
+
+    >>> custom_basis = CustomBasis.get('example_basis')
+
+    Get all custom basis objects:
+
+    >>> custom_bases = CustomBasis.get_all()
+
+    Get the parameter basis:
+
+    >>> parameter_basis = custom_basis.get_parameter_basis()
+
+    '''
+
+    _custom_bases: Dict[str, 'CustomBasis'] = {}  # Class attribute to store all custom bases
+
+    def __init__(self, name: str, parameters: Union[List[str], Dict[str, str]]):
+        """
+        Initialize the CustomBasis class.
+
+        Parameters
+        ----------
+            name : str
+                The name of the custom basis.
+            parameters : list or dict
+                The parameters in the custom basis. If a list is provided, all parameters are assumed to be real. If a dict is provided, the keys are the parameter names and the values are the types (e.g., 'R' for real, 'C' for complex).
+
+        Returns
+        --------
+            None
+
+        Examples
+        --------
+        Initialize a custom basis with real parameters:
+
+        >>> custom_basis = CustomBasis('example_basis', ['param1', 'param2'])
+
+        Initialize a custom basis with mixed parameters:
+
+        >>> custom_basis = CustomBasis('example_basis', {'param1': 'R', 'param2': 'C'})
+        """
+        self.name = name
+        if isinstance(parameters, list):
+            self.parameters = {param: 'R' for param in parameters}
+        elif isinstance(parameters, dict):
+            if not all(value in ['R', 'C'] for value in parameters.values()):
+                raise ValueError("All parameter types must be either 'R' or 'C'.")
+            self.parameters = parameters
+        else:
+            raise ValueError("Parameters must be a list or a dictionary.")
+        self._custom_bases[self.name] = self
+
+    @classmethod
+    def get_all_names(cls) -> List[str]:
+        """
+        Get all custom basis names.
+
+        Returns
+        --------
+            list
+                A list of all custom basis names.
+
+        Examples
+        --------
+            >>> CustomBasis.get_all_names()
+        """
+        return sorted(cls._custom_bases.keys())
+
+    @classmethod
+    def get(cls, name: str) -> 'CustomBasis':
+        """
+        Get a custom basis by name.
+
+        Parameters
+        ----------
+            name : str
+                The name of the custom basis.
+        Returns
+        --------
+            CustomBasis
+                The custom basis object.
+
+        Examples
+        --------
+            >>> CustomBasis.get('example_basis')
+        """
+        return cls._custom_bases.get(name)
+
+    @classmethod
+    def get_all(cls) -> List['CustomBasis']:
+        """
+        Get all custom basis objects.
+
+        Returns
+        --------
+            list
+                A list of all custom basis objects.
+
+        Examples
+        --------
+            >>> CustomBasis.get_all()
+        """
+        return list(cls._custom_bases.values())
+
+    def get_parameter_basis(self, split_re_im=True) -> List:
+        """
+        Get the parameter basis.
+
+        Parameters
+        ----------
+            split_re_im : bool, optional
+                If True, split parameters into real and imaginary parts, otherwise return the parameters directly. Default is True.
+
+        Returns
+        --------
+            list
+                A list containing the parameter basis.
+
+        Examples
+        --------
+            >>> custom_basis = CustomBasis('example_basis', ['param1', 'param2'])
+            >>> custom_basis.get_parameter_basis('example_basis')
+        """
+        if split_re_im:
+            parameter_basis = []
+            for parameter, parameter_type in self.parameters.items():
+                parameter_basis.append((parameter, 'R'))
+                if parameter_type == 'C':
+                    parameter_basis.append((parameter, 'I'))
+        else:
+            parameter_basis = self.parameters.keys()
+        return sorted(parameter_basis)

jelli/core/experimental_correlations.py

@@ -0,0 +1,100 @@
+from typing import Dict, Optional, Iterable
+import h5py
+import numpy as np
+from jax import numpy as jnp
+import os
+from ..utils.data_io import get_json_schema, hash_observable_names
+
+class ExperimentalCorrelations:
+
+    _instances: Dict[str, Dict[str, 'ExperimentalCorrelations']] = {
+        'correlations': {},
+        'central': {},
+        'uncertainties': {}
+    }  # Dictionary to hold instances of ExperimentalCorrelations
+    _covariance_scaled: Dict[str, jnp.ndarray] = {}
+    _pdfxf_h5_versions = {'1.0'} # Set of supported versions of the pdfxf-h5 JSON schema
+
+    def __init__(
+        self,
+        hash_val: str,
+        data_type: str,
+        data: np.ndarray,
+        row_names: Dict[str, int],
+        col_names: Dict[str, int]
+    ) -> None:
+        self.hash_val = hash_val
+        self.data_type = data_type
+        self.data = data
+        self.row_names = row_names
+        self.col_names = col_names
+        self._instances[data_type][hash_val] = self
+
+    @classmethod
+    def _load_file(cls, path: str) -> None:
+        with h5py.File(path, 'r') as f:
+            schema_name, schema_version = get_json_schema(dict(f.attrs))
+            if schema_name == 'pdfxf-h5' and schema_version in cls._pdfxf_h5_versions:
+                for hash_val in f['correlations']:
+                    cls.from_hdf5_group(hash_val, 'correlations', f['correlations'][hash_val])
+                for hash_val in f['central']:
+                    cls.from_hdf5_group(hash_val, 'central', f['central'][hash_val])
+                for hash_val in f['uncertainties']:
+                    cls.from_hdf5_group(hash_val, 'uncertainties', f['uncertainties'][hash_val])
+
+    @classmethod
+    def load(cls, path: str) -> None:
+        # load all hdf5 files in the directory
+        if os.path.isdir(path):
+            for file in os.listdir(path):
+                if file.endswith('.hdf5'):
+                    cls._load_file(os.path.join(path, file))
+        # load single hdf5 file
+        else:
+            cls._load_file(path)
+
+    @classmethod
+    def from_hdf5_group(cls, hash_val: str, data_type: str, hdf5_group: h5py.Group) -> None:
+        data = hdf5_group['data']
+        data = np.array(data[()], dtype=np.float64) * data.attrs.get('scale', 1.0)
+        row_names = {name: i for i, name in enumerate(hdf5_group['row_names'][()].astype(str))}
+        col_names = {name: i for i, name in enumerate(hdf5_group.get('col_names', np.array([]))[()].astype(str))}
+        cls(hash_val, data_type, data, row_names, col_names)
+
+    @classmethod
+    def get_data(
+        cls,
+        data_type: str,
+        row_names: Iterable[str],
+        col_names: Optional[Iterable[str]] = []
+    ):
+        hash_val = hash_observable_names(row_names, col_names)
+        if hash_val in cls._instances[data_type]:
+            data = cls._instances[data_type][hash_val].data
+        else:
+            hash_val = hash_observable_names(col_names, row_names)
+            if hash_val in cls._instances[data_type]:
+                data = cls._instances[data_type][hash_val].data.T
+            else:
+                data = None
+        return data
+
+    @classmethod
+    def get_cov_scaled(
+        cls,
+        row_names: Iterable[str],
+        col_names: Iterable[str],
+        std_exp_scaled_row: np.ndarray,
+        std_exp_scaled_col: np.ndarray,
+    ):
+        hash_val = hash_observable_names(row_names, col_names)
+        if hash_val in cls._covariance_scaled:
+            cov_scaled = cls._covariance_scaled[hash_val]
+        else:
+            corr = cls.get_data('correlations', row_names, col_names)
+            if corr is None:
+                raise ValueError(f"Correlation data for {row_names} and {col_names} not found.")
+            cov_scaled = corr * np.outer(std_exp_scaled_row, std_exp_scaled_col)
+            cov_scaled = jnp.array(cov_scaled, dtype=jnp.float64)
+            cls._covariance_scaled[hash_val] = cov_scaled
+        return cov_scaled