jelli 0.0.1__tar.gz

jelli-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Aleks Smolkovič & Peter Stangl
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jelli-0.0.1/PKG-INFO

@@ -0,0 +1,20 @@
+Metadata-Version: 2.4
+Name: jelli
+Version: 0.0.1
+Summary: A JAX-based EFT likelihood.
+Author-email: Aleks Smolkovič <aleks.smolkovic@ijs.si>, Peter Stangl <peter.stangl@cern.ch>
+License: MIT
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: scipy
+Requires-Dist: jax
+Requires-Dist: wilson
+Requires-Dist: rgevolve
+Requires-Dist: multipledispatch
+Requires-Dist: networkx
+Requires-Dist: h5py
+Dynamic: license-file
+
+# jelli - a JAX-based EFT Likelihood

jelli-0.0.1/README.md

@@ -0,0 +1 @@
+# jelli - a JAX-based EFT Likelihood

jelli-0.0.1/pyproject.toml

@@ -0,0 +1,32 @@
+[build-system]
+requires = ["setuptools >= 61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "jelli"
+authors = [
+    {name = "Aleks Smolkovič", email = "aleks.smolkovic@ijs.si"},
+    {name = "Peter Stangl", email = "peter.stangl@cern.ch"},
+]
+description = "A JAX-based EFT likelihood."
+readme = "README.md"
+requires-python = ">=3.8"
+license = {text = "MIT"}
+dependencies = [
+    "numpy",
+    "scipy",
+    "jax",
+    "wilson",
+    "rgevolve",
+    "multipledispatch",
+    "networkx",
+    "h5py",
+]
+dynamic = ["version"]
+
+[tool.setuptools.packages.find]
+where = ["src/"]
+
+[tool.setuptools.dynamic]
+version = {attr = "jelli._version.__version__"}
+

jelli-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

jelli-0.0.1/setup.py

@@ -0,0 +1,3 @@
+from setuptools import setup
+
+setup()

jelli-0.0.1/src/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-0.0.1/src/jelli/_version.py

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

jelli-0.0.1/src/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-0.0.1/src/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