API

Regressors

Utility classes and functions related to gresit.

Copyright (c) 2025 Robert Bosch GmbH

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.

BackgroundKnowledge

Class to store background knowledge.

Source code in gresit/background_knowledge.py

class BackgroundKnowledge:
    """Class to store background knowledge."""

    def __init__(self, full_data_dict: dict[str, np.ndarray]) -> None:
        """Initializes the BK object.

        Args:
            full_data_dict (dict[str, np.ndarray]): input data without BK.
        """
        self.final_data = full_data_dict

    def target_overwrite(self, target_dict: dict[str, np.ndarray]) -> None:
        """Restrictions on target variable."""
        for target_name, target_values in target_dict.items():
            self.final_data[target_name] = target_values

    def target_remove_by_index(self, target_dict: dict[str, np.ndarray]) -> None:
        """Takes array columns in target key, value pair and removes them.

        Removal is done by (multiple) index array.

        Args:
            target_dict (dict[str, np.ndarray]): Dict with key equal to groups
                and values equal to indices.
        """
        for target_name, target_indices in target_dict.items():
            self.final_data[target_name] = np.delete(
                arr=self.final_data[target_name], obj=target_indices, axis=1
            )

__init__(full_data_dict)

Initializes the BK object.

Parameters:

Name	Type	Description	Default
full_data_dict	dict[str, ndarray]	input data without BK.	required

Source code in gresit/background_knowledge.py

def __init__(self, full_data_dict: dict[str, np.ndarray]) -> None:
    """Initializes the BK object.

    Args:
        full_data_dict (dict[str, np.ndarray]): input data without BK.
    """
    self.final_data = full_data_dict

target_overwrite(target_dict)

Restrictions on target variable.

Source code in gresit/background_knowledge.py

def target_overwrite(self, target_dict: dict[str, np.ndarray]) -> None:
    """Restrictions on target variable."""
    for target_name, target_values in target_dict.items():
        self.final_data[target_name] = target_values

target_remove_by_index(target_dict)

Takes array columns in target key, value pair and removes them.

Removal is done by (multiple) index array.

Parameters:

Name	Type	Description	Default
target_dict	dict[str, ndarray]	Dict with key equal to groups and values equal to indices.	required

Source code in gresit/background_knowledge.py

def target_remove_by_index(self, target_dict: dict[str, np.ndarray]) -> None:
    """Takes array columns in target key, value pair and removes them.

    Removal is done by (multiple) index array.

    Args:
        target_dict (dict[str, np.ndarray]): Dict with key equal to groups
            and values equal to indices.
    """
    for target_name, target_indices in target_dict.items():
        self.final_data[target_name] = np.delete(
            arr=self.final_data[target_name], obj=target_indices, axis=1
        )

Utility classes and functions related to gresit.

Copyright (c) 2025 Robert Bosch GmbH

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.

GroupPC

Bases: LearnAlgo

This class provides tools for causal discovery.

Particularly, in the context where data is known to follow a layered structure.

Source code in gresit/group_pc.py

class GroupPC(LearnAlgo):
    """This class provides tools for causal discovery.

    Particularly, in the context where data is known to follow
    a layered structure.
    """

    def __init__(self, alpha: float = 0.05, test: type[CItest] = FisherZVec) -> None:
        """Initiates VectorPC.

        Args:
            alpha (float, optional): Acts as a tuning parameter. The significance
                threshold for the conditional independence test. The smaller,
                the sparser the resulting graph. Defaults to 0.05.
            test (CItest, optional): Which CI test to use.
        """
        self.layering: dict[str, list[str]] | None = None
        self.alpha: float = alpha
        self.pdag: PDAG = PDAG()
        self.ambiguities: list[tuple[str, str, str]] = []
        self.skel: pd.DataFrame
        self.ci_test = test()

    def learn_graph(
        self,
        data_dict: dict[str, np.ndarray],
        threshold: float = 0.5,
        layering: dict[str, list[str]] | None = None,
    ) -> PDAG:
        """Learns the graph from the given data.

        If layering is provided it is taken to be unambiguous.
        If layering is not Null, then the separation sets may never
        contain variables that appear in future layers to the pair of
        variables considered.

        Args:
            data_dict (dict | ndarray): relevant data.
            threshold (float, optional): The majority vote threshold for deciding
                on ambiguous collider structures. Defaults to 0.5.
            layering (dict[str, list[str]], optional): The layering of the nodes.

        Returns:
            PDAG: Graph estimate.

        """
        self.layering = layering

        self._find_skeleton(data=data_dict, alpha=self.alpha, layering=layering)
        self.maximally_orient(data=data_dict, alpha=self.alpha, threshold=threshold)
        return self.pdag

    def _find_skeleton(
        self,
        data: dict[str, np.ndarray],
        alpha: float = 0.05,
        layering: dict[str, list[str]] | None = None,
    ) -> None:
        """First Phase of layered PC (stable) algorithm.

        If layering is not Null, then the separation sets may never
        contain variables that appear in future layers to the pair of
        variables considered.

        Args:
            data (dict[str, np.ndarray]): _description_
            alpha (float, optional): _description_. Defaults to 0.05.
            layering (dict[str, list[str]], optional): _description_. Defaults to None.
        """
        n_features = len(data)
        node_names = list(data.keys())
        skeleton = pd.DataFrame(
            (np.ones((n_features, n_features)) - np.eye(n_features)),
            columns=node_names,
            index=node_names,
        )

        nodes = sorted(node_names)
        sep_set: dict[int, dict[tuple[str, str], list[str]]] = defaultdict(dict)
        d = -1
        node_sets = list(combinations(nodes, 2))
        node_sets = sorted(node_sets, key=lambda x: x[0])

        while self._adj_size_criterion(skeleton, d):  # until for each adj(C,i)\{j} < l
            d += 1
            c_stable = deepcopy(skeleton)  # needed for stable
            if d not in sep_set:
                sep_set[d] = {}
            for node_i, node_j in node_sets:
                if skeleton.loc[node_i, node_j] == 0:
                    continue
                adj_i = set(c_stable.index[c_stable[node_i] == 1].to_list())
                z = adj_i - {node_j}  # adj(C, i)\{j}
                if len(z) >= d:
                    # |adj(C, i)\{j}| >= l
                    z_list = sorted([*z])
                    k = sorted([*combinations(z_list, d)], reverse=True)
                    for subset in k:
                        sub_z = list(subset)
                        if layering is not None and sub_z:
                            max_layer = self._return_max_layer(
                                layering=layering, node_i=node_i, node_j=node_j
                            )
                            allowed_dict = self._remove_pairs_after_key(
                                layering=layering, max_layer=max_layer
                            )
                            sub_z = self._remove_future_nodes(
                                sep_set=sub_z,
                                allowed_dict=allowed_dict,
                            )
                        if not sub_z:
                            _, p_value = self.ci_test.test(x_data=data[node_i], y_data=data[node_j])
                        else:
                            # find highest layer of node_{i,j} and restrict sep
                            # set such that no nodes can be in layers further
                            # than said highest layer.
                            _, p_value = self.ci_test.test(
                                x_data=data[node_i],
                                y_data=data[node_j],
                                z_data=np.concatenate(
                                    [dat for node, dat in data.items() if node in sub_z], axis=1
                                ),
                            )

                        if p_value >= alpha:
                            skeleton.loc[node_i, node_j] = skeleton.loc[node_j, node_i] = 0
                            sep_set[d][(node_i, node_j)] = sub_z
                            break

        self.skel = skeleton
        self.pdag = PDAG.from_pandas_adjacency(skeleton)

    def _return_max_layer(self, layering: dict[str, list[str]], node_i: str, node_j: str) -> str:
        """Highest layer of the pair of nodes considered.

        Args:
            layering (dict[str, list[str]]): given layering
            node_i (str): first node in the pair
            node_j (str): second node in the pair

        Returns:
            str: key corresponding to highest layer.
        """
        for key, value in reversed(layering.items()):
            if node_i in value or node_j in value:
                max_layer = key
                break
        return max_layer

    def _remove_pairs_after_key(
        self, layering: dict[str, list[str]], max_layer: str
    ) -> dict[str, list[str]]:
        items = list(layering.items())
        # Find the index of the target key
        try:
            target_index = next(i for i, (key, _) in enumerate(items) if key == max_layer)
        except StopIteration:
            # If the target key is not found, return the original dictionary
            return layering

        # Recreate the dictionary up to and including the target key
        return dict(items[: target_index + 1])

    def _remove_future_nodes(
        self, sep_set: list[str], allowed_dict: dict[str, list[str]]
    ) -> list[str]:
        """Remove selected nodes in allowed_dict from sep_set.

        Args:
            sep_set (list[str]): sep set
            allowed_dict (dict[str, list[str]]): allowed dict

        Returns:
            list[str]: list with future nodes removed.
        """
        dict_values = {item for value in allowed_dict.values() for item in value}
        return [item for item in sep_set if item in dict_values]

    def _unshielded_triples(self, pdag: PDAG) -> list[tuple[str, str, str]]:
        """Unshielded triples of a given PDAG.

        Args:
            pdag (PDAG): Skeleton after first phase

        Returns:
            list[tuple[str]]: List of unshielded triples of the form (i,j,k).
        """
        unshielded_triples: list[tuple[str, str, str]] = []
        for node in set(pdag.nodes):
            neighbors = pdag.neighbors(node=node)
            for neighbor in neighbors:
                distant_neighbors = list(pdag.neighbors(neighbor) - {node})
                if distant_neighbors:
                    unshielded_triples.extend(
                        [
                            (node, neighbor, dist_neigh)
                            for dist_neigh in distant_neighbors
                            if (node, dist_neigh) not in pdag.undir_edges
                            and (dist_neigh, node) not in pdag.undir_edges
                        ]
                    )
        unique_triples: list[tuple[str, str, str]] = []
        for tri in unshielded_triples:
            equivalent_tri = tri[::-1]
            if equivalent_tri not in unique_triples:
                unique_triples.append(tri)
        return unique_triples

    def _power_set(self, a_set: set[str]) -> set[frozenset[str]]:
        """Return power set.

        Args:
            a_set (set): Input set.

        Returns:
            frozenset[str]: all combinations of set members.
        """
        length = len(a_set)
        power_set: set[frozenset[str]] = {
            frozenset({e for e, b in zip(a_set, f"{i:{length}b}") if b == "1"})
            for i in range(2**length)
        }
        return power_set

    def _get_conditioning_sets(
        self, triples: tuple[str, str, str], pdag: PDAG
    ) -> set[frozenset[str]]:
        """Return conditioning set.

        Args:
            triples (tuple): triples
            pdag (PDAG): current pdag

        Returns:
            set[tuple[str]]: conditioning set from tuples and pdag.
        """
        set_a = self._power_set(pdag.neighbors(node=triples[0]))
        set_b = self._power_set(pdag.neighbors(node=triples[2]))

        unique_cond_set: set[frozenset[str]] = set_a.union(set_b)
        return unique_cond_set

    def _orient_vstructs_and_flag_amgiguities(
        self,
        data: dict[str, np.ndarray],
        alpha: float = 0.05,
        threshold: float = 0.5,
    ) -> None:
        """Given data orient all unshielded triples to vstrutures if possible.

        Args:
            data (dict[str, np.ndarray]): data
            alpha (float, optional): Significance level of test. Defaults to 0.05.
            threshold (float, optional): Threshold for ambiguity condition. Defaults to 0.5.

        Raises:
            AssertionError: If skeleton changes due to operation error is thrown.
        """
        pdag = self.pdag.copy()
        all_unshielded_triples = self._unshielded_triples(pdag=pdag)
        flag = []
        for triple in all_unshielded_triples:
            node_i, node_j, node_k = triple
            cond_sets = self._get_conditioning_sets(
                triples=triple, pdag=pdag
            )  # getting all adj(X_i) and adj(X_k)

            conditioning_subsetter = []  # initiate candidate conditioning sets

            for cond_set in list(cond_sets):
                sep_set = list(cond_set)
                if not sep_set:
                    _, p_value = self.ci_test.test(x_data=data[node_i], y_data=data[node_k])
                else:
                    _, p_value = self.ci_test.test(
                        x_data=data[node_i],
                        y_data=data[node_k],
                        z_data=np.concatenate(
                            [dat for node, dat in data.items() if node in sep_set], axis=1
                        ),
                    )
                if p_value >= alpha:  # if p_value is large then test statistic is small i.e.
                    # the Null of (cond) independence cannot be rejected.
                    conditioning_subsetter.append(
                        sep_set
                    )  # all subsets that give us cond. independence

            if sum([node_j in sep for sep in conditioning_subsetter]) == threshold * len(
                conditioning_subsetter
            ):
                flag.append(triple)
            elif sum([node_j in sep for sep in conditioning_subsetter]) < threshold * len(
                conditioning_subsetter
            ):
                pdag.undir_to_dir_edge(tail=node_i, head=node_j)
                pdag.undir_to_dir_edge(tail=node_k, head=node_j)

        original_skeleton = nx.from_pandas_adjacency(pdag.adjacency_matrix, create_using=nx.Graph)

        if not self.skeleton.equals(nx.to_pandas_adjacency(original_skeleton)):
            raise AssertionError(
                "Skeleton has changed. This shouldn't be possible. Check your inputs!"
            )
        self.pdag = pdag
        self.ambiguities = flag

    def maximally_orient(
        self, data: dict[str, np.ndarray], alpha: float = 0.05, threshold: float = 0.5
    ) -> None:
        """Given a skeleton, the following orientation steps are taken.

            1. All undirected edges between layers are immediately oriented
                according to the given layering.
            2. Potential v-structures are ordiented.
            3. The remaining undirected edges are oriented according to the four
                Meek rules.

        Args:
            data (dict[str, np.ndarray]): The data.
            alpha (float, optional): The significance
                threshold for the conditional independence test. Defaults to 0.05.
            threshold (float, optional): The majority vote threshold for deciding
                on ambiguous collder structures. Defaults to 0.5.
        """
        # orient immediately according to layering if present
        if self.layering is not None:
            self._orient_between_layers()
        # Orient v-structures
        self._orient_vstructs_and_flag_amgiguities(data=data, alpha=alpha, threshold=threshold)
        # Apply Meek Rules
        self._orient_according_to_meek_rules()

    def _adj_size_criterion(self, skel: pd.DataFrame, ell: int) -> bool:
        r"""Check if |adj(C, X_i) \\ {X_j}| >= l for every pair of adjacent vertices in C.

        Args:
            skel (pd.DataFrame): Skeleton C
            ell (int): size of separating sets

        Returns:
            bool: True if size of adjacency set is larger or equal l and False else.
        """
        assert skel.shape[0] == skel.shape[1]
        columns = skel.columns
        columns = sorted(columns)
        k = list(combinations(columns, 2))
        sorted(k, reverse=True)
        node_pairs = [(x, y) for x, y in k]
        less_l = 0
        for node_i, node_j in node_pairs:
            adj_i = set(skel.index[skel[node_i] != 0].tolist())
            adj_ij = adj_i - {node_j}
            if len(adj_ij) < ell:
                less_l += 1
            else:
                break
        if less_l == len(node_pairs):
            return False
        else:
            return True

    def _between_edges(self) -> list[tuple[str, str]]:
        """Return between edges when layers are present.

        Returns:
            list[tuple]: list of edges between layers.
        """
        if self.layering is None:
            raise ValueError("Layering is not provided. Cannot retrieve between edges.")
        no_orient = []
        undir_edges = self.pdag.undir_edges
        for edge in undir_edges:
            i, j = edge
            for cell_nodes in self.layering.values():
                if set([i, j]).issubset(set(cell_nodes)):
                    no_orient.append(edge)
                    break
        return list(set(undir_edges).difference(set(no_orient)))

    def _orient_between_layers(self) -> None:
        """Orients edges between layers."""
        if self.layering is None:
            raise ValueError("Layering is not provided. Cannot orient between edges.")
        to_orient = self._between_edges()

        flat_mapper: list[str] = []
        for nodelist in self.layering.values():
            flat_mapper.extend(nodelist)

        final_edge_list = []
        for edge_to_orient in to_orient:
            sorted_edge = sorted(edge_to_orient, key=flat_mapper.index)
            final_edge_list.append(tuple(sorted_edge))

        for tail, head in final_edge_list:
            self.pdag.undir_to_dir_edge(tail=tail, head=head)

    def _orient_according_to_meek_rules(self) -> None:
        """Orient edges according to Meek rules."""
        cpdag = self.pdag.copy()
        cpdag = rule_1(pdag=cpdag)
        cpdag = rule_2(pdag=cpdag)
        cpdag = rule_3(pdag=cpdag)
        cpdag = rule_4(pdag=cpdag)
        self.pdag = cpdag

    @property
    def skeleton(self) -> pd.DataFrame:
        """Represent the underlying skeleton as adjacency matrix.

        Returns:
            pd.DataFrame: Adjacency matrix of the skeleton.
        """
        undir_graph = nx.from_pandas_adjacency(self.pdag.adjacency_matrix, create_using=nx.Graph)
        return nx.to_pandas_adjacency(undir_graph)

    @property
    def adjacency_matrix(self) -> pd.DataFrame:
        """Represent the underlying learned PDAG as adjacency matrix.

        Returns:
            pd.DataFrame: Adjacency matrix of the PDAG.
        """
        amat = self.pdag.adjacency_matrix.values
        cpdag_amat = amat.copy()
        upper_triangle_indices = np.triu_indices_from(amat, k=1)
        mask = (amat[upper_triangle_indices] == 1) & (amat[upper_triangle_indices[::-1]] == 1)

        # Set these entries to 2 in both (i, j) and (j, i) locations
        cpdag_amat[upper_triangle_indices[0][mask], upper_triangle_indices[1][mask]] = 2
        cpdag_amat[upper_triangle_indices[1][mask], upper_triangle_indices[0][mask]] = 2

        amat_names = self.pdag.adjacency_matrix.columns

        return pd.DataFrame(cpdag_amat, columns=amat_names, index=amat_names)

    @property
    def causal_order(self) -> list[str] | None:
        """Returns causal order if PDAG is in fact a DAG.

        Else it will return None.

        Returns:
            list[str] | None: causal order if appropriate.
        """
        ordering = None
        if self.pdag.dir_edges and not self.pdag.undir_edges:
            dag = DAG(nodes=self.pdag.nodes, edges=self.pdag.dir_edges)
            ordering = dag.causal_order
        return ordering

adjacency_matrix property

Represent the underlying learned PDAG as adjacency matrix.

Returns:

Type	Description
DataFrame	pd.DataFrame: Adjacency matrix of the PDAG.

causal_order property

Returns causal order if PDAG is in fact a DAG.

Else it will return None.

Returns:

Type	Description
list[str] | None	list[str] | None: causal order if appropriate.

skeleton property

Represent the underlying skeleton as adjacency matrix.

Returns:

Type	Description
DataFrame	pd.DataFrame: Adjacency matrix of the skeleton.

__init__(alpha=0.05, test=FisherZVec)

Initiates VectorPC.

Parameters:

Name	Type	Description	Default
alpha	float	Acts as a tuning parameter. The significance threshold for the conditional independence test. The smaller, the sparser the resulting graph. Defaults to 0.05.	0.05
test	CItest	Which CI test to use.	FisherZVec

Source code in gresit/group_pc.py

def __init__(self, alpha: float = 0.05, test: type[CItest] = FisherZVec) -> None:
    """Initiates VectorPC.

    Args:
        alpha (float, optional): Acts as a tuning parameter. The significance
            threshold for the conditional independence test. The smaller,
            the sparser the resulting graph. Defaults to 0.05.
        test (CItest, optional): Which CI test to use.
    """
    self.layering: dict[str, list[str]] | None = None
    self.alpha: float = alpha
    self.pdag: PDAG = PDAG()
    self.ambiguities: list[tuple[str, str, str]] = []
    self.skel: pd.DataFrame
    self.ci_test = test()

learn_graph(data_dict, threshold=0.5, layering=None)

Learns the graph from the given data.

If layering is provided it is taken to be unambiguous. If layering is not Null, then the separation sets may never contain variables that appear in future layers to the pair of variables considered.

Parameters:

Name	Type	Description	Default
data_dict	dict | ndarray	relevant data.	required
threshold	float	The majority vote threshold for deciding on ambiguous collider structures. Defaults to 0.5.	0.5
layering	dict[str, list[str]]	The layering of the nodes.	None

Returns:

Name	Type	Description
PDAG	PDAG	Graph estimate.

Source code in gresit/group_pc.py

def learn_graph(
    self,
    data_dict: dict[str, np.ndarray],
    threshold: float = 0.5,
    layering: dict[str, list[str]] | None = None,
) -> PDAG:
    """Learns the graph from the given data.

    If layering is provided it is taken to be unambiguous.
    If layering is not Null, then the separation sets may never
    contain variables that appear in future layers to the pair of
    variables considered.

    Args:
        data_dict (dict | ndarray): relevant data.
        threshold (float, optional): The majority vote threshold for deciding
            on ambiguous collider structures. Defaults to 0.5.
        layering (dict[str, list[str]], optional): The layering of the nodes.

    Returns:
        PDAG: Graph estimate.

    """
    self.layering = layering

    self._find_skeleton(data=data_dict, alpha=self.alpha, layering=layering)
    self.maximally_orient(data=data_dict, alpha=self.alpha, threshold=threshold)
    return self.pdag

maximally_orient(data, alpha=0.05, threshold=0.5)

Given a skeleton, the following orientation steps are taken.

1. All undirected edges between layers are immediately oriented
    according to the given layering.
2. Potential v-structures are ordiented.
3. The remaining undirected edges are oriented according to the four
    Meek rules.

Parameters:

Name	Type	Description	Default
data	dict[str, ndarray]	The data.	required
alpha	float	The significance threshold for the conditional independence test. Defaults to 0.05.	0.05
threshold	float	The majority vote threshold for deciding on ambiguous collder structures. Defaults to 0.5.	0.5

Source code in gresit/group_pc.py

def maximally_orient(
    self, data: dict[str, np.ndarray], alpha: float = 0.05, threshold: float = 0.5
) -> None:
    """Given a skeleton, the following orientation steps are taken.

        1. All undirected edges between layers are immediately oriented
            according to the given layering.
        2. Potential v-structures are ordiented.
        3. The remaining undirected edges are oriented according to the four
            Meek rules.

    Args:
        data (dict[str, np.ndarray]): The data.
        alpha (float, optional): The significance
            threshold for the conditional independence test. Defaults to 0.05.
        threshold (float, optional): The majority vote threshold for deciding
            on ambiguous collder structures. Defaults to 0.5.
    """
    # orient immediately according to layering if present
    if self.layering is not None:
        self._orient_between_layers()
    # Orient v-structures
    self._orient_vstructs_and_flag_amgiguities(data=data, alpha=alpha, threshold=threshold)
    # Apply Meek Rules
    self._orient_according_to_meek_rules()

MicroPC

Bases: LearnAlgo

Standard PC stable on micro nodes aggregated after the fact.

Source code in gresit/group_pc.py

class MicroPC(LearnAlgo):
    """Standard PC stable on micro nodes aggregated after the fact."""

    def __init__(self, alpha: float = 0.05) -> None:
        """Inits the object.

        Args:
            alpha (float, optional): Significance level of the test. Defaults to 0.05.
        """
        self.alpha = alpha
        self.graph: DAG | PDAG

    def _causallearn2amat(self, causal_learn_graph: np.ndarray) -> np.ndarray:
        amat = np.zeros(causal_learn_graph.shape)
        for col in range(causal_learn_graph.shape[1]):
            for row in range(causal_learn_graph.shape[0]):
                if causal_learn_graph[row, col] == -1 and causal_learn_graph[col, row] == 1:
                    amat[row, col] = 1
                if causal_learn_graph[row, col] == -1 and causal_learn_graph[col, row] == -1:
                    amat[row, col] = amat[col, row] = 1
                if causal_learn_graph[row, col] == 1 and causal_learn_graph[col, row] == 1:
                    amat[row, col] = amat[col, row] = 1
        return amat

    def learn_graph(self, data_dict: dict[str, np.ndarray], *args: Any, **kwargs: Any) -> GRAPH:
        """Learn graph.

        Args:
            data_dict (_type_): _description_
            *args (Any): additional args.
            **kwargs (Any): additional kwargs.

        Returns:
            PDAG: _description_
        """
        micro_data = np.concatenate([d_data for d_data in data_dict.values()], axis=1)
        micro_pc = pc(
            data=micro_data, alpha=self.alpha, indep_test="fisherz", uc_rule=1, show_progress=True
        )
        micro_amat = self._causallearn2amat(causal_learn_graph=micro_pc.G.graph)
        Q = _make_group_mapping(data_dict=data_dict)
        interim_group_adjacency_matrix = Q @ micro_amat @ Q.T
        np.fill_diagonal(interim_group_adjacency_matrix, 0)
        group_adjacency_matrix = (interim_group_adjacency_matrix > 0).astype(int)

        group_graph: DAG | PDAG
        if not np.any((group_adjacency_matrix == 1) & (group_adjacency_matrix.T == 1)):
            group_graph = DAG.from_pandas_adjacency(
                pd.DataFrame(
                    group_adjacency_matrix, columns=data_dict.keys(), index=data_dict.keys()
                )
            )
        else:
            group_graph = PDAG.from_pandas_adjacency(
                pd.DataFrame(
                    group_adjacency_matrix, columns=data_dict.keys(), index=data_dict.keys()
                )
            )

        self.graph = group_graph
        return group_graph

    @property
    def causal_order(self) -> list[str] | None:
        """Causal order."""
        if isinstance(self.graph, DAG):
            return self.graph.causal_order
        else:
            return None

    @property
    def adjacency_matrix(self) -> pd.DataFrame:
        """Adjacency matrix."""
        return self.graph.adjacency_matrix

adjacency_matrix property

Adjacency matrix.

causal_order property

Causal order.

__init__(alpha=0.05)

Inits the object.

Parameters:

Name	Type	Description	Default
alpha	float	Significance level of the test. Defaults to 0.05.	0.05

Source code in gresit/group_pc.py

def __init__(self, alpha: float = 0.05) -> None:
    """Inits the object.

    Args:
        alpha (float, optional): Significance level of the test. Defaults to 0.05.
    """
    self.alpha = alpha
    self.graph: DAG | PDAG

learn_graph(data_dict, *args, **kwargs)

Learn graph.

Parameters:

Name	Type	Description	Default
data_dict	_type_	description	required
*args	Any	additional args.	()
**kwargs	Any	additional kwargs.	{}

Returns:

Name	Type	Description
PDAG	GRAPH	description

Source code in gresit/group_pc.py

def learn_graph(self, data_dict: dict[str, np.ndarray], *args: Any, **kwargs: Any) -> GRAPH:
    """Learn graph.

    Args:
        data_dict (_type_): _description_
        *args (Any): additional args.
        **kwargs (Any): additional kwargs.

    Returns:
        PDAG: _description_
    """
    micro_data = np.concatenate([d_data for d_data in data_dict.values()], axis=1)
    micro_pc = pc(
        data=micro_data, alpha=self.alpha, indep_test="fisherz", uc_rule=1, show_progress=True
    )
    micro_amat = self._causallearn2amat(causal_learn_graph=micro_pc.G.graph)
    Q = _make_group_mapping(data_dict=data_dict)
    interim_group_adjacency_matrix = Q @ micro_amat @ Q.T
    np.fill_diagonal(interim_group_adjacency_matrix, 0)
    group_adjacency_matrix = (interim_group_adjacency_matrix > 0).astype(int)

    group_graph: DAG | PDAG
    if not np.any((group_adjacency_matrix == 1) & (group_adjacency_matrix.T == 1)):
        group_graph = DAG.from_pandas_adjacency(
            pd.DataFrame(
                group_adjacency_matrix, columns=data_dict.keys(), index=data_dict.keys()
            )
        )
    else:
        group_graph = PDAG.from_pandas_adjacency(
            pd.DataFrame(
                group_adjacency_matrix, columns=data_dict.keys(), index=data_dict.keys()
            )
        )

    self.graph = group_graph
    return group_graph

Utility classes and functions related to gresit.

Copyright (c) 2025 Robert Bosch GmbH

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.

GroupResit

Bases: LearnAlgo

A class representing the groupResit algorithm.

This algorithm is used to learn a DAG based on vector/group valued ANMs.

Source code in gresit/group_resit.py

class GroupResit(LearnAlgo):
    """A class representing the groupResit algorithm.

    This algorithm is used to learn a DAG based on vector/group valued ANMs.

    """

    def __init__(
        self,
        regressor: MultiRegressor,
        test: type[Itest],
        alpha: float = 0.01,
        pruning_method: str = "murgs",
        test_size: float = 0.2,
        local_regression_method: str = "kernel",
    ) -> None:
        """Initialize the GroupResit object.

        Args:
            regressor (MultivariateRegressor): A regressor object.
            test (IndependenceTest): An independence test object.
            alpha (float): Alpha
            pruning_method (str): The pruning method
            test_size (float): Relative size of test-dataset, 0 means no test data
            local_regression_method (str): Type of local linear smoother to use. Options are
                `loess`, `kernel`, soon to be implemented `spline`. Defaults to `kernel`.
        """
        self.regressor = regressor
        self.independence_test = test()
        self.local_regression_method = local_regression_method
        self._pa: dict[str, list[str]] = {}
        self._pa_history: dict[str, dict[float, list[str]]] = {}
        self.alpha_level: float
        self.alpha = alpha
        self.pruning_method = pruning_method
        self.test_size = test_size
        self.DAG: DAG
        self.layering: dict[str, list[str]] = {}
        super().__init__()

    def __repr__(self) -> str:
        """Repr method.

        Returns:
            str: Description of the object.
        """
        return f"GroupResit(regressor={self.regressor}, independence_test={self.independence_test})"

    def __str__(self) -> str:
        """Str method.

        Returns:
            str: Human-readable description of the object.
        """
        method_description = {
            "Regression method: ": self.regressor.__class__.__name__,
            "Independece test:": self.independence_test.__class__.__name__,
            "Inferred causal order: ": "Yes" if self._causal_order else "Not yet",
        }
        s = ""
        for info, info_text in method_description.items():
            s += f"{info:<14}{info_text:>5}\n"

        return s

    def _get_causal_order(
        self,
        data_dict: dict[str, np.ndarray],
        layering: dict[str, list[str]] | None = None,
    ) -> None:
        """Get causal order of the groups respecting the current layering.

        Args:
            data_dict (dict[str, np.ndarray]): A dictionary of np.ndarrays. Key corresponds to
                group name and values to the corresponding data.
            layering (dict[str, list[str]]): A dictionary of layering information. Keys correspond
                to the layer and values to the variable names within each layer.
        """
        if layering is None:
            layering = {"L": list(data_dict.keys())}

        self.layering = layering
        pa: dict[str, list[str]] = {}
        pi: list[str] = []
        data_to_delete_from = data_dict.copy()

        indices = np.arange(data_dict[list(data_dict.keys())[0]].shape[0])
        if self.test_size > 0:
            idx_train, idx_test = train_test_split(
                indices,
                test_size=self.test_size,
                random_state=2024,
            )
        else:
            idx_test = indices
            idx_train = indices

        self._idx_test = idx_test

        for _, vars in reversed(layering.items()):
            if pa:  # add previous layer nodes to current layer parents
                for _, pre_layer_parents in pa.items():
                    pre_layer_parents.extend(vars)

            within_layer_order = vars.copy()
            for _ in vars:
                if len(within_layer_order) == 1:  # in each layer if there's only one node left,
                    # this must be the first in the causal ordering
                    pa[within_layer_order[0]] = []
                    pi.insert(0, within_layer_order[0])
                    del data_to_delete_from[within_layer_order[0]]
                    continue

                test_stats: list[np.float64] = []
                for var in within_layer_order:
                    Y = data_to_delete_from[var].copy()  # remove columns from data!
                    X = np.concatenate(
                        [_d for _group, _d in data_to_delete_from.items() if _group != var],
                        axis=1,
                    )

                    X_train = self._standardize(X[idx_train])
                    Y_train = self._standardize(Y[idx_train])

                    X_test = (X[idx_test] - X[idx_train].mean(axis=0)) / (X[idx_train].std(axis=0))
                    Y_test = (Y[idx_test] - Y[idx_train].mean(axis=0)) / (Y[idx_train].std(axis=0))

                    self.regressor.fit(X=X_train, Y=Y_train)
                    Y_pred = self.regressor.predict(X_test)

                    residuals = Y_test - np.squeeze(np.asarray(Y_pred))
                    residuals = self._standardize(residuals)
                    X_test = self._standardize(X_test)
                    test_stat, _ = self.independence_test.test(x_data=residuals, y_data=X_test)
                    test_stats.append(test_stat)

                k = within_layer_order[np.argmin(test_stats)]  # get most independent group
                within_layer_order.remove(k)  # remove from remaining order
                pa[k] = within_layer_order.copy()  # add all within layer potential parents
                pi.insert(0, k)  # prepend to causal ordering

                del data_to_delete_from[k]

        self._causal_order = pi
        self._pa = pa

    def _standardize(self, x: np.ndarray) -> np.ndarray:
        return (x - x.mean(axis=0)) / x.std(axis=0)

    def _get_causal_order_via_dependence_loss(
        self,
        data_dict: dict[str, np.ndarray],
        layering: dict[str, list[str]] | None = None,
    ) -> None:
        """Get causal order of the groups respecting the current layering.

        Args:
            data_dict (dict[str, np.ndarray]): A dictionary of np.ndarrays. Key corresponds to
                group name and values to the corresponding data.
            layering (dict[str, list[str]]): A dictionary of layering information. Keys correspond
                to the layer and values to the variable names within each layer.
        """
        if layering is None:
            layering = {"L": list(data_dict.keys())}

        self.layering = layering
        pa: dict[str, list[str]] = {}
        pi: list[str] = []
        data_to_delete_from = data_dict.copy()

        indices = np.arange(data_dict[list(data_dict.keys())[0]].shape[0])
        if self.test_size > 0:
            idx_train, idx_test = train_test_split(
                indices,
                test_size=self.test_size,
                random_state=2024,
            )
        else:
            idx_test = indices
            idx_train = indices

        self._idx_test = idx_test

        for _, vars in reversed(layering.items()):
            if pa:  # add previous layer nodes to current layer parents
                for _, pre_layer_parents in pa.items():
                    pre_layer_parents.extend(vars)

            within_layer_order = vars.copy()
            for _ in vars:
                if len(within_layer_order) == 1:  # in each layer if there's only one node left,
                    # this must be the first in the causal ordering
                    pa[within_layer_order[0]] = []
                    pi.insert(0, within_layer_order[0])
                    del data_to_delete_from[within_layer_order[0]]
                    continue

                final_loss_values: list[np.float64] = []
                for var in within_layer_order:
                    Y = data_to_delete_from[var].copy()  # remove columns from data!
                    X = np.concatenate(
                        [_d for _group, _d in data_to_delete_from.items() if _group != var],
                        axis=1,
                    )
                    if isinstance(self.regressor, Multioutcome_MLP):
                        self.regressor.fit(X=X, Y=Y, idx_train=idx_train, idx_test=idx_test)
                        final_loss_values.append(self.regressor.final_loss)
                    else:
                        raise ValueError("Independence loss only implemented in MLP.")

                k = within_layer_order[np.argmin(final_loss_values)]  # get most independent group
                within_layer_order.remove(k)  # remove from remaining order
                pa[k] = within_layer_order.copy()  # add all within layer potential parents
                pi.insert(0, k)  # prepend to causal ordering
                del data_to_delete_from[k]

        self._causal_order = pi
        self._pa = pa

    def _independence_prune(
        self,
        data_dict: dict[str, np.ndarray],
        alpha: float = 0.01,
    ) -> None:
        """Prune according to Peters et al. (2014).

        Args:
            data_dict (dict[str, np.ndarray]): Data in form of dict with groups as keys
                and data as values.
            alpha (float, optional): Significance level. Defaults to 0.01.
        """
        # pruning
        dat = data_dict.copy()
        pa = self._pa
        pi = self._causal_order
        idx_test = self._idx_test

        for k in pi:  # for each but the first node in the causal order
            # Take every parent and check for independence
            parents = pa[k].copy()
            if not parents:
                continue

            for parent in parents:
                Y = dat[k].copy()
                # see whether parent may be removed
                remainders = [var for var in parents if var != parent]
                if not remainders:
                    continue

                X = np.concatenate(
                    [dat[x] for x in remainders],
                    axis=1,
                )

                self.regressor.fit(
                    X=X[idx_test],
                    Y=Y[idx_test],
                )

                y_pred = self.regressor.predict(X[idx_test])
                residual = Y[idx_test] - np.squeeze(np.asarray(y_pred))
                _, decision = self.independence_test.test(x_data=residual, y_data=X[idx_test])
                if isinstance(decision, float):
                    if decision > alpha:  # if indepedence is not rejected, remove the edge
                        pa[k].remove(parent)
                elif isinstance(decision, str):
                    if decision == "Reject Null of vector independence: False":
                        pa[k].remove(parent)
                else:
                    raise ValueError("Test decision is neither float nor string")

        self.alpha_level = alpha
        self._pa = pa

    def _sparse_regression_pruning(
        self, data_dict: dict[str, np.ndarray], nlambda: int = 30
    ) -> None:
        # pruning
        dat = data_dict.copy()
        if self.test_size > 0:
            dat = {nodes: data for nodes, data in data_dict.items()}
            # data[self._idx_test]
        pa = self._pa
        pi = self._causal_order
        for k in pi:  # for each but the first node in the causal order
            # Take every parent and check for independence
            potential_parents = pa[k].copy()
            if not potential_parents:
                continue
            # get Y_data and X_data
            Y_data = dat[k]
            # Create dict
            X_data = {key: dat[key] for key in potential_parents if key in dat}
            # initiate MURGS object
            murgs = MURGS()
            murgs.fit(
                X_data=X_data,
                Y_data=Y_data,
                nlambda=nlambda,
                precalculate_smooths=True,
                local_regression_method=self.local_regression_method,
            )
            # extract zero groups
            zero_groups = murgs.zero_groups
            pa[k] = [parent for i, parent in enumerate(potential_parents) if not zero_groups[i]]
            self._pa_history[k] = {
                penalty: [
                    parent for i, parent in enumerate(potential_parents) if not zero_groups[i]
                ]
                for penalty, zero_groups in murgs.zero_group_history.items()
            }

        self._pa = pa

    def _dict_preprocessing(self, data_dict: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
        """Dict preprocessing dealing with univariate groups.

        Args:
            data_dict (dict[str, np.ndarray]): data dict.

        Returns:
            dict[str, np.ndarray]: data dict with axis added when univariate.
        """
        for key, value in data_dict.items():
            if value.ndim == 1:
                data_dict[key] = value[:, np.newaxis]
        return data_dict

    def learn_graph(
        self,
        data_dict: dict[str, np.ndarray],
        layering: dict[str, list[str]] | None = None,
    ) -> DAG:
        """Learn the causal graph.

        Args:
            data_dict (dict[str, np.ndarray]): A dictionary of np.ndarrays. Key corresponds to
                group name and values to the corresponding data.
            layering (dict[str, list[str]]): A dictionary of layering information. Keys correspond
                to the layer and values to the variable names within each layer.

        Raises:
            NotImplementedError: _description_

        Returns:
            DAG: DAG estimate.
        """
        clean_data_dict = self._dict_preprocessing(data_dict)
        self._get_causal_order(data_dict=clean_data_dict, layering=layering)
        if self.pruning_method == "murgs":
            self._sparse_regression_pruning(data_dict=clean_data_dict)
        elif self.pruning_method == "independence":
            self._independence_prune(data_dict=clean_data_dict, alpha=self.alpha)
        else:
            raise NotImplementedError()

        edge_list = [(parent, child) for child in self._pa for parent in self._pa[child]]

        learned_DAG = DAG(nodes=self._causal_order)
        learned_DAG.add_edges_from(edge_list)
        self._adjacency_matrix = learned_DAG.adjacency_matrix

        self.DAG = learned_DAG
        return learned_DAG

    def _insert_known_causal_order(self, pi: list[str]) -> None:
        """Insert a known causal order.

        Args:
            pi (list[str]): A list of group names in the causal order.
        """
        self._causal_order = pi
        pa = {}
        for node in pi:
            index = pi.index(node)
            pa[node] = pi[0:index]
        self._pa = pa

    def model_selection_with_known_causal_order(
        self,
        pi: list[str],
        data_dict: dict[str, np.ndarray],
        alpha: float = 0.01,
        pruning_method: str = "murgs",
    ) -> None:
        """Given a known causal order perform model selection.

        Args:
            pi (list[str]): Causal ordering. Entries in `pi` need
                to coincide with the keys in `data_dict`.
            data_dict (dict[str, np.ndarray]): A dictionary of np.ndarrays. Key corresponds to
                group name and values to the corresponding data.
            alpha (float, optional): The significance level for the independence test.
                Defaults to 0.1.
            pruning_method (str, optional): The pruning method to use. Defaults to "murgs".
                other options include `"independence"`

        """
        self._insert_known_causal_order(pi=pi)
        if not hasattr(self, "_idx_test"):
            self._idx_test = np.arange(data_dict[list(data_dict.keys())[0]].shape[0])
        if pruning_method == "murgs":
            self._sparse_regression_pruning(data_dict=data_dict)
        elif pruning_method == "independence":
            self._independence_prune(data_dict=data_dict, alpha=alpha)
        else:
            raise NotImplementedError()

        edge_list = [(parent, child) for child in self._pa for parent in self._pa[child]]

        learned_DAG = DAG(nodes=self._causal_order)
        learned_DAG.add_edges_from(edge_list)
        self._adjacency_matrix = learned_DAG.adjacency_matrix
        self.DAG = learned_DAG

    @property
    def causal_order(self) -> list[str] | None:
        """Causal order."""
        return self._causal_order

    @property
    def adjacency_matrix(self) -> pd.DataFrame:
        """Adjacency matrix."""
        return self._adjacency_matrix

    def _add_layered_layout_to_graph(self, nx_graph: nx.DiGraph) -> nx.DiGraph:
        """Add coordinate pos to nx.DiGraph.

        Args:
            nx_graph (nx.DiGraph): DAG in question.

        Returns:
            nx_graph (nx.DiGraph): DAG with attributes set.
        """
        if len(self.layering) > 1:
            for layer, nodes in enumerate(self.layering.values()):
                for node in nodes:
                    nx_graph.nodes[node]["layer"] = layer
        else:
            for layer, nodes in enumerate(nx.topological_generations(nx_graph)):
                for node in nodes:
                    nx_graph.nodes[node]["layer"] = layer

        # Plot multipartite_layout using the "layer" node attribute
        pos = nx.multipartite_layout(nx_graph, subset_key="layer")
        nx.set_node_attributes(G=nx_graph, name="pos", values=pos)
        return nx_graph

    def show(self, title: str = "Group RESIT DAG") -> None:
        """Plot the learned DAG.

        The plot is interactive,
        hovering over the nodes reveals the node labels.
        Colors get brighter the higher the node degree.

        Args:
            title (str, optional): Plot title. Defaults to "Group RESIT DAG".

        Raises:
            AssertionError: Throws error if DAG not yet learned.
        """
        if not hasattr(self, "DAG"):
            raise AssertionError("No graph to plot. Learn the graph first.")

        nx_dag = self.DAG.to_networkx()

        nx_dag = self._add_layered_layout_to_graph(nx_graph=nx_dag)

        edge_x = []
        edge_y = []
        for edge in nx_dag.edges():
            x0, y0 = nx_dag.nodes[edge[0]]["pos"]
            x1, y1 = nx_dag.nodes[edge[1]]["pos"]
            edge_x.append(x0)
            edge_x.append(x1)
            edge_x.append(None)
            edge_y.append(y0)
            edge_y.append(y1)
            edge_y.append(None)

        edge_trace = go.Scatter(
            x=edge_x,
            y=edge_y,
            line=dict(width=0.5, color="#888"),
            hoverinfo="none",
            mode="lines+markers",
            marker=dict(size=10, symbol="arrow-bar-up", angleref="previous"),
        )

        node_x = []
        node_y = []
        for node in nx_dag.nodes():
            x, y = nx_dag.nodes[node]["pos"]
            node_x.append(x)
            node_y.append(y)

        node_trace = go.Scatter(
            x=node_x,
            y=node_y,
            mode="markers",
            hoverinfo="text",
            marker=dict(
                showscale=True,
                # colorscale options
                #'Greys' | 'YlGnBu' | 'Greens' | 'YlOrRd' | 'Bluered' | 'RdBu' |
                #'Reds' | 'Blues' | 'Picnic' | 'Rainbow' | 'Portland' | 'Jet' |
                #'Hot' | 'Blackbody' | 'Earth' | 'Electric' | 'Viridis' |
                colorscale="YlGnBu",
                reversescale=True,
                color=[],
                size=10,
                colorbar=dict(thickness=15, title="node degree", xanchor="left", titleside="right"),
                line_width=2,
            ),
        )

        node_adjacencies = []
        node_text = []
        for node, adjacencies in enumerate(nx_dag.adjacency()):
            node_adjacencies.append(len(adjacencies[1]))
            node_text.append(list(nx_dag.nodes)[node])

        node_trace.marker.color = node_adjacencies
        node_trace.text = node_text

        fig = go.Figure(
            data=[edge_trace, node_trace],
            layout=go.Layout(
                title=title,
                # titlefont_size=16,
                showlegend=False,
                hovermode="closest",
                margin=dict(b=20, l=5, r=5, t=40),
                xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
            ),
        )
        fig.show()

    def _layer_pos(
        self, G: nx.DiGraph, layers: dict[str, list[str]], layer_gap: float = 8.0
    ) -> dict[str | int, np.ndarray]:
        pos = {}
        for i, nodes in enumerate(layers.values()):
            pos.update(
                nx.spring_layout(
                    G.subgraph(nodes), center=np.array([layer_gap * i, 0]), seed=42, k=50
                )
            )
        return pos

    def _create_edge_trace(
        self, G: nx.DiGraph, pos: dict[str | int, np.ndarray], highlight: bool = False
    ) -> go.Scatter:
        edge_x, edge_y = [], []

        if G.edges():
            for edge in G.edges():
                x0, y0 = pos[edge[0]]  # Source node
                x1, y1 = pos[edge[1]]  # Target node
                # Store edge coordinates
                edge_x.extend([x0, x1, None])  # Keep source unchanged
                edge_y.extend([y0, y1, None])

        # Define edge trace with customized width and color
        edge_trace = go.Scatter(
            x=edge_x,
            y=edge_y,
            line=dict(
                width=2 if highlight else 1,  # Thicker for first graph
                color="dimgray" if highlight else "#888",  # Darker first graph edges
            ),
            hoverinfo="none",
            mode="lines+markers",
            marker=dict(size=10, symbol="arrow-bar-up", angleref="previous"),
        )

        return edge_trace

    def _make_graph_list(self) -> list[nx.DiGraph]:
        unique_parent_history_sparser = {}
        unique_parent_history = {}

        for node, possible_parents in self._pa_history.items():
            unique_parent_history[node] = sorted(
                [
                    list(x)
                    for x in {
                        frozenset(sublist)
                        for sublist in [parents for parents in possible_parents.values()]
                    }
                    if len(x) >= len(self._pa[node])
                ],
                key=len,
            )
            unique_parent_history_sparser[node] = sorted(
                [
                    list(x)
                    for x in {
                        frozenset(sublist)
                        for sublist in [parents for parents in possible_parents.values()]
                    }
                    if len(x) < len(self._pa[node])
                ],
                key=len,
            )

        max_length = max([len(penalties) for penalties in unique_parent_history.values()])
        # Pad each list of lists
        padded_dict = unique_parent_history.copy()
        for node, _ in unique_parent_history.items():
            last_element = (
                unique_parent_history[node][-1] if unique_parent_history[node] else []
            )  # Last element of the list

            while len(unique_parent_history[node]) < max_length:
                padded_dict[node].append(last_element)

        max_length_sparser = max(
            [len(penalties) for penalties in unique_parent_history_sparser.values()]
        )
        # Pad each list of lists
        padded_dict_sparser = unique_parent_history_sparser.copy()
        for node, _ in unique_parent_history_sparser.items():
            last_element = (
                unique_parent_history_sparser[node][-1]
                if unique_parent_history_sparser[node]
                else []
            )  # Last element of the list

            while len(unique_parent_history_sparser[node]) < max_length_sparser:
                padded_dict_sparser[node].append(last_element)

        combined_dict = {}
        for key in padded_dict.keys():
            combined_dict[key] = padded_dict_sparser[key] + padded_dict[key]

        graph_list = []
        for i in range(max_length_sparser + max_length):
            pa = {node: parents[i] for node, parents in combined_dict.items()}
            edge_list = [(parent, child) for child in pa for parent in pa[child]]
            learned_DAG = nx.DiGraph()
            learned_DAG.add_nodes_from(list(pa.keys()))
            learned_DAG.add_edges_from(edge_list)
            graph_list.append(learned_DAG)

        return graph_list

    def show_interactive(self, layer_gap: float = 8.0) -> go.Figure:
        """Show interactive plot with slider to select sparsity level.

        Args:
            layer_gap (float, optional): gap between layers when displaying. Defaults to 8.0.
        """
        # Fixed layout for consistency
        graph_list = self._make_graph_list()
        pos = self._layer_pos(G=graph_list[-1], layers=self.layering, layer_gap=layer_gap)
        nodes = list(graph_list[-1].nodes)

        # Extract node coordinates
        node_x, node_y = zip(*[pos[n] for n in nodes])

        node_trace = go.Scatter(
            x=node_x,
            y=node_y,
            mode="markers",
            text=nodes,
            hoverinfo="text",
            marker=dict(size=15, color="lightblue", line=dict(color="black", width=1)),
        )

        # Identify the "optimal" graph
        optimal_graph_id = next(
            (
                i
                for i, graph in enumerate(graph_list)
                if nx.utils.graphs_equal(self.DAG.to_networkx(), graph)
            ),
            None,
        )
        if optimal_graph_id is not None:
            # Create the optimal graph's edge trace
            optimal_edge_trace = self._create_edge_trace(
                G=graph_list[optimal_graph_id], pos=pos, highlight=True
            )
        else:
            raise ValueError("Optimal graph not found in graph list.")

        # Create frames for each graph
        frames = []
        num_frames = len(graph_list)

        for i, G in enumerate(graph_list):
            # highlight = i == optimal_graph_id  # Highlight only the optimal graph
            edge_trace = self._create_edge_trace(G, pos, highlight=False)

            # If we have reached or surpassed the optimal graph, add its highlighted edges
            if i >= optimal_graph_id:
                frames.append(
                    go.Frame(
                        data=[node_trace, edge_trace, optimal_edge_trace],  # Include optimal edges
                        name=f"Graph {i + 1}",
                    )
                )
            else:
                frames.append(
                    go.Frame(
                        data=[node_trace, edge_trace],  # Normal edges only
                        name=f"Graph {i + 1}",
                    )
                )

        # Custom slider labels (first, optimal, and last only)
        slider_labels = [
            "high"
            if i == 0
            else "optimal"
            if i == optimal_graph_id
            else "low"
            if i == num_frames - 1
            else ""  # Empty label for intermediate frames
            for i in range(num_frames)
        ]

        # **Ensure first graph's edges are displayed initially**
        first_edge_trace = self._create_edge_trace(
            graph_list[0], pos, highlight=(0 == optimal_graph_id)
        )

        fig = go.Figure(
            data=[
                node_trace,
                first_edge_trace,  # Now displaying edges initially
            ],
            layout=go.Layout(
                title="",
                showlegend=False,
                hovermode="closest",
                margin=dict(l=20, r=20, t=40, b=20),
                xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                sliders=[
                    {
                        "steps": [
                            {
                                "args": [
                                    [frame.name],
                                    {"frame": {"duration": 0}, "mode": "immediate", "redraw": True},
                                ],
                                "label": slider_labels[i],  # Apply custom labels
                                "method": "animate",
                            }
                            for i, frame in enumerate(frames)
                        ],
                        "active": 0,
                        "currentvalue": {"prefix": "Penalty level: ", "font": {"size": 16}},
                        "pad": {"b": 10, "t": 50},
                    }
                ],
            ),
            frames=frames,  # Ensure frames update correctly
        )

        # Show figure
        return fig

adjacency_matrix property

Adjacency matrix.

causal_order property

Causal order.

__init__(regressor, test, alpha=0.01, pruning_method='murgs', test_size=0.2, local_regression_method='kernel')

Initialize the GroupResit object.

Parameters:

Name	Type	Description	Default
regressor	MultivariateRegressor	A regressor object.	required
test	IndependenceTest	An independence test object.	required
alpha	float	Alpha	0.01
pruning_method	str	The pruning method	'murgs'
test_size	float	Relative size of test-dataset, 0 means no test data	0.2
local_regression_method	str	Type of local linear smoother to use. Options are loess, kernel, soon to be implemented spline. Defaults to kernel.	'kernel'

Source code in gresit/group_resit.py

def __init__(
    self,
    regressor: MultiRegressor,
    test: type[Itest],
    alpha: float = 0.01,
    pruning_method: str = "murgs",
    test_size: float = 0.2,
    local_regression_method: str = "kernel",
) -> None:
    """Initialize the GroupResit object.

    Args:
        regressor (MultivariateRegressor): A regressor object.
        test (IndependenceTest): An independence test object.
        alpha (float): Alpha
        pruning_method (str): The pruning method
        test_size (float): Relative size of test-dataset, 0 means no test data
        local_regression_method (str): Type of local linear smoother to use. Options are
            `loess`, `kernel`, soon to be implemented `spline`. Defaults to `kernel`.
    """
    self.regressor = regressor
    self.independence_test = test()
    self.local_regression_method = local_regression_method
    self._pa: dict[str, list[str]] = {}
    self._pa_history: dict[str, dict[float, list[str]]] = {}
    self.alpha_level: float
    self.alpha = alpha
    self.pruning_method = pruning_method
    self.test_size = test_size
    self.DAG: DAG
    self.layering: dict[str, list[str]] = {}
    super().__init__()

__repr__()

Repr method.

Returns:

Name	Type	Description
str	str	Description of the object.

Source code in gresit/group_resit.py

def __repr__(self) -> str:
    """Repr method.

    Returns:
        str: Description of the object.
    """
    return f"GroupResit(regressor={self.regressor}, independence_test={self.independence_test})"

__str__()

Str method.

Returns:

Name	Type	Description
str	str	Human-readable description of the object.

Source code in gresit/group_resit.py

def __str__(self) -> str:
    """Str method.

    Returns:
        str: Human-readable description of the object.
    """
    method_description = {
        "Regression method: ": self.regressor.__class__.__name__,
        "Independece test:": self.independence_test.__class__.__name__,
        "Inferred causal order: ": "Yes" if self._causal_order else "Not yet",
    }
    s = ""
    for info, info_text in method_description.items():
        s += f"{info:<14}{info_text:>5}\n"

    return s

learn_graph(data_dict, layering=None)

Learn the causal graph.

Parameters:

Name	Type	Description	Default
data_dict	dict[str, ndarray]	A dictionary of np.ndarrays. Key corresponds to group name and values to the corresponding data.	required
layering	dict[str, list[str]]	A dictionary of layering information. Keys correspond to the layer and values to the variable names within each layer.	None

Raises:

Type	Description
NotImplementedError	description

Returns:

Name	Type	Description
DAG	DAG	DAG estimate.

Source code in gresit/group_resit.py

def learn_graph(
    self,
    data_dict: dict[str, np.ndarray],
    layering: dict[str, list[str]] | None = None,
) -> DAG:
    """Learn the causal graph.

    Args:
        data_dict (dict[str, np.ndarray]): A dictionary of np.ndarrays. Key corresponds to
            group name and values to the corresponding data.
        layering (dict[str, list[str]]): A dictionary of layering information. Keys correspond
            to the layer and values to the variable names within each layer.

    Raises:
        NotImplementedError: _description_

    Returns:
        DAG: DAG estimate.
    """
    clean_data_dict = self._dict_preprocessing(data_dict)
    self._get_causal_order(data_dict=clean_data_dict, layering=layering)
    if self.pruning_method == "murgs":
        self._sparse_regression_pruning(data_dict=clean_data_dict)
    elif self.pruning_method == "independence":
        self._independence_prune(data_dict=clean_data_dict, alpha=self.alpha)
    else:
        raise NotImplementedError()

    edge_list = [(parent, child) for child in self._pa for parent in self._pa[child]]

    learned_DAG = DAG(nodes=self._causal_order)
    learned_DAG.add_edges_from(edge_list)
    self._adjacency_matrix = learned_DAG.adjacency_matrix

    self.DAG = learned_DAG
    return learned_DAG

model_selection_with_known_causal_order(pi, data_dict, alpha=0.01, pruning_method='murgs')

Given a known causal order perform model selection.

Parameters:

Name	Type	Description	Default
pi	list[str]	Causal ordering. Entries in pi need to coincide with the keys in data_dict.	required
data_dict	dict[str, ndarray]	A dictionary of np.ndarrays. Key corresponds to group name and values to the corresponding data.	required
alpha	float	The significance level for the independence test. Defaults to 0.1.	0.01
pruning_method	str	The pruning method to use. Defaults to "murgs". other options include "independence"	'murgs'

Source code in gresit/group_resit.py

def model_selection_with_known_causal_order(
    self,
    pi: list[str],
    data_dict: dict[str, np.ndarray],
    alpha: float = 0.01,
    pruning_method: str = "murgs",
) -> None:
    """Given a known causal order perform model selection.

    Args:
        pi (list[str]): Causal ordering. Entries in `pi` need
            to coincide with the keys in `data_dict`.
        data_dict (dict[str, np.ndarray]): A dictionary of np.ndarrays. Key corresponds to
            group name and values to the corresponding data.
        alpha (float, optional): The significance level for the independence test.
            Defaults to 0.1.
        pruning_method (str, optional): The pruning method to use. Defaults to "murgs".
            other options include `"independence"`

    """
    self._insert_known_causal_order(pi=pi)
    if not hasattr(self, "_idx_test"):
        self._idx_test = np.arange(data_dict[list(data_dict.keys())[0]].shape[0])
    if pruning_method == "murgs":
        self._sparse_regression_pruning(data_dict=data_dict)
    elif pruning_method == "independence":
        self._independence_prune(data_dict=data_dict, alpha=alpha)
    else:
        raise NotImplementedError()

    edge_list = [(parent, child) for child in self._pa for parent in self._pa[child]]

    learned_DAG = DAG(nodes=self._causal_order)
    learned_DAG.add_edges_from(edge_list)
    self._adjacency_matrix = learned_DAG.adjacency_matrix
    self.DAG = learned_DAG

show(title='Group RESIT DAG')

Plot the learned DAG.

The plot is interactive, hovering over the nodes reveals the node labels. Colors get brighter the higher the node degree.

Parameters:

Name	Type	Description	Default
title	str	Plot title. Defaults to "Group RESIT DAG".	'Group RESIT DAG'

Raises:

Type	Description
AssertionError	Throws error if DAG not yet learned.

Source code in gresit/group_resit.py

def show(self, title: str = "Group RESIT DAG") -> None:
    """Plot the learned DAG.

    The plot is interactive,
    hovering over the nodes reveals the node labels.
    Colors get brighter the higher the node degree.

    Args:
        title (str, optional): Plot title. Defaults to "Group RESIT DAG".

    Raises:
        AssertionError: Throws error if DAG not yet learned.
    """
    if not hasattr(self, "DAG"):
        raise AssertionError("No graph to plot. Learn the graph first.")

    nx_dag = self.DAG.to_networkx()

    nx_dag = self._add_layered_layout_to_graph(nx_graph=nx_dag)

    edge_x = []
    edge_y = []
    for edge in nx_dag.edges():
        x0, y0 = nx_dag.nodes[edge[0]]["pos"]
        x1, y1 = nx_dag.nodes[edge[1]]["pos"]
        edge_x.append(x0)
        edge_x.append(x1)
        edge_x.append(None)
        edge_y.append(y0)
        edge_y.append(y1)
        edge_y.append(None)

    edge_trace = go.Scatter(
        x=edge_x,
        y=edge_y,
        line=dict(width=0.5, color="#888"),
        hoverinfo="none",
        mode="lines+markers",
        marker=dict(size=10, symbol="arrow-bar-up", angleref="previous"),
    )

    node_x = []
    node_y = []
    for node in nx_dag.nodes():
        x, y = nx_dag.nodes[node]["pos"]
        node_x.append(x)
        node_y.append(y)

    node_trace = go.Scatter(
        x=node_x,
        y=node_y,
        mode="markers",
        hoverinfo="text",
        marker=dict(
            showscale=True,
            # colorscale options
            #'Greys' | 'YlGnBu' | 'Greens' | 'YlOrRd' | 'Bluered' | 'RdBu' |
            #'Reds' | 'Blues' | 'Picnic' | 'Rainbow' | 'Portland' | 'Jet' |
            #'Hot' | 'Blackbody' | 'Earth' | 'Electric' | 'Viridis' |
            colorscale="YlGnBu",
            reversescale=True,
            color=[],
            size=10,
            colorbar=dict(thickness=15, title="node degree", xanchor="left", titleside="right"),
            line_width=2,
        ),
    )

    node_adjacencies = []
    node_text = []
    for node, adjacencies in enumerate(nx_dag.adjacency()):
        node_adjacencies.append(len(adjacencies[1]))
        node_text.append(list(nx_dag.nodes)[node])

    node_trace.marker.color = node_adjacencies
    node_trace.text = node_text

    fig = go.Figure(
        data=[edge_trace, node_trace],
        layout=go.Layout(
            title=title,
            # titlefont_size=16,
            showlegend=False,
            hovermode="closest",
            margin=dict(b=20, l=5, r=5, t=40),
            xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
            yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
        ),
    )
    fig.show()

show_interactive(layer_gap=8.0)

Show interactive plot with slider to select sparsity level.

Parameters:

Name	Type	Description	Default
layer_gap	float	gap between layers when displaying. Defaults to 8.0.	8.0

Source code in gresit/group_resit.py

def show_interactive(self, layer_gap: float = 8.0) -> go.Figure:
    """Show interactive plot with slider to select sparsity level.

    Args:
        layer_gap (float, optional): gap between layers when displaying. Defaults to 8.0.
    """
    # Fixed layout for consistency
    graph_list = self._make_graph_list()
    pos = self._layer_pos(G=graph_list[-1], layers=self.layering, layer_gap=layer_gap)
    nodes = list(graph_list[-1].nodes)

    # Extract node coordinates
    node_x, node_y = zip(*[pos[n] for n in nodes])

    node_trace = go.Scatter(
        x=node_x,
        y=node_y,
        mode="markers",
        text=nodes,
        hoverinfo="text",
        marker=dict(size=15, color="lightblue", line=dict(color="black", width=1)),
    )

    # Identify the "optimal" graph
    optimal_graph_id = next(
        (
            i
            for i, graph in enumerate(graph_list)
            if nx.utils.graphs_equal(self.DAG.to_networkx(), graph)
        ),
        None,
    )
    if optimal_graph_id is not None:
        # Create the optimal graph's edge trace
        optimal_edge_trace = self._create_edge_trace(
            G=graph_list[optimal_graph_id], pos=pos, highlight=True
        )
    else:
        raise ValueError("Optimal graph not found in graph list.")

    # Create frames for each graph
    frames = []
    num_frames = len(graph_list)

    for i, G in enumerate(graph_list):
        # highlight = i == optimal_graph_id  # Highlight only the optimal graph
        edge_trace = self._create_edge_trace(G, pos, highlight=False)

        # If we have reached or surpassed the optimal graph, add its highlighted edges
        if i >= optimal_graph_id:
            frames.append(
                go.Frame(
                    data=[node_trace, edge_trace, optimal_edge_trace],  # Include optimal edges
                    name=f"Graph {i + 1}",
                )
            )
        else:
            frames.append(
                go.Frame(
                    data=[node_trace, edge_trace],  # Normal edges only
                    name=f"Graph {i + 1}",
                )
            )

    # Custom slider labels (first, optimal, and last only)
    slider_labels = [
        "high"
        if i == 0
        else "optimal"
        if i == optimal_graph_id
        else "low"
        if i == num_frames - 1
        else ""  # Empty label for intermediate frames
        for i in range(num_frames)
    ]

    # **Ensure first graph's edges are displayed initially**
    first_edge_trace = self._create_edge_trace(
        graph_list[0], pos, highlight=(0 == optimal_graph_id)
    )

    fig = go.Figure(
        data=[
            node_trace,
            first_edge_trace,  # Now displaying edges initially
        ],
        layout=go.Layout(
            title="",
            showlegend=False,
            hovermode="closest",
            margin=dict(l=20, r=20, t=40, b=20),
            xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
            yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
            sliders=[
                {
                    "steps": [
                        {
                            "args": [
                                [frame.name],
                                {"frame": {"duration": 0}, "mode": "immediate", "redraw": True},
                            ],
                            "label": slider_labels[i],  # Apply custom labels
                            "method": "animate",
                        }
                        for i, frame in enumerate(frames)
                    ],
                    "active": 0,
                    "currentvalue": {"prefix": "Penalty level: ", "font": {"size": 16}},
                    "pad": {"b": 10, "t": 50},
                }
            ],
        ),
        frames=frames,  # Ensure frames update correctly
    )

    # Show figure
    return fig

Utility classes and functions related to gresit.

Copyright (c) 2025 Robert Bosch GmbH

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.

EarlyStopping

Early stopping to conserve compute resources.

Source code in gresit/torch_models.py

class EarlyStopping:
    """Early stopping to conserve compute resources."""

    def __init__(
        self, patience: int = 5, min_delta: float = 0.0, restore_best_weights: bool = True
    ) -> None:
        """Initializes the EarlyStopping class.

        Args:
            patience (int, optional): Number of epochs to wait for the validation error to improve.
                Defaults to 5.
            min_delta (float, optional): Minimum change that should be considered an improvement.
                Defaults to 0.0.
            restore_best_weights (bool, optional): Restores the weights to the values they were when
                the validation set was best. Defaults to True.
        """
        self.patience = patience
        self.min_delta = min_delta
        self.restore_best_weights = restore_best_weights
        self.best_model = None
        self.best_loss: torch.Tensor | None = None
        self.counter = 0
        self.status = ""

    def __call__(self, model: nn.Module, val_loss: float) -> bool:
        """Stopping actions.

        Args:
            model (nn.Module): NN model
            val_loss (float): Value of loss function.

        Returns:
            bool: True if training may be concluded.
        """
        if self.best_loss is None:
            self.best_loss = val_loss
            self.best_model = deepcopy(model.state_dict())
        elif self.best_loss - val_loss >= self.min_delta:
            self.best_model = deepcopy(model.state_dict())
            self.best_loss = val_loss
            self.counter = 0
            self.status = f"Improvement found, counter reset to {self.counter}"
        else:
            self.counter += 1
            self.status = f"No improvement in the last {self.counter} epochs"
            if self.counter >= self.patience:
                self.status = "Early stopping triggered"
                if self.restore_best_weights:
                    model.load_state_dict(self.best_model)
                return True
        return False

__call__(model, val_loss)

Stopping actions.

Parameters:

Name	Type	Description	Default
model	Module	NN model	required
val_loss	float	Value of loss function.	required

Returns:

Name	Type	Description
bool	bool	True if training may be concluded.

Source code in gresit/torch_models.py

def __call__(self, model: nn.Module, val_loss: float) -> bool:
    """Stopping actions.

    Args:
        model (nn.Module): NN model
        val_loss (float): Value of loss function.

    Returns:
        bool: True if training may be concluded.
    """
    if self.best_loss is None:
        self.best_loss = val_loss
        self.best_model = deepcopy(model.state_dict())
    elif self.best_loss - val_loss >= self.min_delta:
        self.best_model = deepcopy(model.state_dict())
        self.best_loss = val_loss
        self.counter = 0
        self.status = f"Improvement found, counter reset to {self.counter}"
    else:
        self.counter += 1
        self.status = f"No improvement in the last {self.counter} epochs"
        if self.counter >= self.patience:
            self.status = "Early stopping triggered"
            if self.restore_best_weights:
                model.load_state_dict(self.best_model)
            return True
    return False

__init__(patience=5, min_delta=0.0, restore_best_weights=True)

Initializes the EarlyStopping class.

Parameters:

Name	Type	Description	Default
patience	int	Number of epochs to wait for the validation error to improve. Defaults to 5.	5
min_delta	float	Minimum change that should be considered an improvement. Defaults to 0.0.	0.0
restore_best_weights	bool	Restores the weights to the values they were when the validation set was best. Defaults to True.	True

Source code in gresit/torch_models.py

def __init__(
    self, patience: int = 5, min_delta: float = 0.0, restore_best_weights: bool = True
) -> None:
    """Initializes the EarlyStopping class.

    Args:
        patience (int, optional): Number of epochs to wait for the validation error to improve.
            Defaults to 5.
        min_delta (float, optional): Minimum change that should be considered an improvement.
            Defaults to 0.0.
        restore_best_weights (bool, optional): Restores the weights to the values they were when
            the validation set was best. Defaults to True.
    """
    self.patience = patience
    self.min_delta = min_delta
    self.restore_best_weights = restore_best_weights
    self.best_model = None
    self.best_loss: torch.Tensor | None = None
    self.counter = 0
    self.status = ""

MLP

Bases: Module

Torch MLP with single-hidden layer and sigmoid non-linearity.

Not too fancy but does what it is supposed to do.

Source code in gresit/torch_models.py

class MLP(nn.Module):  # type: ignore
    """Torch MLP with single-hidden layer and sigmoid non-linearity.

    Not too fancy but does what it is supposed to do.
    """

    def __init__(
        self, input_dim: int, output_dim: int, hidden_dim: int = 100, dropout: float = 0.0
    ) -> None:
        """Initializes the NN.

        Args:
            input_dim (int): dim of input layer
            output_dim (int, optional): dim of output layer.
            hidden_dim (int, optional): dim of hidden layer. Defaults to 100.
            dropout (float, optional): dropout probability. Defaults to 0.0.
        """
        super().__init__()

        self.input_dim = input_dim
        self.output_dim = output_dim
        self.hidden_dim = hidden_dim
        self.dropout = dropout

        # define the hidden layer and final layers
        self.hidden1 = nn.Linear(self.input_dim, self.hidden_dim, dtype=torch.float32)
        self.hidden2 = nn.Linear(self.hidden_dim, self.hidden_dim, dtype=torch.float32)
        self.final = nn.Linear(self.hidden_dim, self.output_dim, dtype=torch.float32, bias=False)
        self.bias = nn.Parameter(
            data=torch.Tensor([0] * self.output_dim),
            requires_grad=False,
        ).to(torch.float32)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward passing the model.

        Args:
            x (torch.Tensor): training data.

        Returns:
            torch.Tensor: prediction after pass through nn.
        """
        x = self.hidden1(x)  # apply hidden layer
        x = torch.tanh(x)  # apply sigmoid non-linearity
        x = nn.Dropout(p=self.dropout)(x)
        x = self.hidden2(x)  # apply hidden layer
        x = nn.Dropout(p=self.dropout)(x)
        x = torch.tanh(x)  # apply sigmoid non-linearity
        x = nn.Dropout(p=self.dropout)(x)
        x = self.final(x)  # apply final layer
        x = x + self.bias  # correct for bias

        return x

    def update_bias(self, bias_value: torch.Tensor) -> None:
        """Update bias due to HSIC location invariance.

        Args:
            bias_value (torch.Tensor): _description_

        Raises:
            ValueError: _description_
        """
        if bias_value.shape == self.bias.shape:
            self.bias.copy_(bias_value)
        else:
            raise ValueError("Shape mismatch")

__init__(input_dim, output_dim, hidden_dim=100, dropout=0.0)

Initializes the NN.

Parameters:

Name	Type	Description	Default
input_dim	int	dim of input layer	required
output_dim	int	dim of output layer.	required
hidden_dim	int	dim of hidden layer. Defaults to 100.	100
dropout	float	dropout probability. Defaults to 0.0.	0.0

Source code in gresit/torch_models.py

def __init__(
    self, input_dim: int, output_dim: int, hidden_dim: int = 100, dropout: float = 0.0
) -> None:
    """Initializes the NN.

    Args:
        input_dim (int): dim of input layer
        output_dim (int, optional): dim of output layer.
        hidden_dim (int, optional): dim of hidden layer. Defaults to 100.
        dropout (float, optional): dropout probability. Defaults to 0.0.
    """
    super().__init__()

    self.input_dim = input_dim
    self.output_dim = output_dim
    self.hidden_dim = hidden_dim
    self.dropout = dropout

    # define the hidden layer and final layers
    self.hidden1 = nn.Linear(self.input_dim, self.hidden_dim, dtype=torch.float32)
    self.hidden2 = nn.Linear(self.hidden_dim, self.hidden_dim, dtype=torch.float32)
    self.final = nn.Linear(self.hidden_dim, self.output_dim, dtype=torch.float32, bias=False)
    self.bias = nn.Parameter(
        data=torch.Tensor([0] * self.output_dim),
        requires_grad=False,
    ).to(torch.float32)

forward(x)

Forward passing the model.

Parameters:

Name	Type	Description	Default
x	Tensor	training data.	required

Returns:

Type	Description
Tensor	torch.Tensor: prediction after pass through nn.

Source code in gresit/torch_models.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward passing the model.

    Args:
        x (torch.Tensor): training data.

    Returns:
        torch.Tensor: prediction after pass through nn.
    """
    x = self.hidden1(x)  # apply hidden layer
    x = torch.tanh(x)  # apply sigmoid non-linearity
    x = nn.Dropout(p=self.dropout)(x)
    x = self.hidden2(x)  # apply hidden layer
    x = nn.Dropout(p=self.dropout)(x)
    x = torch.tanh(x)  # apply sigmoid non-linearity
    x = nn.Dropout(p=self.dropout)(x)
    x = self.final(x)  # apply final layer
    x = x + self.bias  # correct for bias

    return x

update_bias(bias_value)

Update bias due to HSIC location invariance.

Parameters:

Name	Type	Description	Default
bias_value	Tensor	description	required

Raises:

Type	Description
ValueError	description

Source code in gresit/torch_models.py

def update_bias(self, bias_value: torch.Tensor) -> None:
    """Update bias due to HSIC location invariance.

    Args:
        bias_value (torch.Tensor): _description_

    Raises:
        ValueError: _description_
    """
    if bias_value.shape == self.bias.shape:
        self.bias.copy_(bias_value)
    else:
        raise ValueError("Shape mismatch")

MultioutcomeGPR

Bases: MultiRegressor

MultioutcomeGPR Gaussian Process Regression class.

Source code in gresit/torch_models.py

class MultioutcomeGPR(MultiRegressor):
    """MultioutcomeGPR Gaussian Process Regression class."""

    def __init__(
        self,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
        n_epochs: int = 300,
        patience: int = 50,
        learning_rate: float = 0.01,
        val_size: float = 0.2,
        batch_size: int | None = None,
        es: bool = True,
    ) -> None:
        """Initialize MLP.

        Args:
            rng (np.random.Generator, optional): _description_.
                Defaults to np.random.default_rng(seed=2024).
            n_epochs (int): number of times the data gets passed trough the MLP.
                Defaults to 6.
            patience (int, optional): Minimal number of epochs to train before early stopping
                applies
            learning_rate (float, optional): _description_. Defaults to 1e-3.
            val_size (float): Relative size of the validation dataset
            batch_size (int): Batch size.
            es (bool, optional): Early stopping.
        """
        super().__init__(rng)

        self.training_info: pd.DataFrame
        self.prediction: torch.Tensor
        self._Y_test: np.ndarray | torch.Tensor
        self._X_test: np.ndarray | torch.Tensor
        self._model: nn.Module
        self.learning_rate = learning_rate
        self.n_epochs = n_epochs
        self.patience = patience
        self.batch_size = batch_size
        self.val_size = val_size
        self.es = es

        has_mps = torch.backends.mps.is_built()
        self.device = "mps" if has_mps else "cuda" if torch.cuda.is_available() else "cpu"

    def fit(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        idx_train: np.ndarray | None = None,
        idx_test: np.ndarray | None = None,
    ) -> None:
        """Fit the MLP model.

        Args:
            X (np.ndarray): Input data
            Y (np.ndarray): Target data
            idx_train (np.ndarray): training indices
            idx_test (np.ndarray): test indices

        Returns:
            int: epoch at which training was stopped.
        """
        if idx_train is None and idx_test is None:
            X_train, X_val, Y_train, Y_val = self.split_and_standardize(
                X=X,
                Y=Y,
                test_size=self.val_size,
            )
        else:
            X_train, X_val, Y_train, Y_val = (X[idx_train], X[idx_test], Y[idx_train], Y[idx_test])

        self._X_train = torch.from_numpy(X_train).float().to(self.device)
        self._Y_train = torch.from_numpy(Y_train).float().to(self.device)

        self._X_val = torch.from_numpy(X_val).float().to(self.device)
        self._Y_val = torch.from_numpy(Y_val).float().to(self.device)

        num_tasks = self._Y_train.size(1)

        likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)
        model = MultitaskGPModel(self._X_train, self._Y_train, likelihood).to(self.device)
        if self.es:
            es = EarlyStopping(patience=self.patience)

        # Find optimal model hyperparameters
        model.train()
        likelihood.train()
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=self.learning_rate,
        )

        mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

        history_val = []
        for i in range(self.n_epochs):
            optimizer.zero_grad()
            output = model(self._X_train)
            loss = -mll(output, self._Y_train)
            loss.backward()
            optimizer.step()

            # Evaluate Model
            with torch.no_grad(), gpytorch.settings.fast_pred_var():
                model.eval().to(self.device)
                likelihood.eval().to(self.device)
                predictions = likelihood(model(self._X_val))
                Y_pred = predictions.mean
                vloss = loss_mse(self._X_val, Y_pred, self._Y_val, device=self.device)

                history_val.append(
                    {
                        "epoch": i,
                        "task": "regression",
                        "lr": self.learning_rate,
                        "loss": vloss.detach().cpu().numpy(),
                    }
                )
                if self.es:
                    if es(model, vloss):
                        print(f"{es.status}")
                        break

        self.validation_info = pd.DataFrame(history_val)
        self.prediction = make_preds_single(model, self._X_val)
        self._model = model
        self._likelihood = likelihood
        self.final_loss = es.best_loss if self.es else None

    def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
        """Make predictions.

        Returns:
            np.ndarray: Predictions
        """
        if X_test is None:
            prediction = self.prediction.detach().numpy()
        else:
            X_torch_test = torch.from_numpy(X_test).float().to(self.device)
            with torch.no_grad(), gpytorch.settings.fast_pred_var():
                pred_model = self._model.eval().to(self.device)
                pred_likelihood = self._likelihood.eval().to(self.device)
                predictions = pred_likelihood(pred_model(X_torch_test))
                pred = predictions.mean
            prediction = pred.detach().cpu().numpy()
        return prediction

    def plot_training_info(self) -> None:
        """Plot some summary over the MSE during training."""
        sns.lmplot(
            x="trainstep",
            y="loss",
            hue="task",
            lowess=True,
            scatter_kws={"alpha": 0.5},
            line_kws={"linewidth": 2},
            data=self.training_info,
        )

        sns.lmplot(
            x="epoch",
            y="loss",
            hue="task",
            lowess=True,
            scatter_kws={"alpha": 0.5},
            line_kws={"linewidth": 2},
            data=self.validation_info,
        )

__init__(rng=np.random.default_rng(seed=2024), n_epochs=300, patience=50, learning_rate=0.01, val_size=0.2, batch_size=None, es=True)

Initialize MLP.

Parameters:

Name	Type	Description	Default
rng	Generator	description. Defaults to np.random.default_rng(seed=2024).	default_rng(seed=2024)
n_epochs	int	number of times the data gets passed trough the MLP. Defaults to 6.	300
patience	int	Minimal number of epochs to train before early stopping applies	50
learning_rate	float	description. Defaults to 1e-3.	0.01
val_size	float	Relative size of the validation dataset	0.2
batch_size	int	Batch size.	None
es	bool	Early stopping.	True

Source code in gresit/torch_models.py

def __init__(
    self,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
    n_epochs: int = 300,
    patience: int = 50,
    learning_rate: float = 0.01,
    val_size: float = 0.2,
    batch_size: int | None = None,
    es: bool = True,
) -> None:
    """Initialize MLP.

    Args:
        rng (np.random.Generator, optional): _description_.
            Defaults to np.random.default_rng(seed=2024).
        n_epochs (int): number of times the data gets passed trough the MLP.
            Defaults to 6.
        patience (int, optional): Minimal number of epochs to train before early stopping
            applies
        learning_rate (float, optional): _description_. Defaults to 1e-3.
        val_size (float): Relative size of the validation dataset
        batch_size (int): Batch size.
        es (bool, optional): Early stopping.
    """
    super().__init__(rng)

    self.training_info: pd.DataFrame
    self.prediction: torch.Tensor
    self._Y_test: np.ndarray | torch.Tensor
    self._X_test: np.ndarray | torch.Tensor
    self._model: nn.Module
    self.learning_rate = learning_rate
    self.n_epochs = n_epochs
    self.patience = patience
    self.batch_size = batch_size
    self.val_size = val_size
    self.es = es

    has_mps = torch.backends.mps.is_built()
    self.device = "mps" if has_mps else "cuda" if torch.cuda.is_available() else "cpu"

fit(X, Y, idx_train=None, idx_test=None)

Fit the MLP model.

Parameters:

Name	Type	Description	Default
X	ndarray	Input data	required
Y	ndarray	Target data	required
idx_train	ndarray	training indices	None
idx_test	ndarray	test indices	None

Returns:

Name	Type	Description
int	None	epoch at which training was stopped.

Source code in gresit/torch_models.py

def fit(
    self,
    X: np.ndarray,
    Y: np.ndarray,
    idx_train: np.ndarray | None = None,
    idx_test: np.ndarray | None = None,
) -> None:
    """Fit the MLP model.

    Args:
        X (np.ndarray): Input data
        Y (np.ndarray): Target data
        idx_train (np.ndarray): training indices
        idx_test (np.ndarray): test indices

    Returns:
        int: epoch at which training was stopped.
    """
    if idx_train is None and idx_test is None:
        X_train, X_val, Y_train, Y_val = self.split_and_standardize(
            X=X,
            Y=Y,
            test_size=self.val_size,
        )
    else:
        X_train, X_val, Y_train, Y_val = (X[idx_train], X[idx_test], Y[idx_train], Y[idx_test])

    self._X_train = torch.from_numpy(X_train).float().to(self.device)
    self._Y_train = torch.from_numpy(Y_train).float().to(self.device)

    self._X_val = torch.from_numpy(X_val).float().to(self.device)
    self._Y_val = torch.from_numpy(Y_val).float().to(self.device)

    num_tasks = self._Y_train.size(1)

    likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)
    model = MultitaskGPModel(self._X_train, self._Y_train, likelihood).to(self.device)
    if self.es:
        es = EarlyStopping(patience=self.patience)

    # Find optimal model hyperparameters
    model.train()
    likelihood.train()
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=self.learning_rate,
    )

    mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

    history_val = []
    for i in range(self.n_epochs):
        optimizer.zero_grad()
        output = model(self._X_train)
        loss = -mll(output, self._Y_train)
        loss.backward()
        optimizer.step()

        # Evaluate Model
        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            model.eval().to(self.device)
            likelihood.eval().to(self.device)
            predictions = likelihood(model(self._X_val))
            Y_pred = predictions.mean
            vloss = loss_mse(self._X_val, Y_pred, self._Y_val, device=self.device)

            history_val.append(
                {
                    "epoch": i,
                    "task": "regression",
                    "lr": self.learning_rate,
                    "loss": vloss.detach().cpu().numpy(),
                }
            )
            if self.es:
                if es(model, vloss):
                    print(f"{es.status}")
                    break

    self.validation_info = pd.DataFrame(history_val)
    self.prediction = make_preds_single(model, self._X_val)
    self._model = model
    self._likelihood = likelihood
    self.final_loss = es.best_loss if self.es else None

plot_training_info()

Plot some summary over the MSE during training.

Source code in gresit/torch_models.py

def plot_training_info(self) -> None:
    """Plot some summary over the MSE during training."""
    sns.lmplot(
        x="trainstep",
        y="loss",
        hue="task",
        lowess=True,
        scatter_kws={"alpha": 0.5},
        line_kws={"linewidth": 2},
        data=self.training_info,
    )

    sns.lmplot(
        x="epoch",
        y="loss",
        hue="task",
        lowess=True,
        scatter_kws={"alpha": 0.5},
        line_kws={"linewidth": 2},
        data=self.validation_info,
    )

predict(X_test=None)

Make predictions.

Returns:

Type	Description
ndarray	np.ndarray: Predictions

Source code in gresit/torch_models.py

def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
    """Make predictions.

    Returns:
        np.ndarray: Predictions
    """
    if X_test is None:
        prediction = self.prediction.detach().numpy()
    else:
        X_torch_test = torch.from_numpy(X_test).float().to(self.device)
        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            pred_model = self._model.eval().to(self.device)
            pred_likelihood = self._likelihood.eval().to(self.device)
            predictions = pred_likelihood(pred_model(X_torch_test))
            pred = predictions.mean
        prediction = pred.detach().cpu().numpy()
    return prediction

Multioutcome_MLP

Bases: MultiRegressor

Fit simple MLP with one hidden layer.

Source code in gresit/torch_models.py

class Multioutcome_MLP(MultiRegressor):
    """Fit simple MLP with one hidden layer."""

    def __init__(
        self,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
        loss: str = "mse",
        dropout_proba: float = 0.6,
        n_epochs: int = 300,
        patience: int = 50,
        learning_rate: float = 0.01,
        val_size: float = 0.2,
        batch_size: int = 200,
        es: bool = True,
    ) -> None:
        """Initialize MLP.

        Args:
            rng (np.random.Generator, optional): _description_.
                Defaults to np.random.default_rng(seed=2024).
            loss (str, optional): Standard mse loss is default.
                Other options are `hsic` and `disco`.
                Defaults to "mse".
            dropout_proba (float, optional): _description_. Defaults to 0.6.
            n_epochs (int): number of times the data gets passed trough the MLP.
                Defaults to 6.
            patience (int, optional): Minimal number of epochs to train before early stopping
                applies
            learning_rate (float, optional): _description_. Defaults to 1e-3.
            val_size (float): Relative size of the validation dataset
            batch_size (int): Batch size.
            es (bool, optional): Early stopping. Defaults to true.
        """
        super().__init__(rng)

        self.training_info: pd.DataFrame
        self.prediction: torch.Tensor
        self.input_dim: int
        self.output_dim: int
        self._Y_test: np.ndarray | torch.Tensor
        self._X_test: np.ndarray | torch.Tensor
        self._model: nn.Module
        self.dropout_proba = dropout_proba
        self.learning_rate = learning_rate
        self.n_epochs = n_epochs
        self.patience = patience
        self.batch_size = batch_size
        self.val_size = val_size
        self.es = es

        self.loss_name = loss

        has_mps = torch.backends.mps.is_built()
        self.device = "mps" if has_mps else "cuda" if torch.cuda.is_available() else "cpu"

    def _compute_bias(
        self,
        model: nn.Module,
        y_true: torch.Tensor,
        x: torch.Tensor,
    ) -> torch.Tensor:
        y_pred_interim: torch.Tensor = model(x)
        return y_true.mean(dim=0) - y_pred_interim.mean(dim=0)

    def _standardize(self, x: torch.Tensor) -> torch.Tensor:
        return (x - x.mean(dim=0)) / x.std(dim=0)

    def fit(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        idx_train: np.ndarray | None = None,
        idx_test: np.ndarray | None = None,
    ) -> None:
        """Fit the MLP model.

        Args:
            X (np.ndarray): Input data
            Y (np.ndarray): Target data
            idx_train (np.ndarray): training indices
            idx_test (np.ndarray): test indices

        Returns:
            int: epoch at which training was stopped.
        """
        if idx_train is None and idx_test is None:
            X_train, X_val, Y_train, Y_val = self.split_and_standardize(
                X=X,
                Y=Y,
                test_size=self.val_size,
            )
        else:
            X_train, X_val, Y_train, Y_val = (X[idx_train], X[idx_test], Y[idx_train], Y[idx_test])
            X_train = (X_train - X_train.mean(axis=0)) / (X_train.std(axis=0))
            Y_train = (Y_train - Y_train.mean(axis=0)) / (Y_train.std(axis=0))
            X_val = (X_val - X_train.mean(axis=0)) / (X_train.std(axis=0))
            Y_val = (Y_val - Y_train.mean(axis=0)) / (Y_train.std(axis=0))

        if self.loss_name == "hisc":
            loss_fn = loss_hsic
        else:
            loss_fn = loss_mse

        self._X_train = torch.from_numpy(X_train).float().to(self.device)
        self._Y_train = torch.from_numpy(Y_train).float().to(self.device)

        self._X_val = torch.from_numpy(X_val).float().to(self.device)
        self._Y_val = torch.from_numpy(Y_val).float().to(self.device)

        model = MLP(
            input_dim=self._X_train.shape[1],
            output_dim=self._Y_train.shape[1],
            dropout=self.dropout_proba,
        ).to(self.device)

        if self.es:
            es = EarlyStopping(patience=self.patience)

        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=self.learning_rate,
            amsgrad=True,
        )

        dataset_train = TensorDataset(self._X_train, self._Y_train)
        dataloader_train = DataLoader(dataset_train, batch_size=self.batch_size, shuffle=True)

        history_train = []
        history_val = []
        trainstep = 0
        for i in tqdm(range(self.n_epochs), ncols=100):
            for j, (x_batch, y_batch) in enumerate(dataloader_train):
                Y_pred = model(x_batch)
                loss = loss_fn(
                    self._standardize(x_batch),
                    self._standardize(Y_pred),
                    self._standardize(y_batch),
                    device=self.device,
                )

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                history_train.append(
                    {
                        "epoch": i,
                        "minibatch": j,
                        "trainstep": trainstep,
                        "task": "regression",
                        "loss": loss.detach().cpu().numpy(),
                    }
                )
                trainstep += 1

            # Evaluate Model
            with torch.no_grad():
                if self.loss_name != "mse":
                    bias = self._compute_bias(model=model, y_true=self._Y_val, x=self._X_val)
                    model.update_bias(bias)
                Y_pred = model(self._X_val)
                vloss = loss_fn(
                    self._standardize(self._X_val),
                    self._standardize(Y_pred),
                    self._standardize(self._Y_val),
                    device=self.device,
                )

                history_val.append(
                    {
                        "epoch": i,
                        "trainstep": trainstep,
                        "task": "regression",
                        # "lr": lr_scheduler.get_last_lr()[0],
                        "lr": self.learning_rate,
                        "loss": vloss.detach().cpu().numpy(),
                    }
                )
                if self.es:
                    if es(model, vloss):
                        print(f"{es.status}")
                        break

        self.training_info = pd.DataFrame(history_train)
        self.validation_info = pd.DataFrame(history_val)
        self.prediction = make_preds_single(model, self._X_val)
        self._model = model
        self.final_loss = (
            es.best_loss.detach().cpu().numpy() if self.es and es.best_loss else np.infty
        )

    def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
        """Make predictions.

        Returns:
            np.ndarray: Predictions
        """
        if X_test is None:
            prediction = self.prediction.detach().numpy()
        else:
            X_torch_test = torch.from_numpy(X_test).float().to(self.device)
            pred = make_preds_single(self._model, X_torch_test)
            prediction = pred.detach().cpu().numpy()
        return prediction

    def plot_training_info(self) -> None:
        """Plot some summary over the MSE during training."""
        sns.lmplot(
            x="trainstep",
            y="loss",
            hue="task",
            lowess=True,
            scatter_kws={"alpha": 0.5},
            line_kws={"linewidth": 2},
            data=self.training_info,
        )

        sns.lmplot(
            x="epoch",
            y="loss",
            hue="task",
            lowess=True,
            scatter_kws={"alpha": 0.5},
            line_kws={"linewidth": 2},
            data=self.validation_info,
        )

__init__(rng=np.random.default_rng(seed=2024), loss='mse', dropout_proba=0.6, n_epochs=300, patience=50, learning_rate=0.01, val_size=0.2, batch_size=200, es=True)

Initialize MLP.

Parameters:

Name	Type	Description	Default
rng	Generator	description. Defaults to np.random.default_rng(seed=2024).	default_rng(seed=2024)
loss	str	Standard mse loss is default. Other options are hsic and disco. Defaults to "mse".	'mse'
dropout_proba	float	description. Defaults to 0.6.	0.6
n_epochs	int	number of times the data gets passed trough the MLP. Defaults to 6.	300
patience	int	Minimal number of epochs to train before early stopping applies	50
learning_rate	float	description. Defaults to 1e-3.	0.01
val_size	float	Relative size of the validation dataset	0.2
batch_size	int	Batch size.	200
es	bool	Early stopping. Defaults to true.	True

Source code in gresit/torch_models.py

def __init__(
    self,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
    loss: str = "mse",
    dropout_proba: float = 0.6,
    n_epochs: int = 300,
    patience: int = 50,
    learning_rate: float = 0.01,
    val_size: float = 0.2,
    batch_size: int = 200,
    es: bool = True,
) -> None:
    """Initialize MLP.

    Args:
        rng (np.random.Generator, optional): _description_.
            Defaults to np.random.default_rng(seed=2024).
        loss (str, optional): Standard mse loss is default.
            Other options are `hsic` and `disco`.
            Defaults to "mse".
        dropout_proba (float, optional): _description_. Defaults to 0.6.
        n_epochs (int): number of times the data gets passed trough the MLP.
            Defaults to 6.
        patience (int, optional): Minimal number of epochs to train before early stopping
            applies
        learning_rate (float, optional): _description_. Defaults to 1e-3.
        val_size (float): Relative size of the validation dataset
        batch_size (int): Batch size.
        es (bool, optional): Early stopping. Defaults to true.
    """
    super().__init__(rng)

    self.training_info: pd.DataFrame
    self.prediction: torch.Tensor
    self.input_dim: int
    self.output_dim: int
    self._Y_test: np.ndarray | torch.Tensor
    self._X_test: np.ndarray | torch.Tensor
    self._model: nn.Module
    self.dropout_proba = dropout_proba
    self.learning_rate = learning_rate
    self.n_epochs = n_epochs
    self.patience = patience
    self.batch_size = batch_size
    self.val_size = val_size
    self.es = es

    self.loss_name = loss

    has_mps = torch.backends.mps.is_built()
    self.device = "mps" if has_mps else "cuda" if torch.cuda.is_available() else "cpu"

fit(X, Y, idx_train=None, idx_test=None)

Fit the MLP model.

Parameters:

Name	Type	Description	Default
X	ndarray	Input data	required
Y	ndarray	Target data	required
idx_train	ndarray	training indices	None
idx_test	ndarray	test indices	None

Returns:

Name	Type	Description
int	None	epoch at which training was stopped.

Source code in gresit/torch_models.py

def fit(
    self,
    X: np.ndarray,
    Y: np.ndarray,
    idx_train: np.ndarray | None = None,
    idx_test: np.ndarray | None = None,
) -> None:
    """Fit the MLP model.

    Args:
        X (np.ndarray): Input data
        Y (np.ndarray): Target data
        idx_train (np.ndarray): training indices
        idx_test (np.ndarray): test indices

    Returns:
        int: epoch at which training was stopped.
    """
    if idx_train is None and idx_test is None:
        X_train, X_val, Y_train, Y_val = self.split_and_standardize(
            X=X,
            Y=Y,
            test_size=self.val_size,
        )
    else:
        X_train, X_val, Y_train, Y_val = (X[idx_train], X[idx_test], Y[idx_train], Y[idx_test])
        X_train = (X_train - X_train.mean(axis=0)) / (X_train.std(axis=0))
        Y_train = (Y_train - Y_train.mean(axis=0)) / (Y_train.std(axis=0))
        X_val = (X_val - X_train.mean(axis=0)) / (X_train.std(axis=0))
        Y_val = (Y_val - Y_train.mean(axis=0)) / (Y_train.std(axis=0))

    if self.loss_name == "hisc":
        loss_fn = loss_hsic
    else:
        loss_fn = loss_mse

    self._X_train = torch.from_numpy(X_train).float().to(self.device)
    self._Y_train = torch.from_numpy(Y_train).float().to(self.device)

    self._X_val = torch.from_numpy(X_val).float().to(self.device)
    self._Y_val = torch.from_numpy(Y_val).float().to(self.device)

    model = MLP(
        input_dim=self._X_train.shape[1],
        output_dim=self._Y_train.shape[1],
        dropout=self.dropout_proba,
    ).to(self.device)

    if self.es:
        es = EarlyStopping(patience=self.patience)

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=self.learning_rate,
        amsgrad=True,
    )

    dataset_train = TensorDataset(self._X_train, self._Y_train)
    dataloader_train = DataLoader(dataset_train, batch_size=self.batch_size, shuffle=True)

    history_train = []
    history_val = []
    trainstep = 0
    for i in tqdm(range(self.n_epochs), ncols=100):
        for j, (x_batch, y_batch) in enumerate(dataloader_train):
            Y_pred = model(x_batch)
            loss = loss_fn(
                self._standardize(x_batch),
                self._standardize(Y_pred),
                self._standardize(y_batch),
                device=self.device,
            )

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            history_train.append(
                {
                    "epoch": i,
                    "minibatch": j,
                    "trainstep": trainstep,
                    "task": "regression",
                    "loss": loss.detach().cpu().numpy(),
                }
            )
            trainstep += 1

        # Evaluate Model
        with torch.no_grad():
            if self.loss_name != "mse":
                bias = self._compute_bias(model=model, y_true=self._Y_val, x=self._X_val)
                model.update_bias(bias)
            Y_pred = model(self._X_val)
            vloss = loss_fn(
                self._standardize(self._X_val),
                self._standardize(Y_pred),
                self._standardize(self._Y_val),
                device=self.device,
            )

            history_val.append(
                {
                    "epoch": i,
                    "trainstep": trainstep,
                    "task": "regression",
                    # "lr": lr_scheduler.get_last_lr()[0],
                    "lr": self.learning_rate,
                    "loss": vloss.detach().cpu().numpy(),
                }
            )
            if self.es:
                if es(model, vloss):
                    print(f"{es.status}")
                    break

    self.training_info = pd.DataFrame(history_train)
    self.validation_info = pd.DataFrame(history_val)
    self.prediction = make_preds_single(model, self._X_val)
    self._model = model
    self.final_loss = (
        es.best_loss.detach().cpu().numpy() if self.es and es.best_loss else np.infty
    )

plot_training_info()

Plot some summary over the MSE during training.

Source code in gresit/torch_models.py

def plot_training_info(self) -> None:
    """Plot some summary over the MSE during training."""
    sns.lmplot(
        x="trainstep",
        y="loss",
        hue="task",
        lowess=True,
        scatter_kws={"alpha": 0.5},
        line_kws={"linewidth": 2},
        data=self.training_info,
    )

    sns.lmplot(
        x="epoch",
        y="loss",
        hue="task",
        lowess=True,
        scatter_kws={"alpha": 0.5},
        line_kws={"linewidth": 2},
        data=self.validation_info,
    )

predict(X_test=None)

Make predictions.

Returns:

Type	Description
ndarray	np.ndarray: Predictions

Source code in gresit/torch_models.py

def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
    """Make predictions.

    Returns:
        np.ndarray: Predictions
    """
    if X_test is None:
        prediction = self.prediction.detach().numpy()
    else:
        X_torch_test = torch.from_numpy(X_test).float().to(self.device)
        pred = make_preds_single(self._model, X_torch_test)
        prediction = pred.detach().cpu().numpy()
    return prediction

MultitaskGPModel

Bases: ExactGP

Multitask GPR.

Source code in gresit/torch_models.py

class MultitaskGPModel(gpytorch.models.ExactGP):  # type: ignore
    """Multitask GPR."""

    def __init__(
        self,
        train_x: torch.Tensor,
        train_y: torch.Tensor,
        likelihood: gpytorch.likelihoods._GaussianLikelihoodBase,
    ) -> None:
        """Inits the model.

        Args:
            train_x (_type_): _description_
            train_y (_type_): _description_
            likelihood (_type_): _description_
        """
        super().__init__(train_x, train_y, likelihood)
        num_tasks = train_y.size(1)
        self.mean_module = gpytorch.means.MultitaskMean(
            gpytorch.means.ConstantMean(), num_tasks=num_tasks
        )
        self.covar_module = gpytorch.kernels.MultitaskKernel(
            gpytorch.kernels.RBFKernel(), num_tasks=num_tasks, rank=1
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass.

        Args:
            x (_type_): _description_

        Returns:
            _type_: _description_
        """
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)

__init__(train_x, train_y, likelihood)

Inits the model.

Parameters:

Name	Type	Description	Default
train_x	_type_	description	required
train_y	_type_	description	required
likelihood	_type_	description	required

Source code in gresit/torch_models.py

def __init__(
    self,
    train_x: torch.Tensor,
    train_y: torch.Tensor,
    likelihood: gpytorch.likelihoods._GaussianLikelihoodBase,
) -> None:
    """Inits the model.

    Args:
        train_x (_type_): _description_
        train_y (_type_): _description_
        likelihood (_type_): _description_
    """
    super().__init__(train_x, train_y, likelihood)
    num_tasks = train_y.size(1)
    self.mean_module = gpytorch.means.MultitaskMean(
        gpytorch.means.ConstantMean(), num_tasks=num_tasks
    )
    self.covar_module = gpytorch.kernels.MultitaskKernel(
        gpytorch.kernels.RBFKernel(), num_tasks=num_tasks, rank=1
    )

forward(x)

Forward pass.

Parameters:

Name	Type	Description	Default
x	_type_	description	required

Returns:

Name	Type	Description
_type_	Tensor	description

Source code in gresit/torch_models.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass.

    Args:
        x (_type_): _description_

    Returns:
        _type_: _description_
    """
    mean_x = self.mean_module(x)
    covar_x = self.covar_module(x)
    return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)

make_preds_single(model, X)

Helper function for predicting.

Parameters:

Name	Type	Description	Default
model	Model	nn model instance	required
X	Tensor	test data	required

Returns:

Name	Type	Description
_type_	Tensor	description

Source code in gresit/torch_models.py

def make_preds_single(model: nn.Module, X: torch.Tensor) -> torch.Tensor:
    """Helper function for predicting.

    Args:
        model (nn.Model): nn model instance
        X (torch.Tensor): test data

    Returns:
        _type_: _description_
    """
    # helper function to make predictions for a model
    with torch.no_grad():
        y_hat = model(X)
    return y_hat

Independence Tests

Utility classes and functions related to gresit.

Copyright (c) 2025 Robert Bosch GmbH

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.

CItest

Abstract meta class for independence tests.

Source code in gresit/independence_tests.py

class CItest(metaclass=ABCMeta):
    """Abstract meta class for independence tests."""

    def _check_input(
        self,
        x_data: np.ndarray | pd.DataFrame | pd.Series,
        y_data: np.ndarray | pd.DataFrame | pd.Series,
        z_data: np.ndarray | pd.DataFrame | pd.Series | None = None,
    ) -> None:
        if not isinstance(x_data, np.ndarray | pd.DataFrame | pd.Series):
            raise TypeError("x_data must be of type np.ndarray, pd.DataFrame, or pd.Series")
        if not isinstance(y_data, np.ndarray | pd.DataFrame | pd.Series):
            raise TypeError("y_data must be of type np.ndarray, pd.DataFrame, or pd.Series")
        if not isinstance(z_data, np.ndarray | pd.DataFrame | pd.Series) or z_data is None:
            raise TypeError("y_data must be of type np.ndarray, pd.DataFrame, or pd.Series")

    @abstractmethod
    def test(
        self,
        x_data: np.ndarray | pd.DataFrame | pd.Series,
        y_data: np.ndarray | pd.DataFrame | pd.Series,
        z_data: np.ndarray | pd.DataFrame | pd.Series | None = None,
    ) -> tuple[float, float]:
        """Abstract method for independence tests.

        Args:
            x_data (np.ndarray | pd.DataFrame | pd.Series): Variables involved in the test
            y_data (np.ndarray | pd.DataFrame | pd.Series): Variables involved in the test
            z_data (np.ndarray | pd.DataFrame | pd.Series | None): Variables involved in the test


        Returns:
            tuple[float, float]: Test statistic and corresponding pvalue (Test decision).
        """

test(x_data, y_data, z_data=None) abstractmethod

Abstract method for independence tests.

Parameters:

Name	Type	Description	Default
x_data	ndarray | DataFrame | Series	Variables involved in the test	required
y_data	ndarray | DataFrame | Series	Variables involved in the test	required
z_data	ndarray | DataFrame | Series | None	Variables involved in the test	None

Returns:

Type	Description
tuple[float, float]	tuple[float, float]: Test statistic and corresponding pvalue (Test decision).

Source code in gresit/independence_tests.py

@abstractmethod
def test(
    self,
    x_data: np.ndarray | pd.DataFrame | pd.Series,
    y_data: np.ndarray | pd.DataFrame | pd.Series,
    z_data: np.ndarray | pd.DataFrame | pd.Series | None = None,
) -> tuple[float, float]:
    """Abstract method for independence tests.

    Args:
        x_data (np.ndarray | pd.DataFrame | pd.Series): Variables involved in the test
        y_data (np.ndarray | pd.DataFrame | pd.Series): Variables involved in the test
        z_data (np.ndarray | pd.DataFrame | pd.Series | None): Variables involved in the test


    Returns:
        tuple[float, float]: Test statistic and corresponding pvalue (Test decision).
    """

DISCO

Bases: Itest

Simple Wrapper class around the squared distance covariance from the dcor class.

Source code in gresit/independence_tests.py

class DISCO(Itest):
    """Simple Wrapper class around the squared distance covariance from the dcor class."""

    def __init__(self) -> None:
        """Inits the object."""
        pass

    def test(
        self,
        x_data: np.ndarray | pd.DataFrame | pd.Series,
        y_data: np.ndarray | pd.DataFrame | pd.Series,
    ) -> tuple[float, str]:
        """Test for independence between two vectors Finite joint first moments are assumed.

        Args:
            x_data (np.ndarray | pd.DataFrame | pd.Series): x-data involved in the test
            y_data (np.ndarray | pd.DataFrame | pd.Series): y-data involved in the test

        Returns:
            tuple[float, str]: Distance covariance value and some string to comply with format.
        """
        self._check_input(x_data, y_data)
        return u_distance_covariance_sqr(x=x_data, y=y_data), "Squared Distance Covariance"

__init__()

Inits the object.

Source code in gresit/independence_tests.py

def __init__(self) -> None:
    """Inits the object."""
    pass

test(x_data, y_data)

Test for independence between two vectors Finite joint first moments are assumed.

Parameters:

Name	Type	Description	Default
x_data	ndarray | DataFrame | Series	x-data involved in the test	required
y_data	ndarray | DataFrame | Series	y-data involved in the test	required

Returns:

Type	Description
tuple[float, str]	tuple[float, str]: Distance covariance value and some string to comply with format.

Source code in gresit/independence_tests.py

def test(
    self,
    x_data: np.ndarray | pd.DataFrame | pd.Series,
    y_data: np.ndarray | pd.DataFrame | pd.Series,
) -> tuple[float, str]:
    """Test for independence between two vectors Finite joint first moments are assumed.

    Args:
        x_data (np.ndarray | pd.DataFrame | pd.Series): x-data involved in the test
        y_data (np.ndarray | pd.DataFrame | pd.Series): y-data involved in the test

    Returns:
        tuple[float, str]: Distance covariance value and some string to comply with format.
    """
    self._check_input(x_data, y_data)
    return u_distance_covariance_sqr(x=x_data, y=y_data), "Squared Distance Covariance"

FisherZVec

Bases: CItest

Simple extension of standard Fisher-Z test for independence.

Source code in gresit/independence_tests.py

class FisherZVec(CItest):
    """Simple extension of standard Fisher-Z test for independence."""

    def __init__(self) -> None:
        """Init of the object."""
        pass

    def test(
        self,
        x_data: np.ndarray,
        y_data: np.ndarray,
        z_data: np.ndarray | None = None,
        corr_threshold: float = 0.999,
    ) -> tuple[float, float]:
        """Retrieve (composite) p_value using Fisher z-transformation.

        Appropriate when data is jointly Gaussian.

        Args:
            x_data (np.ndarray): X_data.
            y_data (np.ndarray): Y_data.
            z_data (np.ndarray | None): Z_data. defaults to None.
            corr_threshold (float, optional): Threshold to make sure
                r in [-1,1]. Defaults to 0.999.

        Returns:
            tuple[float,float]: test_statistic, p_value
        """
        n = x_data.shape[0]

        if z_data is not None:
            sep_set_length = z_data.shape[1]
            corrdata = np.empty((n, 2 + z_data.shape[1], x_data.shape[1] * y_data.shape[1]))
            k = -1
            for i, j in product(range(x_data.shape[1]), range(y_data.shape[1])):
                k += 1
                corrdata[:, :, k] = np.concatenate(
                    [x_data[:, i][:, np.newaxis], y_data[:, j][:, np.newaxis], z_data], axis=1
                )
            precision_matrices = np.empty(
                (2 + z_data.shape[1], 2 + z_data.shape[1], x_data.shape[1] * y_data.shape[1])
            )
            for k in range(precision_matrices.shape[-1]):
                corrmat = np.corrcoef(corrdata[:, :, k].T)
                try:
                    precision_matrices[:, :, k] = np.linalg.inv(corrmat)
                except np.linalg.LinAlgError as error:
                    raise ValueError(
                        "The correlation matrix of your data is singular. \
                        Partial correlations cannot be estimated. Are there  \
                        collinearities in your data?"
                    ) from error

            r = np.empty(precision_matrices.shape[-1])
            for k in range(precision_matrices.shape[-1]):
                precision_matrix = precision_matrices[:, :, k]
                r[k] = (
                    -1
                    * precision_matrix[0, 1]
                    / np.sqrt(np.abs(precision_matrix[0, 0] * precision_matrix[1, 1]))
                )
        else:
            sep_set_length = 0
            uncond = []
            for i in range(x_data.shape[1]):
                uncond.append(
                    np.corrcoef(np.concatenate([x_data[:, i][:, np.newaxis], y_data], axis=1).T)[
                        1:, 0
                    ]
                )
            r = np.concatenate(uncond)

        r = np.minimum(
            corr_threshold, np.maximum(-1 * corr_threshold, r)
        )  # make r between -1 and 1
        # Fisher’s z-transform
        factor = np.sqrt(n - sep_set_length - 3)
        z_transform = factor * 0.5 * np.log((1 + r) / (1 - r))
        test_stat = factor * z_transform
        p_value = 2 * (1 - norm.cdf(abs(z_transform)))

        final_test_stat = test_stat[np.argmin(np.abs(test_stat))]
        final_p_value = np.max(p_value)

        return (float(final_test_stat), float(final_p_value))

__init__()

Init of the object.

Source code in gresit/independence_tests.py

def __init__(self) -> None:
    """Init of the object."""
    pass

test(x_data, y_data, z_data=None, corr_threshold=0.999)

Retrieve (composite) p_value using Fisher z-transformation.

Appropriate when data is jointly Gaussian.

Parameters:

Name	Type	Description	Default
x_data	ndarray	X_data.	required
y_data	ndarray	Y_data.	required
z_data	ndarray | None	Z_data. defaults to None.	None
corr_threshold	float	Threshold to make sure r in [-1,1]. Defaults to 0.999.	0.999

Returns:

Type	Description
tuple[float, float]	tuple[float,float]: test_statistic, p_value

Source code in gresit/independence_tests.py

def test(
    self,
    x_data: np.ndarray,
    y_data: np.ndarray,
    z_data: np.ndarray | None = None,
    corr_threshold: float = 0.999,
) -> tuple[float, float]:
    """Retrieve (composite) p_value using Fisher z-transformation.

    Appropriate when data is jointly Gaussian.

    Args:
        x_data (np.ndarray): X_data.
        y_data (np.ndarray): Y_data.
        z_data (np.ndarray | None): Z_data. defaults to None.
        corr_threshold (float, optional): Threshold to make sure
            r in [-1,1]. Defaults to 0.999.

    Returns:
        tuple[float,float]: test_statistic, p_value
    """
    n = x_data.shape[0]

    if z_data is not None:
        sep_set_length = z_data.shape[1]
        corrdata = np.empty((n, 2 + z_data.shape[1], x_data.shape[1] * y_data.shape[1]))
        k = -1
        for i, j in product(range(x_data.shape[1]), range(y_data.shape[1])):
            k += 1
            corrdata[:, :, k] = np.concatenate(
                [x_data[:, i][:, np.newaxis], y_data[:, j][:, np.newaxis], z_data], axis=1
            )
        precision_matrices = np.empty(
            (2 + z_data.shape[1], 2 + z_data.shape[1], x_data.shape[1] * y_data.shape[1])
        )
        for k in range(precision_matrices.shape[-1]):
            corrmat = np.corrcoef(corrdata[:, :, k].T)
            try:
                precision_matrices[:, :, k] = np.linalg.inv(corrmat)
            except np.linalg.LinAlgError as error:
                raise ValueError(
                    "The correlation matrix of your data is singular. \
                    Partial correlations cannot be estimated. Are there  \
                    collinearities in your data?"
                ) from error

        r = np.empty(precision_matrices.shape[-1])
        for k in range(precision_matrices.shape[-1]):
            precision_matrix = precision_matrices[:, :, k]
            r[k] = (
                -1
                * precision_matrix[0, 1]
                / np.sqrt(np.abs(precision_matrix[0, 0] * precision_matrix[1, 1]))
            )
    else:
        sep_set_length = 0
        uncond = []
        for i in range(x_data.shape[1]):
            uncond.append(
                np.corrcoef(np.concatenate([x_data[:, i][:, np.newaxis], y_data], axis=1).T)[
                    1:, 0
                ]
            )
        r = np.concatenate(uncond)

    r = np.minimum(
        corr_threshold, np.maximum(-1 * corr_threshold, r)
    )  # make r between -1 and 1
    # Fisher’s z-transform
    factor = np.sqrt(n - sep_set_length - 3)
    z_transform = factor * 0.5 * np.log((1 + r) / (1 - r))
    test_stat = factor * z_transform
    p_value = 2 * (1 - norm.cdf(abs(z_transform)))

    final_test_stat = test_stat[np.argmin(np.abs(test_stat))]
    final_p_value = np.max(p_value)

    return (float(final_test_stat), float(final_p_value))

HSIC

Bases: Itest

Hilbert-Schmidt Independence Criterion (HSIC) test.

Source code in gresit/independence_tests.py

class HSIC(Itest):
    """Hilbert-Schmidt Independence Criterion (HSIC) test."""

    def test(
        self,
        x_data: np.ndarray | pd.DataFrame | pd.Series,
        y_data: np.ndarray | pd.DataFrame | pd.Series,
        bw_method: str = "mdbs",
    ) -> tuple[float, float]:
        """Test for independence between two vectors.

        Args:
            x_data (np.ndarray | pd.DataFrame | pd.Series): x-data involved in the test
            y_data (np.ndarray | pd.DataFrame | pd.Series): y-data involved in the test
            bw_method (str, optional): The method used to calculate the bandwidth of the HSIC.
                * ``mdbs`` : Median distance between samples.
                * ``scott`` : Scott's Rule of Thumb.
                * ``silverman`` : Silverman's Rule of Thumb.. Defaults to "mdbs".


        Returns:
            tuple[float, float]: Test statistic and corresponding pvalue.
        """
        self._check_input(x_data, y_data)
        x_data = x_data if isinstance(x_data, np.ndarray) else x_data.values
        y_data = y_data if isinstance(y_data, np.ndarray) else y_data.values
        test_stat, p_value = self.hsic_test_gamma(X=x_data, Y=y_data, bw_method=bw_method)
        return test_stat, p_value

    def get_kernel_width(self, X: np.ndarray, sample_cut: int = 100) -> np.float64:
        """Calculate the bandwidth to median distance between points.

        Use at most 100 points (since median is only a heuristic,
        and 100 points is sufficient for a robust estimate).

        Args:
            X (np.ndarray): shape (n_samples, n_features) Training data,
            sample_cut (int, optional): Number of samples to use for bandwidth calculation.

        Returns:
            float: The bandwidth parameter.
        """
        n_samples = X.shape[0]
        if n_samples > sample_cut:
            X_med = X[:sample_cut, :]
            n_samples = sample_cut
        else:
            X_med = X

        G = np.sum(X_med * X_med, 1).reshape(n_samples, 1)
        dists = G + G.T - 2 * np.dot(X_med, X_med.T)
        dists = dists - np.tril(dists)
        dists = dists.reshape(n_samples**2, 1)

        return np.sqrt(0.5 * np.median(dists[dists > 0]))

    def _rbf_dot(self, X: np.ndarray, width: float) -> np.ndarray:
        """Calculate rbf dot, in special case with X dot X.

        Args:
            X (np.ndarray): data
            width (float): bandwidth parameter

        Returns:
            np.ndarray: Kernel matrix.
        """
        G = np.sum(X * X, axis=1)
        H = G[None, :] + G[:, None] - 2 * np.dot(X, X.T)
        return np.exp(-H / 2 / (width**2))

    def get_gram_matrix(self, X: np.ndarray, width: float) -> tuple[np.ndarray, np.ndarray]:
        """Get the centered gram matrices.

        Args:
            X (np.ndarray): shape (n_samples, n_features)
                Training data, where ``n_samples`` is the number of samples
                and ``n_features`` is the number of features.
            width (float): The bandwidth parameter.

        Returns:
            tuple[np.ndarray, np.ndarray]: The centered gram matrices.
        """
        n = X.shape[0]

        K = self._rbf_dot(X, width)
        K_colsums = K.sum(axis=0)
        K_rowsums = K.sum(axis=1)
        K_allsum = K_rowsums.sum()
        Kc = K - (K_colsums[None, :] + K_rowsums[:, None]) / n + (K_allsum / n**2)
        return K, Kc

    def hsic_teststat(self, Kc: np.ndarray, Lc: np.ndarray, n: int) -> np.float64:
        """Get the HSIC statistic.

        Args:
            Kc (np.ndarray): Centered gram matrix.
            Lc (np.ndarray): Centered gram matrix.
            n (int): Sample size.

        Returns:
            float: HSIC statistic.
        """
        # test statistic m*HSICb under H1
        return 1 / n * np.sum(Kc.T * Lc)

    def hsic_test_gamma(
        self, X: np.ndarray, Y: np.ndarray, bw_method: str = "mdbs"
    ) -> tuple[float, float]:
        """Get the HSIC statistic and p-value.

        Args:
            X (np.ndarray): data, possibly vector-valued.
            Y (np.ndarray): data, possibly vector-valued.
            bw_method (str, optional): The method used to calculate the bandwidth of the HSIC.
                * ``mdbs`` : Median distance between samples.
                * ``scott`` : Scott's Rule of Thumb.
                * ``silverman`` : Silverman's Rule of Thumb.. Defaults to "mdbs".

        Returns:
            tuple[float, float]: HSIC test statistic and corresponding p-value
        """
        X = X.reshape(-1, 1) if X.ndim == 1 else X
        Y = Y.reshape(-1, 1) if Y.ndim == 1 else Y

        if bw_method == "scott":
            width_x = bandwidths.bw_scott(X)
            width_y = bandwidths.bw_scott(Y)
        elif bw_method == "silverman":
            width_x = bandwidths.bw_silverman(X)
            width_y = bandwidths.bw_silverman(Y)
        # Get kernel width to median distance between points
        else:
            width_x = self.get_kernel_width(X)
            width_y = self.get_kernel_width(Y)

        # these are slightly biased estimates of centered gram matrices
        K, Kc = self.get_gram_matrix(X, width_x)
        L, Lc = self.get_gram_matrix(Y, width_y)

        # test statistic m*HSICb under H1
        n = X.shape[0]
        test_stat = self.hsic_teststat(Kc, Lc, n)

        var = (1 / 6 * Kc * Lc) ** 2
        # second subtracted term is bias correction
        var = 1 / n / (n - 1) * (np.sum(var) - np.trace(var))
        # variance under H0
        var = 72 * (n - 4) * (n - 5) / n / (n - 1) / (n - 2) / (n - 3) * var

        K[np.diag_indices(n)] = 0
        L[np.diag_indices(n)] = 0
        mu_X = 1 / n / (n - 1) * K.sum()
        mu_Y = 1 / n / (n - 1) * L.sum()
        # mean under H0
        mean = 1 / n * (1 + mu_X * mu_Y - mu_X - mu_Y)

        alpha = mean**2 / var
        # threshold for hsicArr*m
        beta = var * n / mean
        p = gamma.sf(test_stat, alpha, scale=beta)

        return test_stat, p

get_gram_matrix(X, width)

Get the centered gram matrices.

Parameters:

Name	Type	Description	Default
X	ndarray	shape (n_samples, n_features) Training data, where n_samples is the number of samples and n_features is the number of features.	required
width	float	The bandwidth parameter.	required

Returns:

Type	Description
tuple[ndarray, ndarray]	tuple[np.ndarray, np.ndarray]: The centered gram matrices.

Source code in gresit/independence_tests.py

def get_gram_matrix(self, X: np.ndarray, width: float) -> tuple[np.ndarray, np.ndarray]:
    """Get the centered gram matrices.

    Args:
        X (np.ndarray): shape (n_samples, n_features)
            Training data, where ``n_samples`` is the number of samples
            and ``n_features`` is the number of features.
        width (float): The bandwidth parameter.

    Returns:
        tuple[np.ndarray, np.ndarray]: The centered gram matrices.
    """
    n = X.shape[0]

    K = self._rbf_dot(X, width)
    K_colsums = K.sum(axis=0)
    K_rowsums = K.sum(axis=1)
    K_allsum = K_rowsums.sum()
    Kc = K - (K_colsums[None, :] + K_rowsums[:, None]) / n + (K_allsum / n**2)
    return K, Kc

get_kernel_width(X, sample_cut=100)

Calculate the bandwidth to median distance between points.

Use at most 100 points (since median is only a heuristic, and 100 points is sufficient for a robust estimate).

Parameters:

Name	Type	Description	Default
X	ndarray	shape (n_samples, n_features) Training data,	required
sample_cut	int	Number of samples to use for bandwidth calculation.	100

Returns:

Name	Type	Description
float	float64	The bandwidth parameter.

Source code in gresit/independence_tests.py

def get_kernel_width(self, X: np.ndarray, sample_cut: int = 100) -> np.float64:
    """Calculate the bandwidth to median distance between points.

    Use at most 100 points (since median is only a heuristic,
    and 100 points is sufficient for a robust estimate).

    Args:
        X (np.ndarray): shape (n_samples, n_features) Training data,
        sample_cut (int, optional): Number of samples to use for bandwidth calculation.

    Returns:
        float: The bandwidth parameter.
    """
    n_samples = X.shape[0]
    if n_samples > sample_cut:
        X_med = X[:sample_cut, :]
        n_samples = sample_cut
    else:
        X_med = X

    G = np.sum(X_med * X_med, 1).reshape(n_samples, 1)
    dists = G + G.T - 2 * np.dot(X_med, X_med.T)
    dists = dists - np.tril(dists)
    dists = dists.reshape(n_samples**2, 1)

    return np.sqrt(0.5 * np.median(dists[dists > 0]))

hsic_test_gamma(X, Y, bw_method='mdbs')

Get the HSIC statistic and p-value.

Parameters:

Name	Type	Description	Default
X	ndarray	data, possibly vector-valued.	required
Y	ndarray	data, possibly vector-valued.	required
bw_method	str	The method used to calculate the bandwidth of the HSIC. * mdbs : Median distance between samples. * scott : Scott's Rule of Thumb. * silverman : Silverman's Rule of Thumb.. Defaults to "mdbs".	'mdbs'

Returns:

Type	Description
tuple[float, float]	tuple[float, float]: HSIC test statistic and corresponding p-value

Source code in gresit/independence_tests.py

def hsic_test_gamma(
    self, X: np.ndarray, Y: np.ndarray, bw_method: str = "mdbs"
) -> tuple[float, float]:
    """Get the HSIC statistic and p-value.

    Args:
        X (np.ndarray): data, possibly vector-valued.
        Y (np.ndarray): data, possibly vector-valued.
        bw_method (str, optional): The method used to calculate the bandwidth of the HSIC.
            * ``mdbs`` : Median distance between samples.
            * ``scott`` : Scott's Rule of Thumb.
            * ``silverman`` : Silverman's Rule of Thumb.. Defaults to "mdbs".

    Returns:
        tuple[float, float]: HSIC test statistic and corresponding p-value
    """
    X = X.reshape(-1, 1) if X.ndim == 1 else X
    Y = Y.reshape(-1, 1) if Y.ndim == 1 else Y

    if bw_method == "scott":
        width_x = bandwidths.bw_scott(X)
        width_y = bandwidths.bw_scott(Y)
    elif bw_method == "silverman":
        width_x = bandwidths.bw_silverman(X)
        width_y = bandwidths.bw_silverman(Y)
    # Get kernel width to median distance between points
    else:
        width_x = self.get_kernel_width(X)
        width_y = self.get_kernel_width(Y)

    # these are slightly biased estimates of centered gram matrices
    K, Kc = self.get_gram_matrix(X, width_x)
    L, Lc = self.get_gram_matrix(Y, width_y)

    # test statistic m*HSICb under H1
    n = X.shape[0]
    test_stat = self.hsic_teststat(Kc, Lc, n)

    var = (1 / 6 * Kc * Lc) ** 2
    # second subtracted term is bias correction
    var = 1 / n / (n - 1) * (np.sum(var) - np.trace(var))
    # variance under H0
    var = 72 * (n - 4) * (n - 5) / n / (n - 1) / (n - 2) / (n - 3) * var

    K[np.diag_indices(n)] = 0
    L[np.diag_indices(n)] = 0
    mu_X = 1 / n / (n - 1) * K.sum()
    mu_Y = 1 / n / (n - 1) * L.sum()
    # mean under H0
    mean = 1 / n * (1 + mu_X * mu_Y - mu_X - mu_Y)

    alpha = mean**2 / var
    # threshold for hsicArr*m
    beta = var * n / mean
    p = gamma.sf(test_stat, alpha, scale=beta)

    return test_stat, p

hsic_teststat(Kc, Lc, n)

Get the HSIC statistic.

Parameters:

Name	Type	Description	Default
Kc	ndarray	Centered gram matrix.	required
Lc	ndarray	Centered gram matrix.	required
n	int	Sample size.	required

Returns:

Name	Type	Description
float	float64	HSIC statistic.

Source code in gresit/independence_tests.py

def hsic_teststat(self, Kc: np.ndarray, Lc: np.ndarray, n: int) -> np.float64:
    """Get the HSIC statistic.

    Args:
        Kc (np.ndarray): Centered gram matrix.
        Lc (np.ndarray): Centered gram matrix.
        n (int): Sample size.

    Returns:
        float: HSIC statistic.
    """
    # test statistic m*HSICb under H1
    return 1 / n * np.sum(Kc.T * Lc)

test(x_data, y_data, bw_method='mdbs')

Test for independence between two vectors.

Parameters:

Name	Type	Description	Default
x_data	ndarray | DataFrame | Series	x-data involved in the test	required
y_data	ndarray | DataFrame | Series	y-data involved in the test	required
bw_method	str	The method used to calculate the bandwidth of the HSIC. * mdbs : Median distance between samples. * scott : Scott's Rule of Thumb. * silverman : Silverman's Rule of Thumb.. Defaults to "mdbs".	'mdbs'

Returns:

Type	Description
tuple[float, float]	tuple[float, float]: Test statistic and corresponding pvalue.

Source code in gresit/independence_tests.py

def test(
    self,
    x_data: np.ndarray | pd.DataFrame | pd.Series,
    y_data: np.ndarray | pd.DataFrame | pd.Series,
    bw_method: str = "mdbs",
) -> tuple[float, float]:
    """Test for independence between two vectors.

    Args:
        x_data (np.ndarray | pd.DataFrame | pd.Series): x-data involved in the test
        y_data (np.ndarray | pd.DataFrame | pd.Series): y-data involved in the test
        bw_method (str, optional): The method used to calculate the bandwidth of the HSIC.
            * ``mdbs`` : Median distance between samples.
            * ``scott`` : Scott's Rule of Thumb.
            * ``silverman`` : Silverman's Rule of Thumb.. Defaults to "mdbs".


    Returns:
        tuple[float, float]: Test statistic and corresponding pvalue.
    """
    self._check_input(x_data, y_data)
    x_data = x_data if isinstance(x_data, np.ndarray) else x_data.values
    y_data = y_data if isinstance(y_data, np.ndarray) else y_data.values
    test_stat, p_value = self.hsic_test_gamma(X=x_data, Y=y_data, bw_method=bw_method)
    return test_stat, p_value

Itest

Abstract meta class for independence tests.

Source code in gresit/independence_tests.py

class Itest(metaclass=ABCMeta):
    """Abstract meta class for independence tests."""

    def _check_input(
        self,
        x_data: np.ndarray | pd.DataFrame | pd.Series,
        y_data: np.ndarray | pd.DataFrame | pd.Series,
    ) -> None:
        if not isinstance(x_data, np.ndarray | pd.DataFrame | pd.Series):
            raise TypeError("x_data must be of type np.ndarray, pd.DataFrame, or pd.Series")
        if not isinstance(y_data, np.ndarray | pd.DataFrame | pd.Series):
            raise TypeError("y_data must be of type np.ndarray, pd.DataFrame, or pd.Series")

    @abstractmethod
    def test(
        self,
        x_data: np.ndarray | pd.DataFrame | pd.Series,
        y_data: np.ndarray | pd.DataFrame | pd.Series,
    ) -> tuple[float, float | str]:
        """Abstract method for independence tests.

        Args:
            x_data (np.ndarray | pd.DataFrame | pd.Series): Variables involved in the test
            y_data (np.ndarray | pd.DataFrame | pd.Series): Variables involved in the test


        Returns:
            tuple[float, float | str]: Test statistic and corresponding pvalue (Test decision).
        """

test(x_data, y_data) abstractmethod

Abstract method for independence tests.

Parameters:

Name	Type	Description	Default
x_data	ndarray | DataFrame | Series	Variables involved in the test	required
y_data	ndarray | DataFrame | Series	Variables involved in the test	required

Returns:

Type	Description
tuple[float, float | str]	tuple[float, float | str]: Test statistic and corresponding pvalue (Test decision).

Source code in gresit/independence_tests.py

@abstractmethod
def test(
    self,
    x_data: np.ndarray | pd.DataFrame | pd.Series,
    y_data: np.ndarray | pd.DataFrame | pd.Series,
) -> tuple[float, float | str]:
    """Abstract method for independence tests.

    Args:
        x_data (np.ndarray | pd.DataFrame | pd.Series): Variables involved in the test
        y_data (np.ndarray | pd.DataFrame | pd.Series): Variables involved in the test


    Returns:
        tuple[float, float | str]: Test statistic and corresponding pvalue (Test decision).
    """

KernelCI

Bases: CItest

Kernel HSIC wrapper around causal-learn.

Source code in gresit/independence_tests.py

class KernelCI(CItest):
    """Kernel HSIC wrapper around causal-learn."""

    def __init__(self) -> None:
        """Init of the object."""
        pass

    def test(
        self,
        x_data: np.ndarray,
        y_data: np.ndarray,
        z_data: np.ndarray | None = None,
    ) -> tuple[float, float]:
        """KCI test wrapper.

        Args:
            x_data (np.ndarray): _description_
            y_data (np.ndarray): _description_
            z_data (np.ndarray | None, optional): _description_. Defaults to None.

        Returns:
            tuple[float, float]: Test statistic and p_value.
        """
        if z_data is None:
            unconditional_test = KCI.KCI_UInd()
            p_value, test_stat = unconditional_test.compute_pvalue(data_x=x_data, data_y=y_data)
        else:
            conditional_test = KCI.KCI_CInd()
            p_value, test_stat = conditional_test.compute_pvalue(
                data_x=x_data, data_y=y_data, data_z=z_data
            )

        return float(test_stat), float(p_value)

__init__()

Init of the object.

Source code in gresit/independence_tests.py

def __init__(self) -> None:
    """Init of the object."""
    pass

test(x_data, y_data, z_data=None)

KCI test wrapper.

Parameters:

Name	Type	Description	Default
x_data	ndarray	description	required
y_data	ndarray	description	required
z_data	ndarray | None	description. Defaults to None.	None

Returns:

Type	Description
tuple[float, float]	tuple[float, float]: Test statistic and p_value.

Source code in gresit/independence_tests.py

def test(
    self,
    x_data: np.ndarray,
    y_data: np.ndarray,
    z_data: np.ndarray | None = None,
) -> tuple[float, float]:
    """KCI test wrapper.

    Args:
        x_data (np.ndarray): _description_
        y_data (np.ndarray): _description_
        z_data (np.ndarray | None, optional): _description_. Defaults to None.

    Returns:
        tuple[float, float]: Test statistic and p_value.
    """
    if z_data is None:
        unconditional_test = KCI.KCI_UInd()
        p_value, test_stat = unconditional_test.compute_pvalue(data_x=x_data, data_y=y_data)
    else:
        conditional_test = KCI.KCI_CInd()
        p_value, test_stat = conditional_test.compute_pvalue(
            data_x=x_data, data_y=y_data, data_z=z_data
        )

    return float(test_stat), float(p_value)

xi_vec_corr(x, y)

Compute the correlation coefficient between x and y.

according to the xi coefficient defined by Chatterjee.

Parameters:

Name	Type	Description	Default
x	ndarray	description	required
y	ndarray	description	required

Source code in gresit/independence_tests.py

def xi_vec_corr(x: np.ndarray, y: np.ndarray) -> float:
    """Compute the correlation coefficient between x and y.

    according to the xi coefficient defined by Chatterjee.

    Args:
        x (np.ndarray): _description_
        y (np.ndarray): _description_
    """

    def rank_order(vector: np.ndarray) -> list[np.ndarray]:
        random_index = np.random.choice(np.arange(length), length, replace=False)
        randomized_vector = vector[random_index]
        ranked_vector = rankdata(randomized_vector, method="ordinal")
        answer = [ranked_vector[j] for _, j in sorted(zip(random_index, range(length)))]
        return answer

    def compute_d_sequence(y: np.ndarray) -> float:
        ell = rankdata([-i for i in y], method="max")
        return float(np.sum(ell * (length - ell)) / (length**3))

    def compute_xi_coefficient(vector: np.ndarray) -> float:
        mean_absolute = np.sum(np.abs([a - b for a, b in zip(vector[:-1], vector[1:])]))
        return float(1 - mean_absolute / (2 * (length**2) * d_sequence))

    def distance_transform(x: np.ndarray) -> np.ndarray:
        n = x.shape[0]
        m = n * (n + 1) / 2

        if x.ndim == 1:
            diff = x[np.newaxis] - x[..., np.newaxis]
            a = np.linalg.norm(diff[..., np.newaxis], axis=-1)
        else:
            diff = x[:, np.newaxis, :] - x[np.newaxis, :, :]
            a = np.linalg.norm(diff, axis=-1)

        H = np.eye(n) - 1.0 / n * np.ones((n, n))
        A = H @ a @ H
        indices = np.triu_indices(A.shape[0])
        D_kx = A[indices]
        if D_kx.shape[0] != m:
            raise ValueError("Shapes do not agree,")
        return D_kx

    x_transformed = distance_transform(x)
    y_transformed = distance_transform(y)

    length = len(x_transformed)
    x_ordered = np.argsort(rank_order(x_transformed))
    y_rank_max = rankdata(y_transformed, method="max")
    x_ordered_max_rank = y_rank_max[x_ordered]
    d_sequence = compute_d_sequence(y_transformed)
    correlation = compute_xi_coefficient(x_ordered_max_rank)

    return float(correlation)

Utility classes and functions related to gresit.

Copyright (c) 2025 Robert Bosch GmbH

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.

ConvergenceError

Bases: Exception

Convenience class for convergence error.

Source code in gresit/model_selection.py

class ConvergenceError(Exception):
    """Convenience class for convergence error."""

    def __init__(self, message: str = "Convergence not reached"):
        """Returns error message."""
        super().__init__(message)

__init__(message='Convergence not reached')

Returns error message.

Source code in gresit/model_selection.py

def __init__(self, message: str = "Convergence not reached"):
    """Returns error message."""
    super().__init__(message)

MURGS

Multi-Response Group Sparse Additive Mode (MURGS) class.

Source code in gresit/model_selection.py

class MURGS:
    """Multi-Response Group Sparse Additive Mode (MURGS) class."""

    def __init__(self, group_names: list[str] | None = None) -> None:
        """Initializes the object."""
        self.zero_groups: list[bool]
        self.zero_group_history: dict[float, list[bool]] = {}
        self.d_g_long: np.ndarray
        self.p_g: int = 0
        self.n_tasks: int
        self.smoothing_matrices: np.ndarray | dict[str, np.ndarray]
        self.f_g_hat: np.ndarray | dict[str, np.ndarray]
        self.chosen_lambda: float = 0.0
        self.lambda_max_value: float | None = None
        self.steps_till_convergence: int = -1
        self.max_iter: int = 10000
        self.tol: float = 1e-8
        if group_names is not None:
            self.group_names = group_names
        else:
            self.group_names = []
        self.predicted_vals: np.ndarray

    def __repr__(self) -> str:
        """Print some useful information.

        Returns:
            str: _description_
        """
        return f"MT-GSpAM on {self.p_g} groups."

    def __str__(self) -> str:
        """Print some useful summary.

        Returns:
            str: _description_
        """
        display_limit = 5
        setting = {
            "Number of Groups: ": self.p_g,
            "Group sizes homogenous: ": True if isinstance(self.f_g_hat, np.ndarray) else False,
            "Steps until convergence: ": self.steps_till_convergence,
            "Final lambda: ": self.chosen_lambda,
            "Number of non-zero groups: ": self.p_g - np.asanyarray(self.zero_groups).sum(),
            "Non-zero group names: ": self.return_nonzero_groups()
            if len(self.return_nonzero_groups()) <= display_limit
            else "Too many to show",
        }
        s = ""
        for info, info_text in setting.items():
            s += f"{info}{info_text}\n"

        return s

    def return_nonzero_groups(self) -> list[str]:
        """Return the group names of the nonzero groups.

        Returns:
            list[str] | str: List of nonzero groups with their actual names if given.
        """
        if len(self.group_names) > 0:
            nonzero_groups = [
                group
                for group, zero_group in zip(self.group_names, self.zero_groups)
                if not zero_group
            ]
        else:
            nonzero_groups = ["You neither provided any group names nor ran the model."]

        return nonzero_groups

    def precalculate_smooths(
        self, X_data: np.ndarray | dict[str, np.ndarray], local_regression_method: str = "kernel"
    ) -> np.ndarray | dict[str, np.ndarray]:
        """Precalculate smoother matrices.

        Input may be both `np.ndarray` and `dict`.

        Args:
            X_data (np.ndarray | dict[str, np.ndarray]): predictor data.
            local_regression_method (str): Method to use to calculate smoother matrix. Options
                currently are `"loess"` and `"kernel"`. When kernel is chosen, the default is set to
                Gaussian kernel regression with the standard deviation rule for bandwidth selection.

        Returns:
            np.ndarray | dict[str, np.ndarray]: smooths.
        """
        if isinstance(X_data, dict):
            X_data = self._dict_preprocessing(X_data=X_data)
            smoothing_matrices = {}
            for group, data in X_data.items():
                smoothing_matrices[group] = np.zeros((data.shape[0], data.shape[0], data.shape[1]))
                for j in range(data.shape[1]):
                    if local_regression_method == "loess":
                        smoothing_matrices[group][:, :, j] = self._make_loess_smoother_matrix(
                            data[:, j]
                        )
                    elif local_regression_method == "kernel":
                        smoothing_matrices[group][:, :, j] = (
                            self._make_gaussian_kernel_smoother_matrix(data[:, j])
                        )
                    else:
                        raise NotImplementedError("This smoothing method is not implemented yet.")
        else:
            smoothing_matrices = np.zeros(
                (X_data.shape[0], X_data.shape[0], X_data.shape[1], X_data.shape[2]),
                dtype=np.float64,
            )
            inner_smoother: np.ndarray = smoothing_matrices
            for num_groups in range(X_data.shape[2]):
                for group_members in range(X_data.shape[1]):
                    if local_regression_method == "loess":
                        inner_smoother[:, :, group_members, num_groups] = (
                            self._make_loess_smoother_matrix(X_data[:, group_members, num_groups])
                        )
                    elif local_regression_method == "kernel":
                        inner_smoother[:, :, group_members, num_groups] = (
                            self._make_gaussian_kernel_smoother_matrix(
                                X_data[:, group_members, num_groups]
                            )
                        )
                    else:
                        raise NotImplementedError("This smoothing method is not implemented yet.")
            smoothing_matrices = inner_smoother

        self.smoothing_matrices = smoothing_matrices
        return smoothing_matrices

    def _init_functions(
        self, X_data: np.ndarray | dict[str, np.ndarray]
    ) -> np.ndarray | dict[str, np.ndarray]:
        """Initiate additive functional components.

        Args:
            X_data (np.ndarray | dict[str, np.ndarray]): predictor data.

        Returns:
            np.ndarray | dict[str, np.ndarray]: functions initiated.
        """
        if isinstance(X_data, dict):
            f_g = {}
            for group, data in X_data.items():
                f_g[group] = np.zeros(data.shape + (self.n_tasks,))
        else:
            f_g = np.zeros(X_data.shape + (self.n_tasks,))
        return f_g

    def _set_dg_pg(self, X_data: np.ndarray | dict[str, np.ndarray]) -> None:
        """Set group sizes and number of groups.

        Args:
            X_data (np.ndarray | dict[str, np.ndarray]): predictor data.
        """
        if isinstance(X_data, dict):
            self.d_g_long = np.array(
                [dat.shape[1] if len(dat.shape) > 1 else 1 for dat in X_data.values()]
            )
            # Number of predictors
            self.p_g = len(self.d_g_long)
        else:
            # Number of predictors
            self.p_g = X_data.shape[-1]

            # Number of variables per group
            self.d_g_long = np.repeat(X_data.shape[1], self.p_g)
            if not len(self.group_names) == self.p_g:
                self.group_names = [str(intgr) for intgr in range(self.p_g)]

    def _dict_preprocessing(self, X_data: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
        """Dict preprocessing dealing with univariate groups.

        Args:
            X_data (dict[str, np.ndarray]): Predictor dict.

        Returns:
            dict[str, np.ndarray]: Predictor dict with axis added when univariate.
        """
        for key, value in X_data.items():
            X_data[key] = self._assure_two_dim(value)
        return X_data

    def _get_m_star(self, s_g_ordered: np.ndarray, penalty: float, d_g: int) -> int:
        """Get m* so as to differentiate the cases.

        Cases involve instanecs where more than one group attains the sup-norm.

        Args:
            s_g_ordered (np.ndarray): Ordered functional norm of group `g`.
            penalty (float): current penalty parameter.
            d_g (int): Number of predictors in group.

        Returns:
            int: m* value.
        """
        m_criterion = s_g_ordered.cumsum() - penalty * np.sqrt(d_g)
        m_prefix = np.array(range(1, m_criterion.shape[0] + 1))
        return int(m_prefix[np.argmax(m_criterion / m_prefix)])

    def smoother_direct_fit(
        self,
        g: int,
        d_g: int,
        R_g: np.ndarray,
        X_data: np.ndarray | dict[str, np.ndarray],
        local_regression_method: str,
    ) -> np.ndarray:
        """Direct fit of smoothing.

        Args:
            g (int): _description_
            d_g (int): _description_
            R_g (np.ndarray): _description_
            X_data (np.ndarray): _description_
            local_regression_method (str): _description_

        Raises:
            NotImplementedError: _description_

        Returns:
            _type_: _description_
        """
        if local_regression_method == "loess":
            smooth_fit = self.loess_direct_fit(g=g, d_g=d_g, R_g=R_g, X_data=X_data)
        elif local_regression_method == "kernel":
            smooth_fit = self.gaussian_kernel_direct_fit(g=g, d_g=d_g, R_g=R_g, X_data=X_data)
        else:
            raise NotImplementedError("Smoothing method not implemented.")
        return smooth_fit

    def _init_or_insert_functions(
        self,
        X_data: np.ndarray | dict[str, np.ndarray],
        warm_start_f_hat: np.ndarray | dict[str, np.ndarray] | None = None,
    ) -> np.ndarray | dict[str, np.ndarray]:
        if warm_start_f_hat is None:
            f_g = self._init_functions(X_data=X_data)
        else:
            f_g = warm_start_f_hat
        return f_g

    def block_coordinate_descent(
        self,
        X_data: np.ndarray | dict[str, np.ndarray],
        Y_data: np.ndarray,
        penalty: float,
        precalculate_smooths: bool = True,
        smoothers: np.ndarray | dict[str, np.ndarray] | None = None,
        local_regression_method: str = "kernel",
        warm_start_f_hat: np.ndarray | dict[str, np.ndarray] | None = None,
    ) -> None:
        """Block coordinate descent for multitask Sparse Group Lasso.

        Args:
            X_data (np.ndarray | dict[str, np.ndarray]): _description_
            Y_data (np.ndarray): _description_
            penalty (float): _description_
            precalculate_smooths (bool, optional): _description_. Defaults to True.
            smoothers (np.ndarray | dict[str, np.ndarray] | None, optional): _description_.
                Defaults to None.
            local_regression_method (str): Defaults to "loess". Other options currently: "kernel".
            warm_start_f_hat (np.ndarray | dict[str, np.ndarray] | None, optional): _description_.
                Defaults to None.
        """
        if smoothers is None:
            smoothers = self.precalculate_smooths(
                X_data=X_data, local_regression_method=local_regression_method
            )

        self.n_tasks = Y_data.shape[1]

        if isinstance(X_data, dict):
            if not self.group_names:
                self.group_names = list(X_data.keys())
            X_data = self._dict_preprocessing(X_data=X_data)

        # Number of predictors and members in groups
        self._set_dg_pg(X_data=X_data)

        f_g = self._init_or_insert_functions(X_data, warm_start_f_hat)

        divergent_max_runs = 10
        divergent_runs = 0
        max_inc_old = 0.0
        new_functional_norm = np.zeros((self.p_g, self.n_tasks))
        # Update loop
        for t in range(self.max_iter):
            old_functional_norm = new_functional_norm
            zero_groups = [False] * self.p_g
            for g in range(self.p_g):
                d_g = self.d_g_long[g]
                # R_g_hat has shape (#samples, #tasks)
                R_g_hat = self.R_g_hat_update(f_g=f_g, Y_data=Y_data, g=g)

                if precalculate_smooths:
                    smooth_fit = self.predict_from_linear_smoother(
                        g=g, smoothing_matrices=smoothers, R_g=R_g_hat
                    )
                else:
                    smooth_fit = self.smoother_direct_fit(
                        g=g,
                        d_g=d_g,
                        X_data=X_data,
                        R_g=R_g_hat,
                        local_regression_method=local_regression_method,
                    )

                # # estimate of || Q^(k)R_g^(k) ||
                # # shape (1, #tasks)
                omega_g = self.functional_norm(smooth_fit)

                # sort in descending order for each task
                s_g_ordered = np.sort(omega_g)[::-1]

                # get m* for the case that more than one group attains the sup-norm.
                m_opt = self._get_m_star(s_g_ordered=s_g_ordered, penalty=penalty, d_g=d_g)

                # calculate nonzero groups
                zero_groups[g] = omega_g.sum() <= penalty * np.sqrt(d_g)
                f_g = self.soft_thresholding_update(
                    g=g,
                    f_g=f_g,
                    zero_groups=zero_groups,
                    smooth_fit=smooth_fit,
                    s_g_ordered=s_g_ordered,
                    penalty=penalty,
                    m_opt=m_opt,
                )

            new_functional_norm = self.omega_hat(f_g)
            max_inc: float = np.sqrt(
                np.square(new_functional_norm - old_functional_norm).sum(axis=(0, 1))
            )
            if max_inc_old < max_inc:
                divergent_runs += 1

            if divergent_runs > divergent_max_runs:
                raise ConvergenceError(
                    f"Model did not converge in group {g}. Penalty: {penalty} too low?"
                )
            max_inc_old = max_inc

            if np.all(max_inc < self.tol / 2):
                break

        self.f_g_hat = f_g
        self.zero_groups = zero_groups
        self.steps_till_convergence = t

    def R_g_hat_update(
        self, f_g: np.ndarray | dict[str, np.ndarray], Y_data: np.ndarray, g: int
    ) -> np.ndarray:
        """Update partial residuals.

        Args:
            f_g (np.ndarray): Current additive components
                of shape (#samples, #groups, #predictors, #tasks)
            Y_data (np.ndarray): Response data of shape (#samples, #tasks)
            g (int): group in question.

        Returns:
            np.ndarray: Partial Residuals.
        """
        if isinstance(f_g, np.ndarray):
            mask = np.ones(self.p_g, bool)
            mask[g] = False
            R_g = Y_data - f_g[:, :, mask, :].sum(axis=(1, 2))
        else:
            groups = list(f_g.keys())
            f_interim_sum = np.asanyarray(
                [value.sum(axis=1) for key, value in f_g.items() if key != groups[g]]
            )
            R_g = Y_data - f_interim_sum.sum(axis=0)

        return R_g

    def predict_from_linear_smoother(
        self,
        g: int,
        smoothing_matrices: np.ndarray | dict[str, np.ndarray],
        R_g: np.ndarray,
    ) -> np.ndarray:
        """Local regression fit.

        Args:
            g (int): group in question.
            R_g (np.ndarray): Partial residuals should be of
                shape (#n_samples. #n_predictors, #n_tasks)
            smoothing_matrices (np.ndarray): Precalculated smoothing matrices
                of shape (#n_samples, #n_samples #n_groups, #n_predictors)

        Returns:
            torch.Tensor: Fitted local regressions.
        """
        if isinstance(smoothing_matrices, np.ndarray):
            smoothing_matrix = smoothing_matrices[:, :, :, g]
        else:
            smoothing_matrix = list(smoothing_matrices.values())[g]

        return torch.einsum(
            "ijk, jl -> ikl", torch.from_numpy(smoothing_matrix), torch.from_numpy(R_g)
        ).numpy()

    def loess_direct_fit(
        self, g: int, d_g: int, R_g: np.ndarray, X_data: np.ndarray | dict[str, np.ndarray]
    ) -> np.ndarray:
        """Local regression fit.

        Args:
            g (int): group in question.
            d_g (int): Number of predictors in group.
            R_g (np.ndarray): Partial residuals should be of
                shape (#n_samples. #n_predictors, #n_tasks)
            X_data (np.ndarray): Training data of shape (#n_samples. #n_groups, #n_predictors)

        Returns:
            np.ndarray: Fitted local regressions.
        """
        if isinstance(X_data, np.ndarray):
            data = X_data[:, :, g]
        else:
            data = list(X_data.values())[g]

        loess_g = np.zeros((R_g.shape[0], d_g, self.n_tasks))
        for k in range(self.n_tasks):
            for j in range(d_g):
                try:
                    loess_obj = loess(data[:, j], R_g[:, k])
                except ValueError:
                    loess_obj = loess(data[:, j], R_g[:, k], surface="direct")
                loess_obj.fit()
                loess_g[:, j, k] = loess_obj.outputs.fitted_values
        return loess_g

    def gaussian_kernel_direct_fit(
        self, g: int, d_g: int, R_g: np.ndarray, X_data: np.ndarray | dict[str, np.ndarray]
    ) -> np.ndarray:
        """Local regression fit.

        Args:
            g (int): group in question.
            d_g (int): Number of predictors in group.
            R_g (np.ndarray): Partial residuals should be of
                shape (#n_samples. #n_predictors, #n_tasks)
            X_data (np.ndarray): Training data of shape (#n_samples. #n_groups, #n_predictors)

        Returns:
            np.ndarray: Fitted local regressions.
        """
        if isinstance(X_data, np.ndarray):
            data = X_data[:, :, g]
        else:
            data = list(X_data.values())[g]

        kernel_reg_g = np.zeros((R_g.shape[0], d_g, self.n_tasks))
        for k in range(self.n_tasks):
            for j in range(d_g):
                kernel_reg_g[:, j, k] = self._torch_gaussian_kernel_regression(
                    x_data=data[:, j],
                    y_data=R_g[:, k],
                    bandwidth=self.plugin_bandwidth(x_j=data[:, j]),
                )

        return kernel_reg_g

    def _torch_gaussian_kernel_regression(
        self, x_data: np.ndarray, y_data: np.ndarray, bandwidth: float
    ) -> np.ndarray:
        """Torch implementation of Gaussian kernel regression.

        Args:
            x_data (np.ndarray): X data
            y_data (np.ndarray): Y data
            bandwidth (float): bandwidth parameter

        Returns:
            np.ndarray: Predicted y values.
        """
        x_data = torch.from_numpy(x_data)[:, None]  # Reshape for broadcasting
        weights = torch.exp(
            -0.5 * ((x_data - x_data.T) / bandwidth) ** 2
        )  # Pairwise Gaussian weights
        weights /= weights.sum(dim=1, keepdim=True)  # Normalize weights along each row

        # Weighted sum to get predictions
        y_pred = weights @ torch.from_numpy(y_data)
        return y_pred.numpy()

    def _norm_2(self, x: np.ndarray) -> np.ndarray:
        return np.sqrt(np.square(x).sum(axis=0))

    def functional_norm(self, array_data: np.ndarray) -> np.ndarray:
        """Sample estimate of functional norm for arrays of shape.

            (#samples, #n_group_entries, #tasks).

        Args:
            array_data (np.ndarray): Array input data.

        Returns:
            np.ndarray: Array of functional norm estimates.
        """
        return np.sqrt(np.square(self._norm_2(array_data)).sum(axis=0) / array_data.shape[0])

    def soft_thresholding_update(
        self,
        g: int,
        f_g: np.ndarray | dict[str, np.ndarray],
        zero_groups: list[bool],
        smooth_fit: np.ndarray,
        s_g_ordered: np.ndarray,
        penalty: float,
        m_opt: np.ndarray,
    ) -> np.ndarray | dict[str, np.ndarray]:
        """Soft-thresholding update.

        Args:
            g (int): _description_
            f_g (np.ndarray | dict[str, np.ndarray]): _description_
            zero_groups (list[bool]): _description_
            smooth_fit (np.ndarray): _description_
            s_g_ordered (np.ndarray): _description_
            penalty (float): _description_
            m_opt (np.ndarray): _description_

        Returns:
            np.ndarray | dict[str, np.ndarray]: _description_
        """
        if isinstance(f_g, np.ndarray):
            f_g_hat = f_g[:, :, g, :]
        else:
            f_g_hat = list(f_g.values())[g]

        if zero_groups[g]:
            f_g_hat = np.zeros(f_g_hat.shape)
        else:
            f_g_hat = self.update_loop(
                f_g_hat=f_g_hat,
                d_g=self.d_g_long[g],
                smooth_fit=smooth_fit,
                m_opt=m_opt,
                s_g_ordered=s_g_ordered,
                penalty=penalty,
            )
        if isinstance(f_g, np.ndarray):
            f_g[:, :, g, :] = f_g_hat
        else:
            f_g[list(f_g.keys())[g]] = f_g_hat
        return f_g

    def update_loop(
        self,
        f_g_hat: np.ndarray,
        d_g: int,
        smooth_fit: np.ndarray,
        s_g_ordered: np.ndarray,
        penalty: float,
        m_opt: np.ndarray,
    ) -> np.ndarray:
        """Inner update loop.

        Args:
            f_g_hat (np.ndarray): _description_
            d_g (int): _description_
            smooth_fit (np.ndarray): _description_
            s_g_ordered (np.ndarray): _description_
            penalty (float): _description_
            m_opt (np.ndarray): _description_

        Returns:
            np.ndarray: _description_
        """
        for k in range(self.n_tasks):
            for j in range(d_g):
                # for each task check whether larger or smaller m*
                if (k + 1) > m_opt:
                    f_g_hat[:, j, k] = smooth_fit[:, j, k]
                else:
                    s_g_sum = s_g_ordered[:m_opt].sum()
                    f_g_hat[:, j, k] = (
                        1
                        / m_opt
                        * (s_g_sum - np.sqrt(d_g) * penalty)
                        * smooth_fit[:, j, k]
                        / s_g_ordered[k]
                    )
                f_g_hat[:, j, k] -= f_g_hat[:, j, k].mean()
        return f_g_hat

    def omega_hat(self, f_g: np.ndarray | dict[str, np.ndarray]) -> np.ndarray:
        """Calculate Omega hat.

        Args:
            f_g (np.ndarray): _description_

        Returns:
            np.ndarray: _description_
        """
        if isinstance(f_g, np.ndarray):
            p_g = f_g.shape[2]
            return np.array([self.functional_norm(f_g[:, :, g, :]) for g in range(p_g)])
        else:
            return np.array([self.functional_norm(f_g_hat) for f_g_hat in f_g.values()])

    def _standardize_data(
        self, X_data: np.ndarray | dict[str, np.ndarray]
    ) -> np.ndarray | dict[str, np.ndarray]:
        """Standardize X data.

        Args:
            X_data (np.ndarray | dict[str, np.ndarray]): Predictor array or dict.

        Returns:
            np.ndarray | dict[str, np.ndarray]: Predictor array or dict standardized.
        """
        if isinstance(X_data, dict):
            for key, value in X_data.items():
                X_data[key] = self._assure_two_dim(value)
                X_data[key] = (value - value.mean(axis=0)) / value.std(axis=0)
        else:
            X_data = (X_data - X_data.mean(axis=0)) / X_data.std(axis=0)

        return X_data

    def _assert_group_sizes_match(self, X_data: np.ndarray | dict[str, np.ndarray]) -> bool:
        """Assert whether all groups have same size.

        Args:
            X_data (np.ndarray | dict[str, np.ndarray]): predictor data.

        Returns:
            bool: True if all groups have same size
        """
        if isinstance(X_data, np.ndarray):
            return False
        else:
            return len({dat.shape[1] if len(dat.shape) > 1 else 1 for dat in X_data.values()}) == 1

    def _assure_two_dim(self, arr: np.ndarray) -> np.ndarray:
        if arr.ndim == 1:
            return arr[:, np.newaxis]
        else:
            return arr

    def fit(
        self,
        X_data: np.ndarray | dict[str, np.ndarray],
        Y_data: np.ndarray,
        nlambda: int = 30,
        lambda_min_ratio: float = 0.005,
        precalculate_smooths: bool = True,
        local_regression_method: str = "kernel",
    ) -> None:
        """Fit the multitask group SpAM.

        Args:
            X_data (np.ndarray): _description_
            Y_data (np.ndarray): _description_
            nlambda (int, optional): _description_. Defaults to 30.
            lambda_min_ratio (float, optional): _description_. Defaults to 0.005.
            precalculate_smooths (bool, optional): _description_. Defaults to True.
            local_regression_method (str): Defaults to "loess". Other options currently: "kernel".
        """
        Y_data = self._assure_two_dim(Y_data)
        Y_data -= Y_data.mean(axis=0)
        if self._assert_group_sizes_match(X_data=X_data):
            self.group_names = list(X_data.keys())
            X_data = np.concatenate([data[..., np.newaxis] for data in X_data.values()], axis=2)
        X_data = self._standardize_data(X_data=X_data)

        smoothers = self.precalculate_smooths(
            X_data=X_data, local_regression_method=local_regression_method
        )

        self.select_penalty(
            X_data=X_data,
            Y_data=Y_data,
            nlambda=nlambda,
            lambda_min_ratio=lambda_min_ratio,
            precalculate_smooths=precalculate_smooths,
            smoothers=smoothers,
            local_regression_method=local_regression_method,
        )

    def predict(self) -> np.ndarray:
        """Predict the model.

        Returns:
            np.ndarray: Predicted values of shape (#n_samples, #n_tasks)
        """
        if isinstance(self.f_g_hat, np.ndarray):
            self.predicted_vals = self.f_g_hat.sum(axis=(1, 2))
        else:
            self.predicted_vals = np.asanyarray(
                [f_g.sum(axis=1) for f_g in self.f_g_hat.values()]
            ).sum(axis=0)
        return self.predicted_vals

    def plugin_bandwidth(self, x_j: np.ndarray) -> float:
        """Plugin bandwidth for Gaussian Kernel.

        Args:
            x_j (np.ndarray): data.

        Returns:
            float: Selected bandwidth.
        """
        return float(0.6 * np.std(x_j) * x_j.shape[0] ** (-1 / 5))

    def _make_gaussian_kernel_smoother_matrix(self, data: np.ndarray) -> np.ndarray:
        torch_data = torch.from_numpy(data)[:, None]  # Reshape for broadcasting
        weights = torch.exp(
            -0.5 * ((torch_data - torch_data.T) / self.plugin_bandwidth(data)) ** 2
        )  # Pairwise Gaussian weights
        weights /= weights.sum(dim=1, keepdim=True)  # Normalize weights along each row

        return weights.numpy()

    def _make_loess_smoother_matrix(self, x: np.ndarray) -> np.ndarray:
        """Make smoothing matrix for one instance based on loess.

        Args:
            x (np.ndarray): data.

        Returns:
            np.ndarray: smoothing matrix
        """
        n = x.shape[0]
        smoother = np.zeros((n, n))

        for i in range(n):
            e_i = np.zeros(n)
            e_i[i] = 1
            loess_obj = loess(x, e_i)
            loess_obj.fit()
            smoother[:, i] = loess_obj.outputs.fitted_values

        return smoother

    def gcv(self, Y_data: np.ndarray) -> float:
        """Generalized cross-validation.

        Args:
            Y_data (np.ndarray): Response data of shape (#n_samples. #n_tasks).

        Returns:
            float: Value of GCV.
        """
        n = Y_data.shape[0]
        numerator: float = np.einsum("ij->", self._quadratic_loss(Y_data))
        denominator: float = np.square(np.square(n * self.n_tasks) - n * self.n_tasks * self._df())
        gcv = numerator / denominator
        return gcv

    def _quadratic_loss(self, Y_data: np.ndarray) -> np.ndarray:
        """Current residual fit in terms of LS criterion.

        Args:
            Y_data (np.ndarray): _description_

        Returns:
            np.ndarray: _description_
        """
        if isinstance(self.f_g_hat, np.ndarray):
            Q_g = np.square(Y_data - self.f_g_hat.sum(axis=(1, 2)))
        else:
            Q_g = np.square(
                Y_data
                - np.array([entry.sum(axis=1) for entry in self.f_g_hat.values()]).sum(axis=0)
            )

        return Q_g

    def _df(self) -> float:
        """Effective degrees of freedom in terms of trace of smoother matrix.

        K multiplication is due to the single-task nature of the predictors.
        In case of multitask this will need to be adapted in the future.

        Returns:
            float: K * sum_g^p sum_j^d_g trace(S_j^k) I(||f_j^k|| not 0)
        """
        v_jk = {}
        for g, d_g in enumerate(self.d_g_long):
            v_jk[g] = np.zeros(d_g)
            for j in range(d_g):
                if self.zero_groups[g]:
                    continue
                elif isinstance(self.smoothing_matrices, np.ndarray):
                    v_jk[g][j] = np.einsum(
                        "ii",
                        self.smoothing_matrices[:, :, j, g],
                    )
                else:
                    v_jk[g][j] = np.einsum(
                        "ii",
                        self.smoothing_matrices[list(self.smoothing_matrices.keys())[g]][:, :, j],
                    )

        v_jk_sum: float = np.asanyarray([df.sum() for df in v_jk.values()]).sum()
        v_jk_sum_sum = self.n_tasks * v_jk_sum
        return v_jk_sum_sum

    def select_penalty(
        self,
        X_data: np.ndarray | dict[str, np.ndarray],
        Y_data: np.ndarray,
        nlambda: int = 30,
        lambda_min_ratio: float = 5e-3,
        smoothers: np.ndarray | dict[str, np.ndarray] | None = None,
        precalculate_smooths: bool = True,
        local_regression_method: str = "kernel",
    ) -> None:
        """GCV model selection procedure.

        Args:
            X_data (np.ndarray): _description_
            Y_data (np.ndarray): _description_
            nlambda (int, optional): _description_. Defaults to 30.
            lambda_min_ratio (float, optional): _description_. Defaults to 5e-3.
            precalculate_smooths (bool, optional): _description_. Defaults to True.
            smoothers (np.ndarray | dict[str, np.ndarray] | None): _description_. Defaults to None.
            local_regression_method (str): Defaults to "loess". Other options currently: "kernel".
        """
        # This is only ever be necessary if the function is called outside fit()
        if smoothers is None:
            smoothers = self.precalculate_smooths(
                X_data=X_data, local_regression_method=local_regression_method
            )

        self._find_lambda_max_value(
            X_data=X_data,
            Y_data=Y_data,
            smoothers=smoothers,
            precalculate_smooths=precalculate_smooths,
            local_regression_method=local_regression_method,
        )
        lambda_scale = np.exp(
            np.linspace(start=np.log(1), stop=np.log(lambda_min_ratio), num=nlambda)
        )
        self.lambda_values = lambda_scale * self.lambda_max_value

        gcv_values = np.empty(len(self.lambda_values))
        f_hat_from_previous_lambda = self._init_functions(X_data=X_data)
        current_min = np.inf
        for i, lambda_value in enumerate(self.lambda_values):
            self._progressBar(i, len(self.lambda_values) - 1, suffix="Finding optimal lambda")
            try:
                self.block_coordinate_descent(
                    X_data=X_data,
                    Y_data=Y_data,
                    penalty=lambda_value,
                    precalculate_smooths=precalculate_smooths,
                    smoothers=smoothers,
                    warm_start_f_hat=f_hat_from_previous_lambda,
                    local_regression_method=local_regression_method,
                )
                gcv_values[i] = self.gcv(Y_data=Y_data)
                f_hat_from_previous_lambda = self.f_g_hat
                self.zero_group_history[lambda_value] = self.zero_groups

                if gcv_values[i] < current_min:
                    final_f_hat = self.f_g_hat  # Save for later
                    final_zero_groups = self.zero_groups  # Save for later
                    final_steps_till_convergence = self.steps_till_convergence
                    current_min = gcv_values[i]

            except ConvergenceError:  # Stop if no convergence
                if i == 0:  # If this is negative even max_lambda didn't converge.
                    raise ValueError(
                        "No lambda value converged. Something wrong with causal order?"
                    )
                break

        self.chosen_lambda = self.lambda_values[np.argmin(gcv_values)]
        self.f_g_hat = final_f_hat
        self.zero_groups = final_zero_groups
        self.steps_till_convergence = final_steps_till_convergence
        self.gcv_values = gcv_values

    def _progressBar(self, count_value: float, total: float, suffix: str = "") -> None:
        bar_length = 100
        filled_up_Length = int(round(bar_length * count_value / float(total)))
        percentage = round(100.0 * count_value / float(total), 1)
        bar = "=" * filled_up_Length + "-" * (bar_length - filled_up_Length)
        sys.stdout.write(f"[{bar}] {percentage}% ...{suffix}\r")
        sys.stdout.flush()

    def _find_lambda_max_value(
        self,
        X_data: np.ndarray | dict[str, np.ndarray],
        Y_data: np.ndarray,
        smoothers: np.ndarray | dict[str, np.ndarray] | None = None,
        precalculate_smooths: bool = True,
        local_regression_method: str = "kernel",
    ) -> None:
        """Identifies largest smallest penalty that just renders all groups \

            to be shrunken to zero.

        Args:
            X_data (np.ndarray | dict[str, np.ndarray]): Predictors.
            Y_data (np.ndarray): Targets.
            smoothers (np.ndarray | dict[str, np.ndarray] | None, optional): Pre-calculated
                Smoothers. Defaults to None.
            precalculate_smooths (bool): Whether to precalculate smooths. Is just for completeness.
                Typically the smooths will have been calculated and provided externally when calling
                this function.
            local_regression_method (str): Defaults to "loess". Other options currently: "kernel".
        """
        self.n_tasks = Y_data.shape[1]

        if isinstance(X_data, dict):
            X_data = self._dict_preprocessing(X_data=X_data)

        # Number of predictors and members in groups
        self._set_dg_pg(X_data=X_data)

        f_g = self._init_functions(X_data=X_data)

        if smoothers is None:
            smoothers = self.precalculate_smooths(
                X_data=X_data, local_regression_method=local_regression_method
            )

        omega_g_sum = np.zeros(self.p_g)
        for g in range(self.p_g):
            d_g = self.d_g_long[g]
            # R_g_hat has shape (#samples, #tasks)
            R_g_hat = self.R_g_hat_update(f_g=f_g, Y_data=Y_data, g=g)

            if precalculate_smooths:
                smooth_fit = self.predict_from_linear_smoother(
                    g=g, smoothing_matrices=smoothers, R_g=R_g_hat
                )
            elif local_regression_method == "loess":
                smooth_fit = self.loess_direct_fit(g=g, d_g=d_g, R_g=R_g_hat, X_data=X_data)
            elif local_regression_method == "kernel":
                smooth_fit = self.gaussian_kernel_direct_fit(
                    g=g, d_g=d_g, R_g=R_g_hat, X_data=X_data
                )
            else:
                raise NotImplementedError("Smoothing method not implemented.")

            # estimate of || Q^(k)R_g^(k) ||
            # shape (#predictors, #tasks)
            omega_g = self.functional_norm(smooth_fit)
            # lambda_max_value must be at least omega_g.sum()/np.sqrt(d_g)
            omega_g_sum[g] = omega_g.sum()

        self.lambda_max_value = np.ceil(np.max(omega_g_sum / np.sqrt(self.d_g_long)))

    def plot_gcv_path(self) -> None:
        """Plot GCV path."""
        if self.gcv_values is None:
            raise ValueError("No GCV values available. Run model selection first.")
        _, ax = plt.subplots()
        ax.plot(self.lambda_values, self.gcv_values)
        plt.axvline(
            x=self.chosen_lambda,
            color="red",
            linestyle="--",
            alpha=0.35,
        )
        plt.xlabel("lambda")
        plt.ylabel("gcv")
        plt.show()

R_g_hat_update(f_g, Y_data, g)

Update partial residuals.

Parameters:

Name	Type	Description	Default
f_g	ndarray	Current additive components of shape (#samples, #groups, #predictors, #tasks)	required
Y_data	ndarray	Response data of shape (#samples, #tasks)	required
g	int	group in question.	required

Returns:

Type	Description
ndarray	np.ndarray: Partial Residuals.

Source code in gresit/model_selection.py

def R_g_hat_update(
    self, f_g: np.ndarray | dict[str, np.ndarray], Y_data: np.ndarray, g: int
) -> np.ndarray:
    """Update partial residuals.

    Args:
        f_g (np.ndarray): Current additive components
            of shape (#samples, #groups, #predictors, #tasks)
        Y_data (np.ndarray): Response data of shape (#samples, #tasks)
        g (int): group in question.

    Returns:
        np.ndarray: Partial Residuals.
    """
    if isinstance(f_g, np.ndarray):
        mask = np.ones(self.p_g, bool)
        mask[g] = False
        R_g = Y_data - f_g[:, :, mask, :].sum(axis=(1, 2))
    else:
        groups = list(f_g.keys())
        f_interim_sum = np.asanyarray(
            [value.sum(axis=1) for key, value in f_g.items() if key != groups[g]]
        )
        R_g = Y_data - f_interim_sum.sum(axis=0)

    return R_g

__init__(group_names=None)

Initializes the object.

Source code in gresit/model_selection.py

def __init__(self, group_names: list[str] | None = None) -> None:
    """Initializes the object."""
    self.zero_groups: list[bool]
    self.zero_group_history: dict[float, list[bool]] = {}
    self.d_g_long: np.ndarray
    self.p_g: int = 0
    self.n_tasks: int
    self.smoothing_matrices: np.ndarray | dict[str, np.ndarray]
    self.f_g_hat: np.ndarray | dict[str, np.ndarray]
    self.chosen_lambda: float = 0.0
    self.lambda_max_value: float | None = None
    self.steps_till_convergence: int = -1
    self.max_iter: int = 10000
    self.tol: float = 1e-8
    if group_names is not None:
        self.group_names = group_names
    else:
        self.group_names = []
    self.predicted_vals: np.ndarray

__repr__()

Print some useful information.

Returns:

Name	Type	Description
str	str	description

Source code in gresit/model_selection.py

def __repr__(self) -> str:
    """Print some useful information.

    Returns:
        str: _description_
    """
    return f"MT-GSpAM on {self.p_g} groups."

__str__()

Print some useful summary.

Returns:

Name	Type	Description
str	str	description

Source code in gresit/model_selection.py

def __str__(self) -> str:
    """Print some useful summary.

    Returns:
        str: _description_
    """
    display_limit = 5
    setting = {
        "Number of Groups: ": self.p_g,
        "Group sizes homogenous: ": True if isinstance(self.f_g_hat, np.ndarray) else False,
        "Steps until convergence: ": self.steps_till_convergence,
        "Final lambda: ": self.chosen_lambda,
        "Number of non-zero groups: ": self.p_g - np.asanyarray(self.zero_groups).sum(),
        "Non-zero group names: ": self.return_nonzero_groups()
        if len(self.return_nonzero_groups()) <= display_limit
        else "Too many to show",
    }
    s = ""
    for info, info_text in setting.items():
        s += f"{info}{info_text}\n"

    return s

block_coordinate_descent(X_data, Y_data, penalty, precalculate_smooths=True, smoothers=None, local_regression_method='kernel', warm_start_f_hat=None)

Block coordinate descent for multitask Sparse Group Lasso.

Parameters:

Name	Type	Description	Default
X_data	ndarray | dict[str, ndarray]	description	required
Y_data	ndarray	description	required
penalty	float	description	required
precalculate_smooths	bool	description. Defaults to True.	True
smoothers	ndarray | dict[str, ndarray] | None	description. Defaults to None.	None
local_regression_method	str	Defaults to "loess". Other options currently: "kernel".	'kernel'
warm_start_f_hat	ndarray | dict[str, ndarray] | None	description. Defaults to None.	None

Source code in gresit/model_selection.py

def block_coordinate_descent(
    self,
    X_data: np.ndarray | dict[str, np.ndarray],
    Y_data: np.ndarray,
    penalty: float,
    precalculate_smooths: bool = True,
    smoothers: np.ndarray | dict[str, np.ndarray] | None = None,
    local_regression_method: str = "kernel",
    warm_start_f_hat: np.ndarray | dict[str, np.ndarray] | None = None,
) -> None:
    """Block coordinate descent for multitask Sparse Group Lasso.

    Args:
        X_data (np.ndarray | dict[str, np.ndarray]): _description_
        Y_data (np.ndarray): _description_
        penalty (float): _description_
        precalculate_smooths (bool, optional): _description_. Defaults to True.
        smoothers (np.ndarray | dict[str, np.ndarray] | None, optional): _description_.
            Defaults to None.
        local_regression_method (str): Defaults to "loess". Other options currently: "kernel".
        warm_start_f_hat (np.ndarray | dict[str, np.ndarray] | None, optional): _description_.
            Defaults to None.
    """
    if smoothers is None:
        smoothers = self.precalculate_smooths(
            X_data=X_data, local_regression_method=local_regression_method
        )

    self.n_tasks = Y_data.shape[1]

    if isinstance(X_data, dict):
        if not self.group_names:
            self.group_names = list(X_data.keys())
        X_data = self._dict_preprocessing(X_data=X_data)

    # Number of predictors and members in groups
    self._set_dg_pg(X_data=X_data)

    f_g = self._init_or_insert_functions(X_data, warm_start_f_hat)

    divergent_max_runs = 10
    divergent_runs = 0
    max_inc_old = 0.0
    new_functional_norm = np.zeros((self.p_g, self.n_tasks))
    # Update loop
    for t in range(self.max_iter):
        old_functional_norm = new_functional_norm
        zero_groups = [False] * self.p_g
        for g in range(self.p_g):
            d_g = self.d_g_long[g]
            # R_g_hat has shape (#samples, #tasks)
            R_g_hat = self.R_g_hat_update(f_g=f_g, Y_data=Y_data, g=g)

            if precalculate_smooths:
                smooth_fit = self.predict_from_linear_smoother(
                    g=g, smoothing_matrices=smoothers, R_g=R_g_hat
                )
            else:
                smooth_fit = self.smoother_direct_fit(
                    g=g,
                    d_g=d_g,
                    X_data=X_data,
                    R_g=R_g_hat,
                    local_regression_method=local_regression_method,
                )

            # # estimate of || Q^(k)R_g^(k) ||
            # # shape (1, #tasks)
            omega_g = self.functional_norm(smooth_fit)

            # sort in descending order for each task
            s_g_ordered = np.sort(omega_g)[::-1]

            # get m* for the case that more than one group attains the sup-norm.
            m_opt = self._get_m_star(s_g_ordered=s_g_ordered, penalty=penalty, d_g=d_g)

            # calculate nonzero groups
            zero_groups[g] = omega_g.sum() <= penalty * np.sqrt(d_g)
            f_g = self.soft_thresholding_update(
                g=g,
                f_g=f_g,
                zero_groups=zero_groups,
                smooth_fit=smooth_fit,
                s_g_ordered=s_g_ordered,
                penalty=penalty,
                m_opt=m_opt,
            )

        new_functional_norm = self.omega_hat(f_g)
        max_inc: float = np.sqrt(
            np.square(new_functional_norm - old_functional_norm).sum(axis=(0, 1))
        )
        if max_inc_old < max_inc:
            divergent_runs += 1

        if divergent_runs > divergent_max_runs:
            raise ConvergenceError(
                f"Model did not converge in group {g}. Penalty: {penalty} too low?"
            )
        max_inc_old = max_inc

        if np.all(max_inc < self.tol / 2):
            break

    self.f_g_hat = f_g
    self.zero_groups = zero_groups
    self.steps_till_convergence = t

fit(X_data, Y_data, nlambda=30, lambda_min_ratio=0.005, precalculate_smooths=True, local_regression_method='kernel')

Fit the multitask group SpAM.

Parameters:

Name	Type	Description	Default
X_data	ndarray	description	required
Y_data	ndarray	description	required
nlambda	int	description. Defaults to 30.	30
lambda_min_ratio	float	description. Defaults to 0.005.	0.005
precalculate_smooths	bool	description. Defaults to True.	True
local_regression_method	str	Defaults to "loess". Other options currently: "kernel".	'kernel'

Source code in gresit/model_selection.py

def fit(
    self,
    X_data: np.ndarray | dict[str, np.ndarray],
    Y_data: np.ndarray,
    nlambda: int = 30,
    lambda_min_ratio: float = 0.005,
    precalculate_smooths: bool = True,
    local_regression_method: str = "kernel",
) -> None:
    """Fit the multitask group SpAM.

    Args:
        X_data (np.ndarray): _description_
        Y_data (np.ndarray): _description_
        nlambda (int, optional): _description_. Defaults to 30.
        lambda_min_ratio (float, optional): _description_. Defaults to 0.005.
        precalculate_smooths (bool, optional): _description_. Defaults to True.
        local_regression_method (str): Defaults to "loess". Other options currently: "kernel".
    """
    Y_data = self._assure_two_dim(Y_data)
    Y_data -= Y_data.mean(axis=0)
    if self._assert_group_sizes_match(X_data=X_data):
        self.group_names = list(X_data.keys())
        X_data = np.concatenate([data[..., np.newaxis] for data in X_data.values()], axis=2)
    X_data = self._standardize_data(X_data=X_data)

    smoothers = self.precalculate_smooths(
        X_data=X_data, local_regression_method=local_regression_method
    )

    self.select_penalty(
        X_data=X_data,
        Y_data=Y_data,
        nlambda=nlambda,
        lambda_min_ratio=lambda_min_ratio,
        precalculate_smooths=precalculate_smooths,
        smoothers=smoothers,
        local_regression_method=local_regression_method,
    )

functional_norm(array_data)

Sample estimate of functional norm for arrays of shape.

(#samples, #n_group_entries, #tasks).

Parameters:

Name	Type	Description	Default
array_data	ndarray	Array input data.	required

Returns:

Type	Description
ndarray	np.ndarray: Array of functional norm estimates.

Source code in gresit/model_selection.py

def functional_norm(self, array_data: np.ndarray) -> np.ndarray:
    """Sample estimate of functional norm for arrays of shape.

        (#samples, #n_group_entries, #tasks).

    Args:
        array_data (np.ndarray): Array input data.

    Returns:
        np.ndarray: Array of functional norm estimates.
    """
    return np.sqrt(np.square(self._norm_2(array_data)).sum(axis=0) / array_data.shape[0])

gaussian_kernel_direct_fit(g, d_g, R_g, X_data)

Local regression fit.

Parameters:

Name	Type	Description	Default
g	int	group in question.	required
d_g	int	Number of predictors in group.	required
R_g	ndarray	Partial residuals should be of shape (#n_samples. #n_predictors, #n_tasks)	required
X_data	ndarray	Training data of shape (#n_samples. #n_groups, #n_predictors)	required

Returns:

Type	Description
ndarray	np.ndarray: Fitted local regressions.

Source code in gresit/model_selection.py

def gaussian_kernel_direct_fit(
    self, g: int, d_g: int, R_g: np.ndarray, X_data: np.ndarray | dict[str, np.ndarray]
) -> np.ndarray:
    """Local regression fit.

    Args:
        g (int): group in question.
        d_g (int): Number of predictors in group.
        R_g (np.ndarray): Partial residuals should be of
            shape (#n_samples. #n_predictors, #n_tasks)
        X_data (np.ndarray): Training data of shape (#n_samples. #n_groups, #n_predictors)

    Returns:
        np.ndarray: Fitted local regressions.
    """
    if isinstance(X_data, np.ndarray):
        data = X_data[:, :, g]
    else:
        data = list(X_data.values())[g]

    kernel_reg_g = np.zeros((R_g.shape[0], d_g, self.n_tasks))
    for k in range(self.n_tasks):
        for j in range(d_g):
            kernel_reg_g[:, j, k] = self._torch_gaussian_kernel_regression(
                x_data=data[:, j],
                y_data=R_g[:, k],
                bandwidth=self.plugin_bandwidth(x_j=data[:, j]),
            )

    return kernel_reg_g

gcv(Y_data)

Generalized cross-validation.

Parameters:

Name	Type	Description	Default
Y_data	ndarray	Response data of shape (#n_samples. #n_tasks).	required

Returns:

Name	Type	Description
float	float	Value of GCV.

Source code in gresit/model_selection.py

def gcv(self, Y_data: np.ndarray) -> float:
    """Generalized cross-validation.

    Args:
        Y_data (np.ndarray): Response data of shape (#n_samples. #n_tasks).

    Returns:
        float: Value of GCV.
    """
    n = Y_data.shape[0]
    numerator: float = np.einsum("ij->", self._quadratic_loss(Y_data))
    denominator: float = np.square(np.square(n * self.n_tasks) - n * self.n_tasks * self._df())
    gcv = numerator / denominator
    return gcv

loess_direct_fit(g, d_g, R_g, X_data)

Local regression fit.

Parameters:

Name	Type	Description	Default
g	int	group in question.	required
d_g	int	Number of predictors in group.	required
R_g	ndarray	Partial residuals should be of shape (#n_samples. #n_predictors, #n_tasks)	required
X_data	ndarray	Training data of shape (#n_samples. #n_groups, #n_predictors)	required

Returns:

Type	Description
ndarray	np.ndarray: Fitted local regressions.

Source code in gresit/model_selection.py

def loess_direct_fit(
    self, g: int, d_g: int, R_g: np.ndarray, X_data: np.ndarray | dict[str, np.ndarray]
) -> np.ndarray:
    """Local regression fit.

    Args:
        g (int): group in question.
        d_g (int): Number of predictors in group.
        R_g (np.ndarray): Partial residuals should be of
            shape (#n_samples. #n_predictors, #n_tasks)
        X_data (np.ndarray): Training data of shape (#n_samples. #n_groups, #n_predictors)

    Returns:
        np.ndarray: Fitted local regressions.
    """
    if isinstance(X_data, np.ndarray):
        data = X_data[:, :, g]
    else:
        data = list(X_data.values())[g]

    loess_g = np.zeros((R_g.shape[0], d_g, self.n_tasks))
    for k in range(self.n_tasks):
        for j in range(d_g):
            try:
                loess_obj = loess(data[:, j], R_g[:, k])
            except ValueError:
                loess_obj = loess(data[:, j], R_g[:, k], surface="direct")
            loess_obj.fit()
            loess_g[:, j, k] = loess_obj.outputs.fitted_values
    return loess_g

omega_hat(f_g)

Calculate Omega hat.

Parameters:

Name	Type	Description	Default
f_g	ndarray	description	required

Returns:

Type	Description
ndarray	np.ndarray: description

Source code in gresit/model_selection.py

def omega_hat(self, f_g: np.ndarray | dict[str, np.ndarray]) -> np.ndarray:
    """Calculate Omega hat.

    Args:
        f_g (np.ndarray): _description_

    Returns:
        np.ndarray: _description_
    """
    if isinstance(f_g, np.ndarray):
        p_g = f_g.shape[2]
        return np.array([self.functional_norm(f_g[:, :, g, :]) for g in range(p_g)])
    else:
        return np.array([self.functional_norm(f_g_hat) for f_g_hat in f_g.values()])

plot_gcv_path()

Plot GCV path.

Source code in gresit/model_selection.py

def plot_gcv_path(self) -> None:
    """Plot GCV path."""
    if self.gcv_values is None:
        raise ValueError("No GCV values available. Run model selection first.")
    _, ax = plt.subplots()
    ax.plot(self.lambda_values, self.gcv_values)
    plt.axvline(
        x=self.chosen_lambda,
        color="red",
        linestyle="--",
        alpha=0.35,
    )
    plt.xlabel("lambda")
    plt.ylabel("gcv")
    plt.show()

plugin_bandwidth(x_j)

Plugin bandwidth for Gaussian Kernel.

Parameters:

Name	Type	Description	Default
x_j	ndarray	data.	required

Returns:

Name	Type	Description
float	float	Selected bandwidth.

Source code in gresit/model_selection.py

def plugin_bandwidth(self, x_j: np.ndarray) -> float:
    """Plugin bandwidth for Gaussian Kernel.

    Args:
        x_j (np.ndarray): data.

    Returns:
        float: Selected bandwidth.
    """
    return float(0.6 * np.std(x_j) * x_j.shape[0] ** (-1 / 5))

precalculate_smooths(X_data, local_regression_method='kernel')

Precalculate smoother matrices.

Input may be both np.ndarray and dict.

Parameters:

Name	Type	Description	Default
X_data	ndarray | dict[str, ndarray]	predictor data.	required
local_regression_method	str	Method to use to calculate smoother matrix. Options currently are "loess" and "kernel". When kernel is chosen, the default is set to Gaussian kernel regression with the standard deviation rule for bandwidth selection.	'kernel'

Returns:

Type	Description
ndarray | dict[str, ndarray]	np.ndarray | dict[str, np.ndarray]: smooths.

Source code in gresit/model_selection.py

def precalculate_smooths(
    self, X_data: np.ndarray | dict[str, np.ndarray], local_regression_method: str = "kernel"
) -> np.ndarray | dict[str, np.ndarray]:
    """Precalculate smoother matrices.

    Input may be both `np.ndarray` and `dict`.

    Args:
        X_data (np.ndarray | dict[str, np.ndarray]): predictor data.
        local_regression_method (str): Method to use to calculate smoother matrix. Options
            currently are `"loess"` and `"kernel"`. When kernel is chosen, the default is set to
            Gaussian kernel regression with the standard deviation rule for bandwidth selection.

    Returns:
        np.ndarray | dict[str, np.ndarray]: smooths.
    """
    if isinstance(X_data, dict):
        X_data = self._dict_preprocessing(X_data=X_data)
        smoothing_matrices = {}
        for group, data in X_data.items():
            smoothing_matrices[group] = np.zeros((data.shape[0], data.shape[0], data.shape[1]))
            for j in range(data.shape[1]):
                if local_regression_method == "loess":
                    smoothing_matrices[group][:, :, j] = self._make_loess_smoother_matrix(
                        data[:, j]
                    )
                elif local_regression_method == "kernel":
                    smoothing_matrices[group][:, :, j] = (
                        self._make_gaussian_kernel_smoother_matrix(data[:, j])
                    )
                else:
                    raise NotImplementedError("This smoothing method is not implemented yet.")
    else:
        smoothing_matrices = np.zeros(
            (X_data.shape[0], X_data.shape[0], X_data.shape[1], X_data.shape[2]),
            dtype=np.float64,
        )
        inner_smoother: np.ndarray = smoothing_matrices
        for num_groups in range(X_data.shape[2]):
            for group_members in range(X_data.shape[1]):
                if local_regression_method == "loess":
                    inner_smoother[:, :, group_members, num_groups] = (
                        self._make_loess_smoother_matrix(X_data[:, group_members, num_groups])
                    )
                elif local_regression_method == "kernel":
                    inner_smoother[:, :, group_members, num_groups] = (
                        self._make_gaussian_kernel_smoother_matrix(
                            X_data[:, group_members, num_groups]
                        )
                    )
                else:
                    raise NotImplementedError("This smoothing method is not implemented yet.")
        smoothing_matrices = inner_smoother

    self.smoothing_matrices = smoothing_matrices
    return smoothing_matrices

predict()

Predict the model.

Returns:

Type	Description
ndarray	np.ndarray: Predicted values of shape (#n_samples, #n_tasks)

Source code in gresit/model_selection.py

def predict(self) -> np.ndarray:
    """Predict the model.

    Returns:
        np.ndarray: Predicted values of shape (#n_samples, #n_tasks)
    """
    if isinstance(self.f_g_hat, np.ndarray):
        self.predicted_vals = self.f_g_hat.sum(axis=(1, 2))
    else:
        self.predicted_vals = np.asanyarray(
            [f_g.sum(axis=1) for f_g in self.f_g_hat.values()]
        ).sum(axis=0)
    return self.predicted_vals

predict_from_linear_smoother(g, smoothing_matrices, R_g)

Local regression fit.

Parameters:

Name	Type	Description	Default
g	int	group in question.	required
R_g	ndarray	Partial residuals should be of shape (#n_samples. #n_predictors, #n_tasks)	required
smoothing_matrices	ndarray	Precalculated smoothing matrices of shape (#n_samples, #n_samples #n_groups, #n_predictors)	required

Returns:

Type	Description
ndarray	torch.Tensor: Fitted local regressions.

Source code in gresit/model_selection.py

def predict_from_linear_smoother(
    self,
    g: int,
    smoothing_matrices: np.ndarray | dict[str, np.ndarray],
    R_g: np.ndarray,
) -> np.ndarray:
    """Local regression fit.

    Args:
        g (int): group in question.
        R_g (np.ndarray): Partial residuals should be of
            shape (#n_samples. #n_predictors, #n_tasks)
        smoothing_matrices (np.ndarray): Precalculated smoothing matrices
            of shape (#n_samples, #n_samples #n_groups, #n_predictors)

    Returns:
        torch.Tensor: Fitted local regressions.
    """
    if isinstance(smoothing_matrices, np.ndarray):
        smoothing_matrix = smoothing_matrices[:, :, :, g]
    else:
        smoothing_matrix = list(smoothing_matrices.values())[g]

    return torch.einsum(
        "ijk, jl -> ikl", torch.from_numpy(smoothing_matrix), torch.from_numpy(R_g)
    ).numpy()

return_nonzero_groups()

Return the group names of the nonzero groups.

Returns:

Type	Description
list[str]	list[str] | str: List of nonzero groups with their actual names if given.

Source code in gresit/model_selection.py

def return_nonzero_groups(self) -> list[str]:
    """Return the group names of the nonzero groups.

    Returns:
        list[str] | str: List of nonzero groups with their actual names if given.
    """
    if len(self.group_names) > 0:
        nonzero_groups = [
            group
            for group, zero_group in zip(self.group_names, self.zero_groups)
            if not zero_group
        ]
    else:
        nonzero_groups = ["You neither provided any group names nor ran the model."]

    return nonzero_groups

select_penalty(X_data, Y_data, nlambda=30, lambda_min_ratio=0.005, smoothers=None, precalculate_smooths=True, local_regression_method='kernel')

GCV model selection procedure.

Parameters:

Name	Type	Description	Default
X_data	ndarray	description	required
Y_data	ndarray	description	required
nlambda	int	description. Defaults to 30.	30
lambda_min_ratio	float	description. Defaults to 5e-3.	0.005
precalculate_smooths	bool	description. Defaults to True.	True
smoothers	ndarray | dict[str, ndarray] | None	description. Defaults to None.	None
local_regression_method	str	Defaults to "loess". Other options currently: "kernel".	'kernel'

Source code in gresit/model_selection.py

def select_penalty(
    self,
    X_data: np.ndarray | dict[str, np.ndarray],
    Y_data: np.ndarray,
    nlambda: int = 30,
    lambda_min_ratio: float = 5e-3,
    smoothers: np.ndarray | dict[str, np.ndarray] | None = None,
    precalculate_smooths: bool = True,
    local_regression_method: str = "kernel",
) -> None:
    """GCV model selection procedure.

    Args:
        X_data (np.ndarray): _description_
        Y_data (np.ndarray): _description_
        nlambda (int, optional): _description_. Defaults to 30.
        lambda_min_ratio (float, optional): _description_. Defaults to 5e-3.
        precalculate_smooths (bool, optional): _description_. Defaults to True.
        smoothers (np.ndarray | dict[str, np.ndarray] | None): _description_. Defaults to None.
        local_regression_method (str): Defaults to "loess". Other options currently: "kernel".
    """
    # This is only ever be necessary if the function is called outside fit()
    if smoothers is None:
        smoothers = self.precalculate_smooths(
            X_data=X_data, local_regression_method=local_regression_method
        )

    self._find_lambda_max_value(
        X_data=X_data,
        Y_data=Y_data,
        smoothers=smoothers,
        precalculate_smooths=precalculate_smooths,
        local_regression_method=local_regression_method,
    )
    lambda_scale = np.exp(
        np.linspace(start=np.log(1), stop=np.log(lambda_min_ratio), num=nlambda)
    )
    self.lambda_values = lambda_scale * self.lambda_max_value

    gcv_values = np.empty(len(self.lambda_values))
    f_hat_from_previous_lambda = self._init_functions(X_data=X_data)
    current_min = np.inf
    for i, lambda_value in enumerate(self.lambda_values):
        self._progressBar(i, len(self.lambda_values) - 1, suffix="Finding optimal lambda")
        try:
            self.block_coordinate_descent(
                X_data=X_data,
                Y_data=Y_data,
                penalty=lambda_value,
                precalculate_smooths=precalculate_smooths,
                smoothers=smoothers,
                warm_start_f_hat=f_hat_from_previous_lambda,
                local_regression_method=local_regression_method,
            )
            gcv_values[i] = self.gcv(Y_data=Y_data)
            f_hat_from_previous_lambda = self.f_g_hat
            self.zero_group_history[lambda_value] = self.zero_groups

            if gcv_values[i] < current_min:
                final_f_hat = self.f_g_hat  # Save for later
                final_zero_groups = self.zero_groups  # Save for later
                final_steps_till_convergence = self.steps_till_convergence
                current_min = gcv_values[i]

        except ConvergenceError:  # Stop if no convergence
            if i == 0:  # If this is negative even max_lambda didn't converge.
                raise ValueError(
                    "No lambda value converged. Something wrong with causal order?"
                )
            break

    self.chosen_lambda = self.lambda_values[np.argmin(gcv_values)]
    self.f_g_hat = final_f_hat
    self.zero_groups = final_zero_groups
    self.steps_till_convergence = final_steps_till_convergence
    self.gcv_values = gcv_values

smoother_direct_fit(g, d_g, R_g, X_data, local_regression_method)

Direct fit of smoothing.

Parameters:

Name	Type	Description	Default
g	int	description	required
d_g	int	description	required
R_g	ndarray	description	required
X_data	ndarray	description	required
local_regression_method	str	description	required

Raises:

Type	Description
NotImplementedError	description

Returns:

Name	Type	Description
_type_	ndarray	description

Source code in gresit/model_selection.py

def smoother_direct_fit(
    self,
    g: int,
    d_g: int,
    R_g: np.ndarray,
    X_data: np.ndarray | dict[str, np.ndarray],
    local_regression_method: str,
) -> np.ndarray:
    """Direct fit of smoothing.

    Args:
        g (int): _description_
        d_g (int): _description_
        R_g (np.ndarray): _description_
        X_data (np.ndarray): _description_
        local_regression_method (str): _description_

    Raises:
        NotImplementedError: _description_

    Returns:
        _type_: _description_
    """
    if local_regression_method == "loess":
        smooth_fit = self.loess_direct_fit(g=g, d_g=d_g, R_g=R_g, X_data=X_data)
    elif local_regression_method == "kernel":
        smooth_fit = self.gaussian_kernel_direct_fit(g=g, d_g=d_g, R_g=R_g, X_data=X_data)
    else:
        raise NotImplementedError("Smoothing method not implemented.")
    return smooth_fit

soft_thresholding_update(g, f_g, zero_groups, smooth_fit, s_g_ordered, penalty, m_opt)

Soft-thresholding update.

Parameters:

Name	Type	Description	Default
g	int	description	required
f_g	ndarray | dict[str, ndarray]	description	required
zero_groups	list[bool]	description	required
smooth_fit	ndarray	description	required
s_g_ordered	ndarray	description	required
penalty	float	description	required
m_opt	ndarray	description	required

Returns:

Type	Description
ndarray | dict[str, ndarray]	np.ndarray | dict[str, np.ndarray]: description

Source code in gresit/model_selection.py

def soft_thresholding_update(
    self,
    g: int,
    f_g: np.ndarray | dict[str, np.ndarray],
    zero_groups: list[bool],
    smooth_fit: np.ndarray,
    s_g_ordered: np.ndarray,
    penalty: float,
    m_opt: np.ndarray,
) -> np.ndarray | dict[str, np.ndarray]:
    """Soft-thresholding update.

    Args:
        g (int): _description_
        f_g (np.ndarray | dict[str, np.ndarray]): _description_
        zero_groups (list[bool]): _description_
        smooth_fit (np.ndarray): _description_
        s_g_ordered (np.ndarray): _description_
        penalty (float): _description_
        m_opt (np.ndarray): _description_

    Returns:
        np.ndarray | dict[str, np.ndarray]: _description_
    """
    if isinstance(f_g, np.ndarray):
        f_g_hat = f_g[:, :, g, :]
    else:
        f_g_hat = list(f_g.values())[g]

    if zero_groups[g]:
        f_g_hat = np.zeros(f_g_hat.shape)
    else:
        f_g_hat = self.update_loop(
            f_g_hat=f_g_hat,
            d_g=self.d_g_long[g],
            smooth_fit=smooth_fit,
            m_opt=m_opt,
            s_g_ordered=s_g_ordered,
            penalty=penalty,
        )
    if isinstance(f_g, np.ndarray):
        f_g[:, :, g, :] = f_g_hat
    else:
        f_g[list(f_g.keys())[g]] = f_g_hat
    return f_g

update_loop(f_g_hat, d_g, smooth_fit, s_g_ordered, penalty, m_opt)

Inner update loop.

Parameters:

Name	Type	Description	Default
f_g_hat	ndarray	description	required
d_g	int	description	required
smooth_fit	ndarray	description	required
s_g_ordered	ndarray	description	required
penalty	float	description	required
m_opt	ndarray	description	required

Returns:

Type	Description
ndarray	np.ndarray: description

Source code in gresit/model_selection.py

def update_loop(
    self,
    f_g_hat: np.ndarray,
    d_g: int,
    smooth_fit: np.ndarray,
    s_g_ordered: np.ndarray,
    penalty: float,
    m_opt: np.ndarray,
) -> np.ndarray:
    """Inner update loop.

    Args:
        f_g_hat (np.ndarray): _description_
        d_g (int): _description_
        smooth_fit (np.ndarray): _description_
        s_g_ordered (np.ndarray): _description_
        penalty (float): _description_
        m_opt (np.ndarray): _description_

    Returns:
        np.ndarray: _description_
    """
    for k in range(self.n_tasks):
        for j in range(d_g):
            # for each task check whether larger or smaller m*
            if (k + 1) > m_opt:
                f_g_hat[:, j, k] = smooth_fit[:, j, k]
            else:
                s_g_sum = s_g_ordered[:m_opt].sum()
                f_g_hat[:, j, k] = (
                    1
                    / m_opt
                    * (s_g_sum - np.sqrt(d_g) * penalty)
                    * smooth_fit[:, j, k]
                    / s_g_ordered[k]
                )
            f_g_hat[:, j, k] -= f_g_hat[:, j, k].mean()
    return f_g_hat

Simulation

Utility classes and functions related to gresit.

Copyright (c) 2025 Robert Bosch GmbH

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.

Equation

Abstract class for non-linear equations.

Source code in gresit/synthetic_data.py

class Equation(metaclass=ABCMeta):
    """Abstract class for non-linear equations."""

    def __init__(
        self,
        group_size: int,
        input_dim: int,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
    ):
        """Abstract class for non-linear equations.

        Args:
            input_dim (int): parent dimension
            group_size (int, optional): _description_. Defaults to 2.
            rng (np.random.Generator, optional): Random number generator. Defaults to default rng.
        """
        self.group_size = group_size
        self.input_dim = input_dim
        self.rng = rng

    @abstractmethod
    def __call__(self, x: np.ndarray) -> np.ndarray:
        """Abstract call method for equation."""

__call__(x) abstractmethod

Abstract call method for equation.

Source code in gresit/synthetic_data.py

@abstractmethod
def __call__(self, x: np.ndarray) -> np.ndarray:
    """Abstract call method for equation."""

__init__(group_size, input_dim, rng=np.random.default_rng(seed=2024))

Abstract class for non-linear equations.

Parameters:

Name	Type	Description	Default
input_dim	int	parent dimension	required
group_size	int	description. Defaults to 2.	required
rng	Generator	Random number generator. Defaults to default rng.	default_rng(seed=2024)

Source code in gresit/synthetic_data.py

def __init__(
    self,
    group_size: int,
    input_dim: int,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
):
    """Abstract class for non-linear equations.

    Args:
        input_dim (int): parent dimension
        group_size (int, optional): _description_. Defaults to 2.
        rng (np.random.Generator, optional): Random number generator. Defaults to default rng.
    """
    self.group_size = group_size
    self.input_dim = input_dim
    self.rng = rng

FCNN

Bases: Equation

A randomly initialized fully connected neural net.

Source code in gresit/synthetic_data.py

class FCNN(Equation):
    """A randomly initialized fully connected neural net."""

    def __init__(
        self,
        group_size: int,
        input_dim: int,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
        hidden_layer: int = 10,
    ) -> None:
        """A randomly initialized fully connected neural net.

        Args:
            input_dim (int): parent dimension
            group_size (int, optional): _description_. Defaults to 2.
            rng (np.random.Generator, optional): Random number generator. Defaults to None.
            hidden_layer (int, optional): Size of hidden dimension.
        """
        super().__init__(group_size, input_dim, rng=rng)

        layers = [
            nn.Linear(self.group_size * self.input_dim, hidden_layer, dtype=torch.float32),
            nn.Sigmoid(),
            nn.Linear(hidden_layer, self.group_size, dtype=torch.float32),
        ]
        self.fcn = nn.Sequential(*layers)
        self.fcn.apply(self.init_weights)

    def init_weights(self, m: nn.Module) -> None:
        """Initializes the weights of the neural net.

        Args:
            m (nn.Module): The layer to be initialized
        """
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0, std=1)
            nn.init.normal_(m.bias, mean=0, std=1)

    def __call__(self, x: np.ndarray) -> np.ndarray:
        """Computes the right hand side of the equation.

        Args:
            x (np.ndarray): Input data

        Returns:
            np.ndarray: output data of size n x multioutput_dim
        """
        X = torch.from_numpy(x).flatten(1, -1).float()

        with torch.no_grad():
            return self.fcn(X).detach().cpu().numpy()

__call__(x)

Computes the right hand side of the equation.

Parameters:

Name	Type	Description	Default
x	ndarray	Input data	required

Returns:

Type	Description
ndarray	np.ndarray: output data of size n x multioutput_dim

Source code in gresit/synthetic_data.py

def __call__(self, x: np.ndarray) -> np.ndarray:
    """Computes the right hand side of the equation.

    Args:
        x (np.ndarray): Input data

    Returns:
        np.ndarray: output data of size n x multioutput_dim
    """
    X = torch.from_numpy(x).flatten(1, -1).float()

    with torch.no_grad():
        return self.fcn(X).detach().cpu().numpy()

__init__(group_size, input_dim, rng=np.random.default_rng(seed=2024), hidden_layer=10)

A randomly initialized fully connected neural net.

Parameters:

Name	Type	Description	Default
input_dim	int	parent dimension	required
group_size	int	description. Defaults to 2.	required
rng	Generator	Random number generator. Defaults to None.	default_rng(seed=2024)
hidden_layer	int	Size of hidden dimension.	10

Source code in gresit/synthetic_data.py

def __init__(
    self,
    group_size: int,
    input_dim: int,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
    hidden_layer: int = 10,
) -> None:
    """A randomly initialized fully connected neural net.

    Args:
        input_dim (int): parent dimension
        group_size (int, optional): _description_. Defaults to 2.
        rng (np.random.Generator, optional): Random number generator. Defaults to None.
        hidden_layer (int, optional): Size of hidden dimension.
    """
    super().__init__(group_size, input_dim, rng=rng)

    layers = [
        nn.Linear(self.group_size * self.input_dim, hidden_layer, dtype=torch.float32),
        nn.Sigmoid(),
        nn.Linear(hidden_layer, self.group_size, dtype=torch.float32),
    ]
    self.fcn = nn.Sequential(*layers)
    self.fcn.apply(self.init_weights)

init_weights(m)

Initializes the weights of the neural net.

Parameters:

Name	Type	Description	Default
m	Module	The layer to be initialized	required

Source code in gresit/synthetic_data.py

def init_weights(self, m: nn.Module) -> None:
    """Initializes the weights of the neural net.

    Args:
        m (nn.Module): The layer to be initialized
    """
    if isinstance(m, nn.Linear):
        nn.init.normal_(m.weight, mean=0, std=1)
        nn.init.normal_(m.bias, mean=0, std=1)

GaussianProcesses

Bases: Equation

A weighted sum of Gaussian processes.

Source code in gresit/synthetic_data.py

class GaussianProcesses(Equation):
    """A weighted sum of Gaussian processes."""

    def __init__(
        self,
        group_size: int,
        input_dim: int,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
        n_gp: int = 5,
        sigma: float = 0.3,
    ) -> None:
        """A weighted sum of Gaussian processes.

        Args:
            input_dim (int): parent dimension
            group_size (int, optional): _description_. Defaults to 2.
            rng (np.random.Generator, optional): Random number generator. Defaults to None.
            n_gp (int): Number of processes
            sigma (float): The standard deviation of GPs. Note that all GPS have the same std.

        Raises:
            ValueError: Gets thrown when input dimensions don't match.
        """
        super().__init__(group_size, input_dim, rng=rng)

        self.n_gp = n_gp

        mean_vec = self.rng.random((self.group_size, self.input_dim, self.n_gp)) * 2.0 - 0.5
        # Weights in GPs.
        w1 = self.rng.random((self.group_size, self.input_dim, self.n_gp))
        w1 /= w1.sum(axis=-1, keepdims=True)
        # Weights for a linear aggregation of GPs.
        w2 = self.rng.random((self.group_size, self.input_dim))
        w2 /= w2.sum(axis=-1, keepdims=True)

        (_, n1, _) = mean_vec.shape
        (_, n2, _) = w1.shape
        (_, n3) = w2.shape

        if n1 != n2 or n1 != n3:
            raise ValueError("Input variable dimensions need to match.")

        # add first axis corresponding to the number of samples.
        self.mean_vec = mean_vec[np.newaxis]
        self.w1 = w1[np.newaxis]
        self.w2 = w2[np.newaxis]
        self.sigma = sigma

    @staticmethod
    def from_params(
        mean_vec: np.ndarray,
        w1: np.ndarray,
        w2: np.ndarray,
        sigma: float = 0.3,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
    ) -> "GaussianProcesses":
        """Creates a gaussian process from parameters.

        Args:
            mean_vec (np.ndarray): Mean vectors of Gaussian processes with
                shape=(group_size, num_parents, num_kernels), where num_parents is
                the dimension of the input variable and num_kernels is
                the number of Gaussian kernels in one GP.
            w1 (np.ndarray): Weights of GPs with shape=(num_parents, num_kernels).
            w2 (np.ndarray): Weights for aggregating GPs with shape=(n,)
            sigma (float): The standard deviation of GPs. Note that all GPS have the same std.
            rng (np.random.Generator, optional): The random number generator to use.
        """
        gp = GaussianProcesses(group_size=mean_vec.shape[0], input_dim=mean_vec.shape[1], rng=rng)

        gp.mean_vec = mean_vec[np.newaxis]
        gp.w1 = w1[np.newaxis]
        gp.w2 = w2[np.newaxis]
        gp.sigma = sigma
        return gp

    def __call__(self, x: np.ndarray) -> np.ndarray:
        """Depending on the shape of x dimension will be adjusted.

        Args:
            x (float | np.ndarray): Input data

        Raises:
            ValueError: Gets thrown when input dimensions don't match.

        Returns:
            float | np.ndarray: output data of size n x multioutput_dim
        """
        kernel_list = []
        for i in range(self.mean_vec.shape[-1]):
            kernel_list.append(
                self._gaussian_kernel_function(
                    x_1=x, x_2=self.mean_vec[:, :, :, i], sigma=self.sigma
                )
            )

        kernel = np.stack(kernel_list, axis=x.ndim)
        if not kernel.shape == (*x.shape, self.mean_vec.shape[-1]):
            raise ValueError(
                f"""something went wrong, shape of kernel function should be \
                {(*x.shape, self.mean_vec.shape[-1])}, \
                but in fact it is {kernel.shape}."""
            )
        # kernel should have shape #samples, #group_size, #parents #kernels
        # Weighted sum of q Gaussian kernels.
        gps: np.ndarray = (kernel * self.w1).sum(axis=-1)
        # Weighted sum of GPs (doesn't do anythin if there is only one parent)
        gp: np.ndarray = (gps * self.w2).sum(axis=-1)

        if not gp.shape == (x.shape[0], x.shape[1]):
            raise ValueError(
                f"""something went wrong, shape of function output should be \
                {x.shape}, but in fact it is {gp.shape}."""
            )

        return gp

    def _gaussian_kernel_function(
        self, x_1: np.ndarray, x_2: np.ndarray, sigma: np.ndarray | float
    ) -> np.ndarray:
        """Gaussian kernel."""
        return np.exp(-0.5 * np.square(x_1 - x_2) / np.square(sigma))

__call__(x)

Depending on the shape of x dimension will be adjusted.

Parameters:

Name	Type	Description	Default
x	float | ndarray	Input data	required

Raises:

Type	Description
ValueError	Gets thrown when input dimensions don't match.

Returns:

Type	Description
ndarray	float | np.ndarray: output data of size n x multioutput_dim

Source code in gresit/synthetic_data.py

def __call__(self, x: np.ndarray) -> np.ndarray:
    """Depending on the shape of x dimension will be adjusted.

    Args:
        x (float | np.ndarray): Input data

    Raises:
        ValueError: Gets thrown when input dimensions don't match.

    Returns:
        float | np.ndarray: output data of size n x multioutput_dim
    """
    kernel_list = []
    for i in range(self.mean_vec.shape[-1]):
        kernel_list.append(
            self._gaussian_kernel_function(
                x_1=x, x_2=self.mean_vec[:, :, :, i], sigma=self.sigma
            )
        )

    kernel = np.stack(kernel_list, axis=x.ndim)
    if not kernel.shape == (*x.shape, self.mean_vec.shape[-1]):
        raise ValueError(
            f"""something went wrong, shape of kernel function should be \
            {(*x.shape, self.mean_vec.shape[-1])}, \
            but in fact it is {kernel.shape}."""
        )
    # kernel should have shape #samples, #group_size, #parents #kernels
    # Weighted sum of q Gaussian kernels.
    gps: np.ndarray = (kernel * self.w1).sum(axis=-1)
    # Weighted sum of GPs (doesn't do anythin if there is only one parent)
    gp: np.ndarray = (gps * self.w2).sum(axis=-1)

    if not gp.shape == (x.shape[0], x.shape[1]):
        raise ValueError(
            f"""something went wrong, shape of function output should be \
            {x.shape}, but in fact it is {gp.shape}."""
        )

    return gp

__init__(group_size, input_dim, rng=np.random.default_rng(seed=2024), n_gp=5, sigma=0.3)

A weighted sum of Gaussian processes.

Parameters:

Name	Type	Description	Default
input_dim	int	parent dimension	required
group_size	int	description. Defaults to 2.	required
rng	Generator	Random number generator. Defaults to None.	default_rng(seed=2024)
n_gp	int	Number of processes	5
sigma	float	The standard deviation of GPs. Note that all GPS have the same std.	0.3

Raises:

Type	Description
ValueError	Gets thrown when input dimensions don't match.

Source code in gresit/synthetic_data.py

def __init__(
    self,
    group_size: int,
    input_dim: int,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
    n_gp: int = 5,
    sigma: float = 0.3,
) -> None:
    """A weighted sum of Gaussian processes.

    Args:
        input_dim (int): parent dimension
        group_size (int, optional): _description_. Defaults to 2.
        rng (np.random.Generator, optional): Random number generator. Defaults to None.
        n_gp (int): Number of processes
        sigma (float): The standard deviation of GPs. Note that all GPS have the same std.

    Raises:
        ValueError: Gets thrown when input dimensions don't match.
    """
    super().__init__(group_size, input_dim, rng=rng)

    self.n_gp = n_gp

    mean_vec = self.rng.random((self.group_size, self.input_dim, self.n_gp)) * 2.0 - 0.5
    # Weights in GPs.
    w1 = self.rng.random((self.group_size, self.input_dim, self.n_gp))
    w1 /= w1.sum(axis=-1, keepdims=True)
    # Weights for a linear aggregation of GPs.
    w2 = self.rng.random((self.group_size, self.input_dim))
    w2 /= w2.sum(axis=-1, keepdims=True)

    (_, n1, _) = mean_vec.shape
    (_, n2, _) = w1.shape
    (_, n3) = w2.shape

    if n1 != n2 or n1 != n3:
        raise ValueError("Input variable dimensions need to match.")

    # add first axis corresponding to the number of samples.
    self.mean_vec = mean_vec[np.newaxis]
    self.w1 = w1[np.newaxis]
    self.w2 = w2[np.newaxis]
    self.sigma = sigma

from_params(mean_vec, w1, w2, sigma=0.3, rng=np.random.default_rng(seed=2024)) staticmethod

Creates a gaussian process from parameters.

Parameters:

Name	Type	Description	Default
mean_vec	ndarray	Mean vectors of Gaussian processes with shape=(group_size, num_parents, num_kernels), where num_parents is the dimension of the input variable and num_kernels is the number of Gaussian kernels in one GP.	required
w1	ndarray	Weights of GPs with shape=(num_parents, num_kernels).	required
w2	ndarray	Weights for aggregating GPs with shape=(n,)	required
sigma	float	The standard deviation of GPs. Note that all GPS have the same std.	0.3
rng	Generator	The random number generator to use.	default_rng(seed=2024)

Source code in gresit/synthetic_data.py

@staticmethod
def from_params(
    mean_vec: np.ndarray,
    w1: np.ndarray,
    w2: np.ndarray,
    sigma: float = 0.3,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
) -> "GaussianProcesses":
    """Creates a gaussian process from parameters.

    Args:
        mean_vec (np.ndarray): Mean vectors of Gaussian processes with
            shape=(group_size, num_parents, num_kernels), where num_parents is
            the dimension of the input variable and num_kernels is
            the number of Gaussian kernels in one GP.
        w1 (np.ndarray): Weights of GPs with shape=(num_parents, num_kernels).
        w2 (np.ndarray): Weights for aggregating GPs with shape=(n,)
        sigma (float): The standard deviation of GPs. Note that all GPS have the same std.
        rng (np.random.Generator, optional): The random number generator to use.
    """
    gp = GaussianProcesses(group_size=mean_vec.shape[0], input_dim=mean_vec.shape[1], rng=rng)

    gp.mean_vec = mean_vec[np.newaxis]
    gp.w1 = w1[np.newaxis]
    gp.w2 = w2[np.newaxis]
    gp.sigma = sigma
    return gp

GenChainedData

Bases: GenData

Class to generate chain DAG with nonlinear data following multivariate ANM (mANM).

Source code in gresit/synthetic_data.py

class GenChainedData(GenData):
    """Class to generate chain DAG with nonlinear data following multivariate ANM (mANM)."""

    def __init__(
        self,
        number_of_nodes: int = 15,
        equation_cls: type[Equation] = GaussianProcesses,
        equation_kwargs: dict[str, Any] | None = None,
        group_size: int = 2,
        edge_density: float = 0.2,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
        snr: float = 1.0,
        noise_distribution: str = "gaussian",
    ) -> None:
        """Initiate the layered DAG.

        Args:
            rng (np.random.Generator, optional): Random number generator. Defaults to None.
            number_of_nodes (int, optional): _description_. Defaults to 15.
            equation_cls (Equation): The type of the equation that should be used.
            equation_kwargs (dict): Arguments for equation.
            group_size (int, optional): Number of entries in groups. Defaults to 2.
            edge_density (float, optional): _description_. Defaults to 0.2.
            snr (float, optional): Signal to noise ratio. Defaults to 1.0.
            noise_distribution (str, optional): Which distribution to choose for the noise.
                defaults to Gaussian noise. Options include: `lognormal`
        """
        super().__init__(
            number_of_nodes=number_of_nodes,
            equation_cls=equation_cls,
            equation_kwargs=equation_kwargs,
            group_size=group_size,
            edge_density=edge_density,
            rng=rng,
            snr=snr,
            noise_distribution=noise_distribution,
        )
        self.dag = DAG()
        self._initiate()
        self.causal_order = self.dag.causal_order

    def _make_dag(self) -> None:
        """Generate an chain Directed Acyclic Graph (DAG)."""
        # Assign a random topological ordering
        nodes = [f"X_{i}" for i in range(self.number_of_nodes)]
        self.rng.shuffle(nodes)
        self.causal_order = nodes
        chain_edges = [(nodes[i], nodes[i + 1]) for i in range(self.number_of_nodes - 1)]
        G = DAG(nodes=nodes, edges=chain_edges)
        for i in range(self.number_of_nodes):
            for j in range(i + 2, self.number_of_nodes):
                if self.rng.random() < self.edge_density:
                    G.add_edge(edge=(nodes[i], nodes[j]))
        self.dag = G

__init__(number_of_nodes=15, equation_cls=GaussianProcesses, equation_kwargs=None, group_size=2, edge_density=0.2, rng=np.random.default_rng(seed=2024), snr=1.0, noise_distribution='gaussian')

Initiate the layered DAG.

Parameters:

Name	Type	Description	Default
rng	Generator	Random number generator. Defaults to None.	default_rng(seed=2024)
number_of_nodes	int	description. Defaults to 15.	15
equation_cls	Equation	The type of the equation that should be used.	GaussianProcesses
equation_kwargs	dict	Arguments for equation.	None
group_size	int	Number of entries in groups. Defaults to 2.	2
edge_density	float	description. Defaults to 0.2.	0.2
snr	float	Signal to noise ratio. Defaults to 1.0.	1.0
noise_distribution	str	Which distribution to choose for the noise. defaults to Gaussian noise. Options include: lognormal	'gaussian'

Source code in gresit/synthetic_data.py

def __init__(
    self,
    number_of_nodes: int = 15,
    equation_cls: type[Equation] = GaussianProcesses,
    equation_kwargs: dict[str, Any] | None = None,
    group_size: int = 2,
    edge_density: float = 0.2,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
    snr: float = 1.0,
    noise_distribution: str = "gaussian",
) -> None:
    """Initiate the layered DAG.

    Args:
        rng (np.random.Generator, optional): Random number generator. Defaults to None.
        number_of_nodes (int, optional): _description_. Defaults to 15.
        equation_cls (Equation): The type of the equation that should be used.
        equation_kwargs (dict): Arguments for equation.
        group_size (int, optional): Number of entries in groups. Defaults to 2.
        edge_density (float, optional): _description_. Defaults to 0.2.
        snr (float, optional): Signal to noise ratio. Defaults to 1.0.
        noise_distribution (str, optional): Which distribution to choose for the noise.
            defaults to Gaussian noise. Options include: `lognormal`
    """
    super().__init__(
        number_of_nodes=number_of_nodes,
        equation_cls=equation_cls,
        equation_kwargs=equation_kwargs,
        group_size=group_size,
        edge_density=edge_density,
        rng=rng,
        snr=snr,
        noise_distribution=noise_distribution,
    )
    self.dag = DAG()
    self._initiate()
    self.causal_order = self.dag.causal_order

GenData

Parent class to GenDate classes.

Source code in gresit/synthetic_data.py

class GenData:
    """Parent class to GenDate classes."""

    def __init__(
        self,
        number_of_nodes: int = 15,
        equation_cls: type[Equation] = GaussianProcesses,
        equation_kwargs: dict[str, Any] | None = None,
        group_size: int = 2,
        edge_density: float = 0.2,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
        snr: float = 1.0,
        noise_distribution: str = "gaussian",
    ) -> None:
        """Initiate parent class DAG.

        Args:
            rng (np.random.Generator, optional): Random number generator. Defaults to None.
            number_of_nodes (int, optional): _description_. Defaults to 15.
            equation_cls (Equation): The type of the equation that should be used.
            equation_kwargs (dict): Arguments for equation.
            group_size (int, optional): Number of entries in groups. Defaults to 2.
            edge_density (float, optional): _description_. Defaults to 0.2.
            snr (float, optional): Signal to noise ratio. Defaults to 1.0.
            noise_distribution (str, optional): Which distribution to choose for the noise.
                defaults to Gaussian noise. Options include: `lognormal`
        """
        self.number_of_nodes = number_of_nodes
        self.group_size = group_size
        self.edge_density = edge_density
        self.rng = rng
        self.equation_cls = equation_cls
        self.equation_kwargs = equation_kwargs
        self.dag = DAG()
        self.causal_order: list[str] = []
        self.snr = snr
        self.noise_distribution = noise_distribution

    def _initiate(self) -> None:
        """Initiate DAG and random functions."""
        self._make_dag()
        self._initiate_ANM()

    def _make_dag(self) -> None:
        """Make DAG.

        Raises:
            NotImplementedError: Needs to be overwritten.
        """
        raise NotImplementedError

    def _initiate_ANM(self) -> None:
        """Initiate the ANM equations."""
        equations: dict[str, MultiOutputANM] = {}
        for node in self.dag.nodes:
            parents = self.dag.parents(of_node=node)
            equations[node] = MultiOutputANM(
                input_dim=len(parents),
                equation_cls=self.equation_cls,
                equation_kwargs=self.equation_kwargs,
                group_size=self.group_size,
                rng=self.rng,
                snr=self.snr,
            )

        self.equations = equations

    def _generate_random_correlation_matrix(self, group_size: int) -> np.ndarray:
        """Generate random pd correlation matrix.

        Args:
            group_size (int): dimension of matrix.

        Returns:
            np.ndarray: correlation matrix.
        """
        L = np.tril(self.rng.uniform(-0.8, 0.8, (group_size, group_size)), k=-1)
        np.fill_diagonal(L, 1)  # Set diagonal to 1

        # Compute symmetric positive definite matrix
        Sigma = L @ L.T

        # Normalize diagonal to 1
        D = np.sqrt(np.diag(Sigma))
        Sigma = Sigma / np.outer(D, D)

        return Sigma

    def generate_data(self, num_samples: int = 1000) -> tuple[dict[str, np.ndarray], np.ndarray]:
        """Sample data from the layered DAG."""
        noise_data = []
        for _ in self.dag.nodes:
            mean = self.rng.uniform(-0.8, 0.8, self.group_size)
            corr = self._generate_random_correlation_matrix(self.group_size)
            if self.noise_distribution == "gaussian":
                noise_data.append(
                    self.rng.multivariate_normal(mean=mean, cov=corr, size=num_samples)
                )
            else:
                noise_data.append(
                    np.exp(self.rng.multivariate_normal(mean=mean, cov=corr, size=num_samples))
                )
        reshaped_noise = np.moveaxis(np.asanyarray(noise_data), 0, -1)

        data = np.zeros(reshaped_noise.shape)
        data_dict = {}
        for i, node in enumerate(self.causal_order):
            parents = self.dag.parents(of_node=node)
            parent_indices = [self.dag.causal_order.index(parent) for parent in parents]
            data[:, :, i] = self.equations[node].apply_rhs(
                parent_data=data[:, :, parent_indices], noise_data=reshaped_noise[:, :, i]
            )
            data[:, :, i] -= data[:, :, i].mean(axis=0)
            data[:, :, i] /= data[:, :, i].std(axis=0)
            data_dict[node] = data[:, :, i]
        return data_dict, data

__init__(number_of_nodes=15, equation_cls=GaussianProcesses, equation_kwargs=None, group_size=2, edge_density=0.2, rng=np.random.default_rng(seed=2024), snr=1.0, noise_distribution='gaussian')

Initiate parent class DAG.

Parameters:

Name	Type	Description	Default
rng	Generator	Random number generator. Defaults to None.	default_rng(seed=2024)
number_of_nodes	int	description. Defaults to 15.	15
equation_cls	Equation	The type of the equation that should be used.	GaussianProcesses
equation_kwargs	dict	Arguments for equation.	None
group_size	int	Number of entries in groups. Defaults to 2.	2
edge_density	float	description. Defaults to 0.2.	0.2
snr	float	Signal to noise ratio. Defaults to 1.0.	1.0
noise_distribution	str	Which distribution to choose for the noise. defaults to Gaussian noise. Options include: lognormal	'gaussian'

Source code in gresit/synthetic_data.py

def __init__(
    self,
    number_of_nodes: int = 15,
    equation_cls: type[Equation] = GaussianProcesses,
    equation_kwargs: dict[str, Any] | None = None,
    group_size: int = 2,
    edge_density: float = 0.2,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
    snr: float = 1.0,
    noise_distribution: str = "gaussian",
) -> None:
    """Initiate parent class DAG.

    Args:
        rng (np.random.Generator, optional): Random number generator. Defaults to None.
        number_of_nodes (int, optional): _description_. Defaults to 15.
        equation_cls (Equation): The type of the equation that should be used.
        equation_kwargs (dict): Arguments for equation.
        group_size (int, optional): Number of entries in groups. Defaults to 2.
        edge_density (float, optional): _description_. Defaults to 0.2.
        snr (float, optional): Signal to noise ratio. Defaults to 1.0.
        noise_distribution (str, optional): Which distribution to choose for the noise.
            defaults to Gaussian noise. Options include: `lognormal`
    """
    self.number_of_nodes = number_of_nodes
    self.group_size = group_size
    self.edge_density = edge_density
    self.rng = rng
    self.equation_cls = equation_cls
    self.equation_kwargs = equation_kwargs
    self.dag = DAG()
    self.causal_order: list[str] = []
    self.snr = snr
    self.noise_distribution = noise_distribution

generate_data(num_samples=1000)

Sample data from the layered DAG.

Source code in gresit/synthetic_data.py

def generate_data(self, num_samples: int = 1000) -> tuple[dict[str, np.ndarray], np.ndarray]:
    """Sample data from the layered DAG."""
    noise_data = []
    for _ in self.dag.nodes:
        mean = self.rng.uniform(-0.8, 0.8, self.group_size)
        corr = self._generate_random_correlation_matrix(self.group_size)
        if self.noise_distribution == "gaussian":
            noise_data.append(
                self.rng.multivariate_normal(mean=mean, cov=corr, size=num_samples)
            )
        else:
            noise_data.append(
                np.exp(self.rng.multivariate_normal(mean=mean, cov=corr, size=num_samples))
            )
    reshaped_noise = np.moveaxis(np.asanyarray(noise_data), 0, -1)

    data = np.zeros(reshaped_noise.shape)
    data_dict = {}
    for i, node in enumerate(self.causal_order):
        parents = self.dag.parents(of_node=node)
        parent_indices = [self.dag.causal_order.index(parent) for parent in parents]
        data[:, :, i] = self.equations[node].apply_rhs(
            parent_data=data[:, :, parent_indices], noise_data=reshaped_noise[:, :, i]
        )
        data[:, :, i] -= data[:, :, i].mean(axis=0)
        data[:, :, i] /= data[:, :, i].std(axis=0)
        data_dict[node] = data[:, :, i]
    return data_dict, data

GenERData

Bases: GenData

Class to generate general nonlinear data following Erdos-Renyi (ER) graph.

Source code in gresit/synthetic_data.py

class GenERData(GenData):
    """Class to generate general nonlinear data following Erdos-Renyi (ER) graph."""

    def __init__(
        self,
        number_of_nodes: int = 15,
        equation_cls: type[Equation] = GaussianProcesses,
        equation_kwargs: dict[str, Any] | None = None,
        group_size: int = 2,
        edge_density: float = 0.2,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
        snr: float = 1.0,
        noise_distribution: str = "gaussian",
    ) -> None:
        """Initiate the ER DAG.

        Args:
            rng (np.random.Generator, optional): Random number generator. Defaults to None.
            number_of_nodes (int, optional): _description_. Defaults to 15.
            equation_cls (Equation): The type of the equation that should be used.
            equation_kwargs (dict): Arguments for equation.
            group_size (int, optional): Number of entries in groups. Defaults to 2.
            edge_density (float, optional): _description_. Defaults to 0.2.
            snr (float, optional): Signal to noise ratio. Defaults to 1.0.
            noise_distribution (str, optional): Which distribution to choose for the noise.
                defaults to Gaussian noise. Options include: `lognormal`
        """
        super().__init__(
            number_of_nodes=number_of_nodes,
            equation_cls=equation_cls,
            equation_kwargs=equation_kwargs,
            group_size=group_size,
            edge_density=edge_density,
            rng=rng,
            snr=snr,
            noise_distribution=noise_distribution,
        )
        self.dag = DAG()
        self._initiate()
        self.causal_order = self.dag.causal_order

    def _make_dag(self) -> None:
        """Generate an Erdős-Rényi Directed Acyclic Graph (DAG)."""
        # Assign a random topological ordering
        nodes = [f"X_{i}" for i in range(self.number_of_nodes)]
        self.rng.shuffle(nodes)

        G = DAG(nodes=nodes)

        # Add edges based on Erdős-Rényi model
        for i, j in combinations(range(self.number_of_nodes), 2):
            # Ensure edges go from lower to higher index
            if self.rng.random() < self.edge_density:
                G.add_edge(edge=(nodes[i], nodes[j]))

        self.dag = G

__init__(number_of_nodes=15, equation_cls=GaussianProcesses, equation_kwargs=None, group_size=2, edge_density=0.2, rng=np.random.default_rng(seed=2024), snr=1.0, noise_distribution='gaussian')

Initiate the ER DAG.

Parameters:

Name	Type	Description	Default
rng	Generator	Random number generator. Defaults to None.	default_rng(seed=2024)
number_of_nodes	int	description. Defaults to 15.	15
equation_cls	Equation	The type of the equation that should be used.	GaussianProcesses
equation_kwargs	dict	Arguments for equation.	None
group_size	int	Number of entries in groups. Defaults to 2.	2
edge_density	float	description. Defaults to 0.2.	0.2
snr	float	Signal to noise ratio. Defaults to 1.0.	1.0
noise_distribution	str	Which distribution to choose for the noise. defaults to Gaussian noise. Options include: lognormal	'gaussian'

Source code in gresit/synthetic_data.py

def __init__(
    self,
    number_of_nodes: int = 15,
    equation_cls: type[Equation] = GaussianProcesses,
    equation_kwargs: dict[str, Any] | None = None,
    group_size: int = 2,
    edge_density: float = 0.2,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
    snr: float = 1.0,
    noise_distribution: str = "gaussian",
) -> None:
    """Initiate the ER DAG.

    Args:
        rng (np.random.Generator, optional): Random number generator. Defaults to None.
        number_of_nodes (int, optional): _description_. Defaults to 15.
        equation_cls (Equation): The type of the equation that should be used.
        equation_kwargs (dict): Arguments for equation.
        group_size (int, optional): Number of entries in groups. Defaults to 2.
        edge_density (float, optional): _description_. Defaults to 0.2.
        snr (float, optional): Signal to noise ratio. Defaults to 1.0.
        noise_distribution (str, optional): Which distribution to choose for the noise.
            defaults to Gaussian noise. Options include: `lognormal`
    """
    super().__init__(
        number_of_nodes=number_of_nodes,
        equation_cls=equation_cls,
        equation_kwargs=equation_kwargs,
        group_size=group_size,
        edge_density=edge_density,
        rng=rng,
        snr=snr,
        noise_distribution=noise_distribution,
    )
    self.dag = DAG()
    self._initiate()
    self.causal_order = self.dag.causal_order

GenLayeredData

Bases: GenData

Class to generate general nonlinear data following multivariate ANM (mANM).

Source code in gresit/synthetic_data.py

class GenLayeredData(GenData):
    """Class to generate general nonlinear data following multivariate ANM (mANM)."""

    def __init__(
        self,
        number_of_nodes: int = 15,
        number_of_layers: int = 3,
        equation_cls: type[Equation] = GaussianProcesses,
        equation_kwargs: dict[str, Any] | None = None,
        group_size: int = 2,
        edge_density: float = 0.2,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
        snr: float = 1.0,
        noise_distribution: str = "gaussian",
    ) -> None:
        """Initiate the layered DAG.

        Args:
            rng (np.random.Generator, optional): Random number generator. Defaults to None.
            number_of_nodes (int, optional): _description_. Defaults to 15.
            number_of_layers (int, optional): _description_. Defaults to 3.
            equation_cls (Equation): The type of the equation that should be used.
            equation_kwargs (dict): Arguments for equation.
            group_size (int, optional): Number of entries in groups. Defaults to 2.
            edge_density (float, optional): _description_. Defaults to 0.2.
            snr (float, optional): Signal to noise ratio. Defaults to 1.0.
            noise_distribution (str, optional): Which distribution to choose for the noise.
                defaults to Gaussian noise. Options include: `lognormal`
        """
        super().__init__(
            number_of_nodes=number_of_nodes,
            equation_cls=equation_cls,
            equation_kwargs=equation_kwargs,
            group_size=group_size,
            edge_density=edge_density,
            rng=rng,
            snr=snr,
            noise_distribution=noise_distribution,
        )
        self.number_of_layers = number_of_layers
        self.dag: LayeredDAG = LayeredDAG()
        self._initiate()
        self.layering = self.dag.layering

    def _divide_equal(self, number: int, pieces: int) -> list[int]:
        """Divides an integer into a specified number of pieces, distributing the remainder.

        Args:
            number: The integer to divide.
            pieces: The desired number of pieces.

        Returns:
            A list containing the values of each piece.
        """
        base_size = number // pieces
        remainder = number % pieces
        return [base_size + 1 if i < remainder else base_size for i in range(pieces)]

    def all_causal_orderings(self) -> list[list[str]]:
        """Valid causal ordering of the layered DAG.

        Raises:
            ValueError: If there are no nodes or edges.

        Returns:
            list[list[str]]: list of causal orders.
        """
        if not self.dag.nodes:
            raise ValueError("There are no nodes in the graph")

        if not self.layering:
            raise ValueError("Layering must be provided.")

        all_within_orders = [
            list(nx.all_topological_sorts(self.dag.layer_induced_subgraph(layer).to_networkx()))
            for layer in self.layering.values()
        ]

        all_orderings_tupled = list(product(*all_within_orders))
        all_orderings = []
        for order_tuple in all_orderings_tupled:
            all_orderings.append([x for xs in [portion for portion in order_tuple] for x in xs])
        return all_orderings

    def _make_dag(self) -> None:
        """Create a layered DAG."""
        nodes_per_layer = self._divide_equal(self.number_of_nodes, self.number_of_layers)

        layering = defaultdict(list)
        for i, layer_i in enumerate(nodes_per_layer):
            for j in range(layer_i):
                layering[f"L_{i}"].append(f"X_{i}_{j}")

        for list_of_nodes in layering.values():
            self.rng.shuffle(list_of_nodes)

        self.dag.layering = layering

        causal_ordering = [x for xs in [k for k in layering.values()] for x in xs]
        self.causal_order = causal_ordering
        # Add edges based on Erdős-Rényi model

        self.dag.add_nodes_from(causal_ordering)

        for i in range(len(causal_ordering)):
            for j in range(i + 1, len(causal_ordering)):
                if self.rng.random() < self.edge_density:
                    self.dag.add_edge(edge=(causal_ordering[i], causal_ordering[j]))

__init__(number_of_nodes=15, number_of_layers=3, equation_cls=GaussianProcesses, equation_kwargs=None, group_size=2, edge_density=0.2, rng=np.random.default_rng(seed=2024), snr=1.0, noise_distribution='gaussian')

Initiate the layered DAG.

Parameters:

Name	Type	Description	Default
rng	Generator	Random number generator. Defaults to None.	default_rng(seed=2024)
number_of_nodes	int	description. Defaults to 15.	15
number_of_layers	int	description. Defaults to 3.	3
equation_cls	Equation	The type of the equation that should be used.	GaussianProcesses
equation_kwargs	dict	Arguments for equation.	None
group_size	int	Number of entries in groups. Defaults to 2.	2
edge_density	float	description. Defaults to 0.2.	0.2
snr	float	Signal to noise ratio. Defaults to 1.0.	1.0
noise_distribution	str	Which distribution to choose for the noise. defaults to Gaussian noise. Options include: lognormal	'gaussian'

Source code in gresit/synthetic_data.py

def __init__(
    self,
    number_of_nodes: int = 15,
    number_of_layers: int = 3,
    equation_cls: type[Equation] = GaussianProcesses,
    equation_kwargs: dict[str, Any] | None = None,
    group_size: int = 2,
    edge_density: float = 0.2,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
    snr: float = 1.0,
    noise_distribution: str = "gaussian",
) -> None:
    """Initiate the layered DAG.

    Args:
        rng (np.random.Generator, optional): Random number generator. Defaults to None.
        number_of_nodes (int, optional): _description_. Defaults to 15.
        number_of_layers (int, optional): _description_. Defaults to 3.
        equation_cls (Equation): The type of the equation that should be used.
        equation_kwargs (dict): Arguments for equation.
        group_size (int, optional): Number of entries in groups. Defaults to 2.
        edge_density (float, optional): _description_. Defaults to 0.2.
        snr (float, optional): Signal to noise ratio. Defaults to 1.0.
        noise_distribution (str, optional): Which distribution to choose for the noise.
            defaults to Gaussian noise. Options include: `lognormal`
    """
    super().__init__(
        number_of_nodes=number_of_nodes,
        equation_cls=equation_cls,
        equation_kwargs=equation_kwargs,
        group_size=group_size,
        edge_density=edge_density,
        rng=rng,
        snr=snr,
        noise_distribution=noise_distribution,
    )
    self.number_of_layers = number_of_layers
    self.dag: LayeredDAG = LayeredDAG()
    self._initiate()
    self.layering = self.dag.layering

all_causal_orderings()

Valid causal ordering of the layered DAG.

Raises:

Type	Description
ValueError	If there are no nodes or edges.

Returns:

Type	Description
list[list[str]]	list[list[str]]: list of causal orders.

Source code in gresit/synthetic_data.py

def all_causal_orderings(self) -> list[list[str]]:
    """Valid causal ordering of the layered DAG.

    Raises:
        ValueError: If there are no nodes or edges.

    Returns:
        list[list[str]]: list of causal orders.
    """
    if not self.dag.nodes:
        raise ValueError("There are no nodes in the graph")

    if not self.layering:
        raise ValueError("Layering must be provided.")

    all_within_orders = [
        list(nx.all_topological_sorts(self.dag.layer_induced_subgraph(layer).to_networkx()))
        for layer in self.layering.values()
    ]

    all_orderings_tupled = list(product(*all_within_orders))
    all_orderings = []
    for order_tuple in all_orderings_tupled:
        all_orderings.append([x for xs in [portion for portion in order_tuple] for x in xs])
    return all_orderings

MultiOutputANM

Class to construct nonlinear multi-outcome data that follows an ANM.

Source code in gresit/synthetic_data.py

class MultiOutputANM:
    """Class to construct nonlinear multi-outcome data that follows an ANM."""

    def __init__(
        self,
        input_dim: int,
        equation_cls: type[Equation] = GaussianProcesses,
        equation_kwargs: dict[str, Any] | None = None,
        group_size: int = 2,
        snr: float = 1.0,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
    ) -> None:
        """Initiates the ANM object.

        Args:
            input_dim (int): parent dimension
            group_size (int, optional): _description_. Defaults to 2.
            rng (np.random.Generator, optional): Random number generator. Defaults to None.
            equation_cls (Equation): The type of the equation that should be used.
            equation_kwargs (dict): Arguments for equation.
            snr (float, optional): Signal-to-noise ratio. Defaults to 1.0.
        """
        self.input_dim = input_dim
        self.group_size = group_size
        self.snr = snr
        self.rng = rng
        self._lambda = self.snr / (1.0 + self.snr)
        self.equation_cls = equation_cls
        if not equation_kwargs:
            equation_kwargs = {}
        self.equation_kwargs = equation_kwargs
        if self.input_dim > 0:
            self.f_nonlinear = self.generate_nonlinear_map()

    def generate_nonlinear_map(self) -> Equation:
        """Nonlinear function."""
        return self.equation_cls(
            group_size=self.group_size,
            input_dim=self.input_dim,
            rng=self.rng,
            **self.equation_kwargs,
        )

    def apply_rhs(self, parent_data: np.ndarray | None, noise_data: np.ndarray) -> np.ndarray:
        """See whether this works."""
        noise_standard_dev = 1 / np.std(noise_data, axis=0)
        noise_std = noise_standard_dev * noise_data

        if self.input_dim > 0:
            y_interim = self.f_nonlinear(parent_data)
            y_standard_dev = 1 / np.std(y_interim, axis=0)
            y_std = y_interim * y_standard_dev
            y = np.sqrt(self._lambda) * y_std + np.sqrt(1.0 - self._lambda) * noise_std
            y = y - np.mean(y, axis=0)
        else:
            y = noise_std
        return y

__init__(input_dim, equation_cls=GaussianProcesses, equation_kwargs=None, group_size=2, snr=1.0, rng=np.random.default_rng(seed=2024))

Initiates the ANM object.

Parameters:

Name	Type	Description	Default
input_dim	int	parent dimension	required
group_size	int	description. Defaults to 2.	2
rng	Generator	Random number generator. Defaults to None.	default_rng(seed=2024)
equation_cls	Equation	The type of the equation that should be used.	GaussianProcesses
equation_kwargs	dict	Arguments for equation.	None
snr	float	Signal-to-noise ratio. Defaults to 1.0.	1.0

Source code in gresit/synthetic_data.py

def __init__(
    self,
    input_dim: int,
    equation_cls: type[Equation] = GaussianProcesses,
    equation_kwargs: dict[str, Any] | None = None,
    group_size: int = 2,
    snr: float = 1.0,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
) -> None:
    """Initiates the ANM object.

    Args:
        input_dim (int): parent dimension
        group_size (int, optional): _description_. Defaults to 2.
        rng (np.random.Generator, optional): Random number generator. Defaults to None.
        equation_cls (Equation): The type of the equation that should be used.
        equation_kwargs (dict): Arguments for equation.
        snr (float, optional): Signal-to-noise ratio. Defaults to 1.0.
    """
    self.input_dim = input_dim
    self.group_size = group_size
    self.snr = snr
    self.rng = rng
    self._lambda = self.snr / (1.0 + self.snr)
    self.equation_cls = equation_cls
    if not equation_kwargs:
        equation_kwargs = {}
    self.equation_kwargs = equation_kwargs
    if self.input_dim > 0:
        self.f_nonlinear = self.generate_nonlinear_map()

apply_rhs(parent_data, noise_data)

See whether this works.

Source code in gresit/synthetic_data.py

def apply_rhs(self, parent_data: np.ndarray | None, noise_data: np.ndarray) -> np.ndarray:
    """See whether this works."""
    noise_standard_dev = 1 / np.std(noise_data, axis=0)
    noise_std = noise_standard_dev * noise_data

    if self.input_dim > 0:
        y_interim = self.f_nonlinear(parent_data)
        y_standard_dev = 1 / np.std(y_interim, axis=0)
        y_std = y_interim * y_standard_dev
        y = np.sqrt(self._lambda) * y_std + np.sqrt(1.0 - self._lambda) * noise_std
        y = y - np.mean(y, axis=0)
    else:
        y = noise_std
    return y

generate_nonlinear_map()

Nonlinear function.

Source code in gresit/synthetic_data.py

def generate_nonlinear_map(self) -> Equation:
    """Nonlinear function."""
    return self.equation_cls(
        group_size=self.group_size,
        input_dim=self.input_dim,
        rng=self.rng,
        **self.equation_kwargs,
    )

Utility classes and functions related to gresit.

Copyright (c) 2025 Robert Bosch GmbH

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.

BenchMarker

Class to instantiate and run simulations.

Source code in gresit/simulation_utils.py

class BenchMarker:
    """Class to instantiate and run simulations."""

    @staticmethod
    def _emptygraph(size: int) -> np.ndarray:
        return np.zeros((size, size))

    def _cpdag_dag_processing(
        self,
        gt: pd.DataFrame,
        learned_graph: GRAPH,
        cpdag_strategy: str = "random_dag",
    ) -> tuple[pd.DataFrame, list[str] | None]:
        """Converts CPDAGs to DAGs according to CPDAG strategy.

        If PDAG is not a proper CPDAG, and orientations lead to cycles,
        we don't do anything.

        Args:
            gt (pd.DataFrame): ground truth adjacency matrix
            learned_graph (GRAPH): learned graph
            cpdag_strategy (str, optional): Conversion strategy. Defaults to "random_dag".

        Raises:
            NotImplementedError: _description_
            ValueError: _description_

        Returns:
            tuple[pd.DataFrame, list[str] | None]: Tuple containing amat and causal order (or None).
        """
        learned_dag: PDAG | DAG
        if isinstance(learned_graph, PDAG):
            if cpdag_strategy == "random_dag":
                try:
                    learned_dag = learned_graph.to_random_dag()
                except ValueError:
                    learned_dag = learned_graph
            elif cpdag_strategy == "best_dag":
                try:
                    all_dags = learned_graph.to_allDAGs()
                    original_order = gt.columns
                    shd_array = np.full(len(all_dags), np.inf)
                    for i, each_dag in enumerate(all_dags):
                        shd_array[i] = shd(
                            gt.to_numpy(dtype=np.int8),
                            each_dag.adjacency_matrix.reindex(
                                index=original_order, columns=original_order
                            ).to_numpy(dtype=np.int8),
                        )[1]
                    learned_dag = all_dags[shd_array.argmin()]
                except ValueError:
                    learned_dag = learned_graph
            else:
                raise NotImplementedError()

            learned_adjacency_matrix = learned_dag.adjacency_matrix
            learned_causal_order = learned_dag.causal_order

        elif isinstance(learned_graph, DAG):
            learned_adjacency_matrix = learned_graph.adjacency_matrix
            learned_causal_order = learned_graph.causal_order

        else:
            raise ValueError("Something went wrong, result is neither `DAG` nor `CPDAG`.")

        return learned_adjacency_matrix, learned_causal_order

    def run_benchmark(
        self,
        params: ExperimentParams,
        num_runs: int = 30,
        metrics: list[str] = ["shd", "sid", "ancestor_aid", "ancester_ordering_aid"],
        cpdag_strategy: str = "random_dag",
    ) -> defaultdict[str, defaultdict[str, list[int | float | None]]]:
        """Benchmark run.

        Args:
            num_runs (int, optional): _description_. Defaults to 30.
            params (ExperimenParams): Params to benchmark
            metrics (list[str], optional): _description_. Defaults to
                ["shd", "sid", "ancestor_aid", "ancester_ordering_aid"].
            cpdag_strategy (str): How to convert CPDAG to DAG to enable fair comparison.
                Options are `random_dag` (default) and `best_dag`. In the former, a random
                DAG is chosen from the MEC, in the latter the SHD is calculated for all DAGs
                in the MEC and the best one is chosen.
        """
        results: defaultdict[str, defaultdict[str, list[int | float | None]]] = defaultdict(
            partial(defaultdict, list)
        )
        for run in range(num_runs):
            # Generate data to also generate a new DAG
            try:
                ldat = params.make_data()
                data, _ = ldat.generate_data(num_samples=params.number_of_samples)

                layering = getattr(ldat.dag, "layering", None)
                gt = ldat.dag.adjacency_matrix

                for algo, name in params:
                    learned_graph = algo.learn_graph(
                        data_dict=data,
                        layering=layering,
                    )

                    learned_adjacency_matrix, learned_causal_order = self._cpdag_dag_processing(
                        gt=gt, learned_graph=learned_graph, cpdag_strategy=cpdag_strategy
                    )

                    for metric in metrics:
                        _, resulting_metric = _metric_mapping[metric](
                            gt,
                            learned_adjacency_matrix,
                            learned_causal_order,
                        )

                        if isinstance(resulting_metric, np.ndarray):
                            resulting_metric = None

                        results[name][metric].append(resulting_metric)
            except Exception as e:
                print(f"Error in run {run}: {e}. Skipping this run.")
                continue
        return results

    def write_results(
        self, results: defaultdict[str, defaultdict[str, list[int | float | None]]], path: str
    ) -> None:
        """Write results to JSON files with time signature.

        Args:
            results (defaultdict[str, defaultdict[str, list[int | float | None]]]): Result
                dict from a benchmark run.
            path (str): Destination path where to save files.
        """
        now = datetime.now()
        current_time = now.strftime("%Y-%m-%d_%H-%M")
        with open(path + f"/results_{current_time}.json", "w", encoding="utf-8") as outfile:
            json.dump(results, outfile)

run_benchmark(params, num_runs=30, metrics=['shd', 'sid', 'ancestor_aid', 'ancester_ordering_aid'], cpdag_strategy='random_dag')

Benchmark run.

Parameters:

Name	Type	Description	Default
num_runs	int	description. Defaults to 30.	30
params	ExperimenParams	Params to benchmark	required
metrics	list[str]	description. Defaults to ["shd", "sid", "ancestor_aid", "ancester_ordering_aid"].	['shd', 'sid', 'ancestor_aid', 'ancester_ordering_aid']
cpdag_strategy	str	How to convert CPDAG to DAG to enable fair comparison. Options are random_dag (default) and best_dag. In the former, a random DAG is chosen from the MEC, in the latter the SHD is calculated for all DAGs in the MEC and the best one is chosen.	'random_dag'

Source code in gresit/simulation_utils.py

def run_benchmark(
    self,
    params: ExperimentParams,
    num_runs: int = 30,
    metrics: list[str] = ["shd", "sid", "ancestor_aid", "ancester_ordering_aid"],
    cpdag_strategy: str = "random_dag",
) -> defaultdict[str, defaultdict[str, list[int | float | None]]]:
    """Benchmark run.

    Args:
        num_runs (int, optional): _description_. Defaults to 30.
        params (ExperimenParams): Params to benchmark
        metrics (list[str], optional): _description_. Defaults to
            ["shd", "sid", "ancestor_aid", "ancester_ordering_aid"].
        cpdag_strategy (str): How to convert CPDAG to DAG to enable fair comparison.
            Options are `random_dag` (default) and `best_dag`. In the former, a random
            DAG is chosen from the MEC, in the latter the SHD is calculated for all DAGs
            in the MEC and the best one is chosen.
    """
    results: defaultdict[str, defaultdict[str, list[int | float | None]]] = defaultdict(
        partial(defaultdict, list)
    )
    for run in range(num_runs):
        # Generate data to also generate a new DAG
        try:
            ldat = params.make_data()
            data, _ = ldat.generate_data(num_samples=params.number_of_samples)

            layering = getattr(ldat.dag, "layering", None)
            gt = ldat.dag.adjacency_matrix

            for algo, name in params:
                learned_graph = algo.learn_graph(
                    data_dict=data,
                    layering=layering,
                )

                learned_adjacency_matrix, learned_causal_order = self._cpdag_dag_processing(
                    gt=gt, learned_graph=learned_graph, cpdag_strategy=cpdag_strategy
                )

                for metric in metrics:
                    _, resulting_metric = _metric_mapping[metric](
                        gt,
                        learned_adjacency_matrix,
                        learned_causal_order,
                    )

                    if isinstance(resulting_metric, np.ndarray):
                        resulting_metric = None

                    results[name][metric].append(resulting_metric)
        except Exception as e:
            print(f"Error in run {run}: {e}. Skipping this run.")
            continue
    return results

write_results(results, path)

Write results to JSON files with time signature.

Parameters:

Name	Type	Description	Default
results	defaultdict[str, defaultdict[str, list[int | float | None]]]	Result dict from a benchmark run.	required
path	str	Destination path where to save files.	required

Source code in gresit/simulation_utils.py

def write_results(
    self, results: defaultdict[str, defaultdict[str, list[int | float | None]]], path: str
) -> None:
    """Write results to JSON files with time signature.

    Args:
        results (defaultdict[str, defaultdict[str, list[int | float | None]]]): Result
            dict from a benchmark run.
        path (str): Destination path where to save files.
    """
    now = datetime.now()
    current_time = now.strftime("%Y-%m-%d_%H-%M")
    with open(path + f"/results_{current_time}.json", "w", encoding="utf-8") as outfile:
        json.dump(results, outfile)

draw_result_boxplots(result_dict, title='', file_path=None, file_name='example_name')

Draw full set of result boxplots.

Parameters:

Name	Type	Description	Default
result_dict	dict	description	required
title	str	title of the plot.	''
file_path	str | None	description. Defaults to None.	None
file_name	str | None	description. Defaults to None.	'example_name'

Source code in gresit/simulation_utils.py

def draw_result_boxplots(
    result_dict: defaultdict[str, defaultdict[str, list[float | None]]],
    title: str = "",
    file_path: str | None = None,
    file_name: str | None = "example_name",
) -> None:
    """Draw full set of result boxplots.

    Args:
        result_dict (dict): _description_
        title (str, optional): title of the plot.
        file_path (str | None, optional): _description_. Defaults to None.
        file_name (str | None, optional): _description_. Defaults to None.
    """
    num_plots = len(result_dict[list(result_dict.keys())[0]].keys())
    MAX_PLOT_PER_ROW = 3
    if num_plots > MAX_PLOT_PER_ROW:
        num_rows = 2
        num_cols = int(np.ceil(num_plots / num_rows))

    fig, ax = plt.subplots(
        figsize=(18, 4.2),
        nrows=num_rows,
        ncols=num_cols,
        sharex=False,
        sharey=True,
    )

    valid_axes = [axi for axi in ax.flat[:] if axi in fig.axes]
    if (num_plots % 2) != 0:
        fig.delaxes(ax[0, -1])
        valid_axes = [axi for axi in ax.flat[:] if axi in fig.axes]

    for axi, metric in zip(
        valid_axes,
        result_dict[list(result_dict.keys())[0]].keys(),
    ):
        _draw_full_boxplot(
            result_dict=result_dict,
            position=axi,
            title="",
            x_label=_metric_names[metric],
            metric=metric,
        )

    plt.figtext(
        0.16,
        1.01,
        r"$\leftarrow$ lower is better",
        horizontalalignment="right",
        fontsize="x-small",
    )
    plt.figtext(
        0.41,
        1.01,
        r"$\rightarrow$ higher is better",
        horizontalalignment="right",
        fontsize="x-small",
    )

    fig.tight_layout()
    fig.suptitle(title)
    fig.subplots_adjust(top=0.88)
    if file_path is not None:
        plt.savefig(f"{file_path}/{file_name}.pdf", bbox_inches="tight")
    plt.show()

Utils

Utility classes and functions related to gresit.

Copyright (c) 2025 Robert Bosch GmbH

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.

Loss functions for training the model.

loss_hsic() taken from https://github.com/danielgreenfeld3/XIC/blob/master/hsic.py as used in paper https://arxiv.org/abs/1910.00270

GaussianKernelMatrix(x, sigma=None)

Get gaussian kernel matrix.

Parameters:

Name	Type	Description	Default
x	Tensor	description	required
sigma	float	description. Defaults the median heuristic when None.	None

Returns:

Name	Type	Description
_type_	Tensor	description

Source code in gresit/losses.py

def GaussianKernelMatrix(x: torch.Tensor, sigma: float | None = None) -> torch.Tensor:
    """Get gaussian kernel matrix.

    Args:
        x (torch.Tensor): _description_
        sigma (float, optional): _description_. Defaults the median heuristic when None.

    Returns:
        _type_: _description_
    """
    pairwise_distances_ = pairwise_distances(x)
    if sigma is None:
        sigma = _median_heuristic(pairwise_distances_)
    return torch.exp(-pairwise_distances_ / sigma)

HSIC(x, y, s_x=None, s_y=None, device='cpu')

Get test statistic.

Parameters:

Name	Type	Description	Default
x	Tensor	description	required
y	Tensor	description	required
s_x	float	description. Defaults to None.	None
s_y	float	description. Defaults to None.	None
device	str	description. Defaults to "cpu".	'cpu'

Returns:

Type	Description
Tensor	torch.Tensor: description

Source code in gresit/losses.py

def HSIC(
    x: torch.Tensor,
    y: torch.Tensor,
    s_x: float | None = None,
    s_y: float | None = None,
    device: str = "cpu",
) -> torch.Tensor:
    """Get test statistic.

    Args:
        x (torch.Tensor): _description_
        y (torch.Tensor): _description_
        s_x (float, optional): _description_. Defaults to None.
        s_y (float, optional): _description_. Defaults to None.
        device (str, optional): _description_. Defaults to "cpu".

    Returns:
        torch.Tensor: _description_
    """
    m = x.shape[0]
    K = GaussianKernelMatrix(x, s_x)
    L = GaussianKernelMatrix(y, s_y)

    H = torch.eye(m) - 1.0 / m * torch.ones((m, m))
    H = H.float().to(device)
    return torch.trace(torch.mm(L, torch.mm(H, torch.mm(K, H)))) / ((m - 1) ** 2)

loss_disco(x, y_pred, y_true, device='cpu')

Get UCRT loss.

Parameters:

Name	Type	Description	Default
x	Tensor	description	required
y_pred	Tensor	description	required
y_true	Tensor	description	required
device	str	description. Defaults to "cpu".	'cpu'

Returns:

Type	Description
Tensor	torch.Tensor: description

Source code in gresit/losses.py

def loss_disco(
    x: torch.Tensor,
    y_pred: torch.Tensor,
    y_true: torch.Tensor,
    device: str = "cpu",
) -> torch.Tensor:
    """Get UCRT loss.

    Args:
        x (torch.Tensor): _description_
        y_pred (torch.Tensor): _description_
        y_true (torch.Tensor): _description_
        device (str, optional): _description_. Defaults to "cpu".

    Returns:
        torch.Tensor: _description_
    """
    r = y_pred - y_true
    return u_distance_cov_squared(a=x, b=r, device=device)

loss_hsic(x, y_pred, y_true, device='cpu')

Get HSIC loss.

Parameters:

Name	Type	Description	Default
x	Tensor	description	required
y_pred	Tensor	description	required
y_true	Tensor	description	required
device	str	description. Defaults to "cpu".	'cpu'

Returns:

Type	Description
Tensor	torch.Tensor: description

Source code in gresit/losses.py

def loss_hsic(
    x: torch.Tensor,
    y_pred: torch.Tensor,
    y_true: torch.Tensor,
    device: str = "cpu",
) -> torch.Tensor:
    """Get HSIC loss.

    Args:
        x (torch.Tensor): _description_
        y_pred (torch.Tensor): _description_
        y_true (torch.Tensor): _description_
        device (str, optional): _description_. Defaults to "cpu".

    Returns:
        torch.Tensor: _description_
    """
    r = y_pred - y_true
    return HSIC(x=x, y=r, device=device)

loss_mse(x, y_pred, y_true, device='cpu')

MSE loss.

Returns:

Type	Description
Tensor	torch.Tensor: description

Source code in gresit/losses.py

def loss_mse(
    x: torch.Tensor,
    y_pred: torch.Tensor,
    y_true: torch.Tensor,
    device: str = "cpu",
) -> torch.Tensor:
    """MSE loss.

    Returns:
        torch.Tensor: _description_
    """
    _ = x  # Ignored on purpose, kept for compatibility
    return ((y_pred - y_true) ** 2).mean()

pairwise_distances(x)

Get pairwise distance.

Parameters:

Name	Type	Description	Default
x	Tensor	description	required

Returns:

Type	Description
Tensor	torch.Tensor: description

Source code in gresit/losses.py

def pairwise_distances(x: torch.Tensor) -> torch.Tensor:
    """Get pairwise distance.

    Args:
        x (torch.Tensor): _description_

    Returns:
        torch.Tensor: _description_
    """
    instances_norm = torch.sum(x**2, -1).reshape((-1, 1))
    return -2 * torch.mm(x, x.t()) + instances_norm + instances_norm.t()

u_distance_cov_squared(a, b, device='cpu', robust=False)

Unbiased squared distance covariance.

Parameters:

Name	Type	Description	Default
a	Tensor	description	required
b	Tensor	description	required
device	str	device to run on.	'cpu'
robust	bool	if you want to make sure that the joint finite first moment condition for the distance covariance holds, setting robust to True transforms data to univariate ranks.	False

Returns:

Type	Description
Tensor	torch.Tensor: description

Source code in gresit/losses.py

def u_distance_cov_squared(
    a: torch.Tensor, b: torch.Tensor, device: str = "cpu", robust: bool = False
) -> torch.Tensor:
    """Unbiased squared distance covariance.

    Args:
        a (torch.Tensor): _description_
        b (torch.Tensor): _description_
        device (str): device to run on.
        robust (bool): if you want to make sure that the joint finite first moment
            condition for the distance covariance holds, setting `robust` to `True`
            transforms data to univariate ranks.

    Returns:
        torch.Tensor: _description_
    """
    if robust:
        a = _univariate_ranks(a)
        b = _univariate_ranks(b)
    n = a.shape[0]
    a_distance = _pairwise_distance_matrix(a)
    b_distance = _pairwise_distance_matrix(b)

    a_ij = _u_centered_matrix(a_distance)
    b_ij = _u_centered_matrix(b_distance)

    return (a_ij * b_ij).sum(dtype=torch.float32) / (n * (n - 3))

Utility classes and functions related to gresit.

Copyright (c) 2025 Robert Bosch GmbH

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.

BoostedRegressionTrees

Bases: MultiRegressor

Boosted multi-outcome regression trees.

This is simply a wrapper around the xgboost XGBRegressor class.

Source code in gresit/regression_techniques.py

class BoostedRegressionTrees(MultiRegressor):
    """Boosted multi-outcome regression trees.

    This is simply a wrapper around the xgboost `XGBRegressor` class.
    """

    def __init__(self) -> None:
        """Initializes the `XGBRegressor` object."""
        self.clf = xgb.XGBRegressor(tree_method="hist", multi_strategy="multi_output_tree")

    def fit(self, X: np.ndarray, Y: np.ndarray) -> None:
        """Fit boosted multi-outcome regression trees to training data.

        Args:
            X (np.ndarray): training data predictors
            Y (np.ndarray): training data (multi-) targets
        """
        self.clf.fit(X, Y)
        self._X_test = X

    def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
        """Predict using fitted boosted multi-outcome regression trees.

        Args:
            X_test (np.ndarray): Test data to predict on.

        Returns:
            np.ndarray: Predicted values.
        """
        if X_test is None:
            X_test = self._X_test
        return self.clf.predict(X_test)

__init__()

Initializes the XGBRegressor object.

Source code in gresit/regression_techniques.py

def __init__(self) -> None:
    """Initializes the `XGBRegressor` object."""
    self.clf = xgb.XGBRegressor(tree_method="hist", multi_strategy="multi_output_tree")

fit(X, Y)

Fit boosted multi-outcome regression trees to training data.

Parameters:

Name	Type	Description	Default
X	ndarray	training data predictors	required
Y	ndarray	training data (multi-) targets	required

Source code in gresit/regression_techniques.py

def fit(self, X: np.ndarray, Y: np.ndarray) -> None:
    """Fit boosted multi-outcome regression trees to training data.

    Args:
        X (np.ndarray): training data predictors
        Y (np.ndarray): training data (multi-) targets
    """
    self.clf.fit(X, Y)
    self._X_test = X

predict(X_test=None)

Predict using fitted boosted multi-outcome regression trees.

Parameters:

Name	Type	Description	Default
X_test	ndarray	Test data to predict on.	None

Returns:

Type	Description
ndarray	np.ndarray: Predicted values.

Source code in gresit/regression_techniques.py

def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
    """Predict using fitted boosted multi-outcome regression trees.

    Args:
        X_test (np.ndarray): Test data to predict on.

    Returns:
        np.ndarray: Predicted values.
    """
    if X_test is None:
        X_test = self._X_test
    return self.clf.predict(X_test)

CurdsWhey

Bases: MultiRegressor

Breiman and Friedman's curds and whey multivariate regression model.

When regression problem is of multivariate nature and the outcome variables are related among another, more accurate predictions may be obtained by using a linear combination of the OLS predictors.

Source code in gresit/regression_techniques.py

class CurdsWhey(MultiRegressor):
    """Breiman and Friedman's curds and whey multivariate regression model.

    When regression problem is of multivariate nature and the outcome variables are
    related among another, more accurate predictions may be obtained by using a linear combination
    of the OLS predictors.
    """

    def __init__(self) -> None:
        """Initializes C&W linear shrinkage method."""
        self._ols: sklearn.linear_model.LinearRegression
        self._T: np.ndarray
        self._D: np.ndarray
        self._response_transformed: StandardScaler

    def fit(self, X: np.ndarray, Y: np.ndarray) -> None:
        """Fits the C&W model to the multivariate X and Y.

        Args:
            X (np.ndarray): Matrix of predictors
            Y (np.ndarray): Matrix of responses
        """
        n = X.shape[0]
        p = X.shape[1]
        r = p / n
        q = Y.shape[1]

        Y_transform = StandardScaler().fit(Y)
        Y_std = Y_transform.transform(Y)

        cca = CCA(n_components=q)
        cca.fit(X, Y_std)
        X_c, Y_c = cca.transform(X, Y_std)
        c_k = np.corrcoef(X_c.T, Y_c.T).diagonal(offset=q)
        T = cca.y_rotations_

        # Computing Diagonal Shrinkage
        denom_1 = (c_k**2) * ((1 - r) ** 2)
        denom_2 = (1 - c_k**2) * (r**2)
        numer = (1 - r) * (c_k**2 - r)
        ds = numer / (denom_1 + denom_2)
        ds[ds < 0] = 0
        D = np.diag(ds)

        # Predicting values
        ols_cw = LinearRegression().fit(X, Y_std)
        self._ols = ols_cw
        self._T = T
        self._D = D
        self._response_transformed = Y_transform
        self._X_test = X

    def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
        """Predict response matrix based on X data.

        Args:
            X_test (np.ndarray): Matrix of responses

        Returns:
            np.ndarray: predicted values.
        """
        if X_test is None:
            X_test = self._X_test
        Y_tilde = self._ols.predict(X=X_test)
        return self._response_transformed.inverse_transform(
            Y_tilde @ self._T @ self._D @ np.linalg.inv(self._T)
        )

__init__()

Initializes C&W linear shrinkage method.

Source code in gresit/regression_techniques.py

def __init__(self) -> None:
    """Initializes C&W linear shrinkage method."""
    self._ols: sklearn.linear_model.LinearRegression
    self._T: np.ndarray
    self._D: np.ndarray
    self._response_transformed: StandardScaler

fit(X, Y)

Fits the C&W model to the multivariate X and Y.

Parameters:

Name	Type	Description	Default
X	ndarray	Matrix of predictors	required
Y	ndarray	Matrix of responses	required

Source code in gresit/regression_techniques.py

def fit(self, X: np.ndarray, Y: np.ndarray) -> None:
    """Fits the C&W model to the multivariate X and Y.

    Args:
        X (np.ndarray): Matrix of predictors
        Y (np.ndarray): Matrix of responses
    """
    n = X.shape[0]
    p = X.shape[1]
    r = p / n
    q = Y.shape[1]

    Y_transform = StandardScaler().fit(Y)
    Y_std = Y_transform.transform(Y)

    cca = CCA(n_components=q)
    cca.fit(X, Y_std)
    X_c, Y_c = cca.transform(X, Y_std)
    c_k = np.corrcoef(X_c.T, Y_c.T).diagonal(offset=q)
    T = cca.y_rotations_

    # Computing Diagonal Shrinkage
    denom_1 = (c_k**2) * ((1 - r) ** 2)
    denom_2 = (1 - c_k**2) * (r**2)
    numer = (1 - r) * (c_k**2 - r)
    ds = numer / (denom_1 + denom_2)
    ds[ds < 0] = 0
    D = np.diag(ds)

    # Predicting values
    ols_cw = LinearRegression().fit(X, Y_std)
    self._ols = ols_cw
    self._T = T
    self._D = D
    self._response_transformed = Y_transform
    self._X_test = X

predict(X_test=None)

Predict response matrix based on X data.

Parameters:

Name	Type	Description	Default
X_test	ndarray	Matrix of responses	None

Returns:

Type	Description
ndarray	np.ndarray: predicted values.

Source code in gresit/regression_techniques.py

def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
    """Predict response matrix based on X data.

    Args:
        X_test (np.ndarray): Matrix of responses

    Returns:
        np.ndarray: predicted values.
    """
    if X_test is None:
        X_test = self._X_test
    Y_tilde = self._ols.predict(X=X_test)
    return self._response_transformed.inverse_transform(
        Y_tilde @ self._T @ self._D @ np.linalg.inv(self._T)
    )

MultiRegressor

Abstract class for multivariate regression.

Source code in gresit/regression_techniques.py

class MultiRegressor(metaclass=ABCMeta):
    """Abstract class for multivariate regression."""

    def __init__(
        self,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
        **kwargs: dict[str, Any],
    ) -> None:
        """Base class for regressors."""
        self.rng = rng

    @abstractmethod
    def fit(self, X: np.ndarray, Y: np.ndarray) -> None:
        """Fit the model.

        Args:
            X (np.ndarray): Matrix of predictors
            Y (np.ndarray): Matrix of responses
        """

    @abstractmethod
    def predict(self, X_test: np.ndarray) -> np.ndarray:
        """Predict given data matrix.

        Args:
            X_test (np.ndarray): data matrix to predict on.

        Returns:
            np.ndarray: (Matrix of) predicted values
        """

    def mse(self, Y_test: np.ndarray, X_test: np.ndarray) -> float:
        """Mean squared error.

        Args:
            Y_test (np.ndarray): Test response
            X_test (np.ndarray): Test predictors

        Returns:
            float: MSE
        """
        Yhat = self.predict(X_test)
        mse: float = (np.square(Y_test - Yhat) / np.prod(Y_test.shape)).mean()
        return mse

    def standardize(self, a: np.ndarray) -> np.ndarray:
        """Standardize data.

        Args:
            a (np.ndarray): _description_

        Returns:
            np.ndarray: _description_
        """
        return (a - np.mean(a, axis=0)) / np.std(a, axis=0)

    def split_and_standardize(
        self,
        X: np.ndarray,
        Y: np.ndarray,
        test_size: float = 0.2,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """Train test split and standardize data.

        Args:
            X (np.ndarray): predictors
            Y (np.ndarray): targets
            test_size (float): Size of test data

        Returns:
            tuple[np.ndarray]: X_train, X_test, Y_train, Y_test
        """
        X_train, X_test, Y_train, Y_test = train_test_split(
            X, Y, test_size=test_size, random_state=2024
        )

        mean = np.mean(X_train, axis=0)
        std = np.std(X_train, axis=0)
        X_train = (X_train - mean) / std
        X_test = (X_test - mean) / std

        return X_train, X_test, Y_train, Y_test

__init__(rng=np.random.default_rng(seed=2024), **kwargs)

Base class for regressors.

Source code in gresit/regression_techniques.py

def __init__(
    self,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
    **kwargs: dict[str, Any],
) -> None:
    """Base class for regressors."""
    self.rng = rng

fit(X, Y) abstractmethod

Fit the model.

Parameters:

Name	Type	Description	Default
X	ndarray	Matrix of predictors	required
Y	ndarray	Matrix of responses	required

Source code in gresit/regression_techniques.py

@abstractmethod
def fit(self, X: np.ndarray, Y: np.ndarray) -> None:
    """Fit the model.

    Args:
        X (np.ndarray): Matrix of predictors
        Y (np.ndarray): Matrix of responses
    """

mse(Y_test, X_test)

Mean squared error.

Parameters:

Name	Type	Description	Default
Y_test	ndarray	Test response	required
X_test	ndarray	Test predictors	required

Returns:

Name	Type	Description
float	float	MSE

Source code in gresit/regression_techniques.py

def mse(self, Y_test: np.ndarray, X_test: np.ndarray) -> float:
    """Mean squared error.

    Args:
        Y_test (np.ndarray): Test response
        X_test (np.ndarray): Test predictors

    Returns:
        float: MSE
    """
    Yhat = self.predict(X_test)
    mse: float = (np.square(Y_test - Yhat) / np.prod(Y_test.shape)).mean()
    return mse

predict(X_test) abstractmethod

Predict given data matrix.

Parameters:

Name	Type	Description	Default
X_test	ndarray	data matrix to predict on.	required

Returns:

Type	Description
ndarray	np.ndarray: (Matrix of) predicted values

Source code in gresit/regression_techniques.py

@abstractmethod
def predict(self, X_test: np.ndarray) -> np.ndarray:
    """Predict given data matrix.

    Args:
        X_test (np.ndarray): data matrix to predict on.

    Returns:
        np.ndarray: (Matrix of) predicted values
    """

split_and_standardize(X, Y, test_size=0.2)

Train test split and standardize data.

Parameters:

Name	Type	Description	Default
X	ndarray	predictors	required
Y	ndarray	targets	required
test_size	float	Size of test data	0.2

Returns:

Type	Description
tuple[ndarray, ndarray, ndarray, ndarray]	tuple[np.ndarray]: X_train, X_test, Y_train, Y_test

Source code in gresit/regression_techniques.py

def split_and_standardize(
    self,
    X: np.ndarray,
    Y: np.ndarray,
    test_size: float = 0.2,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Train test split and standardize data.

    Args:
        X (np.ndarray): predictors
        Y (np.ndarray): targets
        test_size (float): Size of test data

    Returns:
        tuple[np.ndarray]: X_train, X_test, Y_train, Y_test
    """
    X_train, X_test, Y_train, Y_test = train_test_split(
        X, Y, test_size=test_size, random_state=2024
    )

    mean = np.mean(X_train, axis=0)
    std = np.std(X_train, axis=0)
    X_train = (X_train - mean) / std
    X_test = (X_test - mean) / std

    return X_train, X_test, Y_train, Y_test

standardize(a)

Standardize data.

Parameters:

Name	Type	Description	Default
a	ndarray	description	required

Returns:

Type	Description
ndarray	np.ndarray: description

Source code in gresit/regression_techniques.py

def standardize(self, a: np.ndarray) -> np.ndarray:
    """Standardize data.

    Args:
        a (np.ndarray): _description_

    Returns:
        np.ndarray: _description_
    """
    return (a - np.mean(a, axis=0)) / np.std(a, axis=0)

ReducedRankRegressor

Bases: MultiRegressor

Kernel Reduced Rank Ridge Regression.

Source code in gresit/regression_techniques.py

class ReducedRankRegressor(MultiRegressor):  # , BaseEstimator):
    """Kernel Reduced Rank Ridge Regression."""

    def __init__(self, rank: int, alpha: np.float64 = 1.0) -> None:
        """Initializes the model.

        Args:
            rank (int): Rank constraint.
            alpha (np.float64, optional): Regularization parameter. Defaults to 1.0.
        """
        self.rank = rank
        self.alpha = alpha
        self._P_rr: np.ndarray
        self._Q_fr: np.ndarray
        self._X_train: np.ndarray
        self._X_test: np.ndarray
        self._Y_test: np.ndarray | None = None

    def __str__(self) -> str:
        """Method print."""
        return f"kernel Reduced Rank Ridge Regression \
            by Mukherjee (rank = f{self.rank})"

    def fit(self, X: np.ndarray, Y: np.ndarray) -> None:
        """Fit kRRR model to data.

        Args:
            X (np.ndarray): training data predictors
            Y (np.ndarray): training (multivariate) response
        """
        K_X: np.ndarray = np.dot(X, X.T)
        tmp_1 = self.alpha * np.identity(K_X.shape[0]) + K_X
        Q_fr = np.linalg.solve(tmp_1, Y)
        P_fr = np.linalg.eig(np.dot(Y.T, np.dot(K_X, Q_fr)))[1].real
        P_rr = np.dot(P_fr[:, 0 : self.rank], P_fr[:, 0 : self.rank].T)

        self._Q_fr = Q_fr
        self._P_rr = P_rr
        self._X_train = X
        self._X_test = X

    def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
        """Predict fitted kRRR model.

        Args:
            X_test (np.ndarray): Test data to predict on.

        Returns:
            np.ndarray: Predicted values.
        """
        if X_test is None:
            X_test = self._X_test
        K_Xx = np.dot(X_test, self._X_train.T)

        return np.dot(K_Xx, np.dot(self._Q_fr, self._P_rr))

__init__(rank, alpha=1.0)

Initializes the model.

Parameters:

Name	Type	Description	Default
rank	int	Rank constraint.	required
alpha	float64	Regularization parameter. Defaults to 1.0.	1.0

Source code in gresit/regression_techniques.py

def __init__(self, rank: int, alpha: np.float64 = 1.0) -> None:
    """Initializes the model.

    Args:
        rank (int): Rank constraint.
        alpha (np.float64, optional): Regularization parameter. Defaults to 1.0.
    """
    self.rank = rank
    self.alpha = alpha
    self._P_rr: np.ndarray
    self._Q_fr: np.ndarray
    self._X_train: np.ndarray
    self._X_test: np.ndarray
    self._Y_test: np.ndarray | None = None

__str__()

Method print.

Source code in gresit/regression_techniques.py

def __str__(self) -> str:
    """Method print."""
    return f"kernel Reduced Rank Ridge Regression \
        by Mukherjee (rank = f{self.rank})"

fit(X, Y)

Fit kRRR model to data.

Parameters:

Name	Type	Description	Default
X	ndarray	training data predictors	required
Y	ndarray	training (multivariate) response	required

Source code in gresit/regression_techniques.py

def fit(self, X: np.ndarray, Y: np.ndarray) -> None:
    """Fit kRRR model to data.

    Args:
        X (np.ndarray): training data predictors
        Y (np.ndarray): training (multivariate) response
    """
    K_X: np.ndarray = np.dot(X, X.T)
    tmp_1 = self.alpha * np.identity(K_X.shape[0]) + K_X
    Q_fr = np.linalg.solve(tmp_1, Y)
    P_fr = np.linalg.eig(np.dot(Y.T, np.dot(K_X, Q_fr)))[1].real
    P_rr = np.dot(P_fr[:, 0 : self.rank], P_fr[:, 0 : self.rank].T)

    self._Q_fr = Q_fr
    self._P_rr = P_rr
    self._X_train = X
    self._X_test = X

predict(X_test=None)

Predict fitted kRRR model.

Parameters:

Name	Type	Description	Default
X_test	ndarray	Test data to predict on.	None

Returns:

Type	Description
ndarray	np.ndarray: Predicted values.

Source code in gresit/regression_techniques.py

def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
    """Predict fitted kRRR model.

    Args:
        X_test (np.ndarray): Test data to predict on.

    Returns:
        np.ndarray: Predicted values.
    """
    if X_test is None:
        X_test = self._X_test
    K_Xx = np.dot(X_test, self._X_train.T)

    return np.dot(K_Xx, np.dot(self._Q_fr, self._P_rr))

SimultaneousLinearModel

Bases: MultiRegressor

Class for performing multivariate linear Regression.

Source code in gresit/regression_techniques.py

class SimultaneousLinearModel(MultiRegressor):
    """Class for performing multivariate linear Regression."""

    def __init__(
        self,
        rng: np.random.Generator = np.random.default_rng(seed=2024),
        alpha: float = 0.1,
    ) -> None:
        """Initializes with a ridge penalty equal to 0.1.

        Args:
            rng (np.random.Generator): A random generator
            alpha (float, optional): Penalty term. Defaults to 0.1.
        """
        super().__init__(rng)

        self.alpha: float = alpha
        self.reg: Ridge

    def fit(self, X: np.ndarray, Y: np.ndarray) -> None:
        """Fit multivariate linear regression.

        Args:
            X (np.ndarray): Matrix of predictors
            Y (np.ndarray): Matrix of responses
        """
        self._X_test: np.ndarray = X
        self.reg = Ridge(alpha=self.alpha).fit(X, Y)

    def __str__(self) -> str:
        """Class description."""
        return "Multivariate Linear Regression"

    def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
        """Predict given data matrix.

        Args:
            X_test (np.ndarray): data matrix to predict on.

        Returns:
            np.ndarray: (Matrix of) predicted values
        """
        if X_test is None:
            X_test = self._X_test
        return self.reg.predict(X_test)

__init__(rng=np.random.default_rng(seed=2024), alpha=0.1)

Initializes with a ridge penalty equal to 0.1.

Parameters:

Name	Type	Description	Default
rng	Generator	A random generator	default_rng(seed=2024)
alpha	float	Penalty term. Defaults to 0.1.	0.1

Source code in gresit/regression_techniques.py

def __init__(
    self,
    rng: np.random.Generator = np.random.default_rng(seed=2024),
    alpha: float = 0.1,
) -> None:
    """Initializes with a ridge penalty equal to 0.1.

    Args:
        rng (np.random.Generator): A random generator
        alpha (float, optional): Penalty term. Defaults to 0.1.
    """
    super().__init__(rng)

    self.alpha: float = alpha
    self.reg: Ridge

__str__()

Class description.

Source code in gresit/regression_techniques.py

def __str__(self) -> str:
    """Class description."""
    return "Multivariate Linear Regression"

fit(X, Y)

Fit multivariate linear regression.

Parameters:

Name	Type	Description	Default
X	ndarray	Matrix of predictors	required
Y	ndarray	Matrix of responses	required

Source code in gresit/regression_techniques.py

def fit(self, X: np.ndarray, Y: np.ndarray) -> None:
    """Fit multivariate linear regression.

    Args:
        X (np.ndarray): Matrix of predictors
        Y (np.ndarray): Matrix of responses
    """
    self._X_test: np.ndarray = X
    self.reg = Ridge(alpha=self.alpha).fit(X, Y)

predict(X_test=None)

Predict given data matrix.

Parameters:

Name	Type	Description	Default
X_test	ndarray	data matrix to predict on.	None

Returns:

Type	Description
ndarray	np.ndarray: (Matrix of) predicted values

Source code in gresit/regression_techniques.py

def predict(self, X_test: np.ndarray | None = None) -> np.ndarray:
    """Predict given data matrix.

    Args:
        X_test (np.ndarray): data matrix to predict on.

    Returns:
        np.ndarray: (Matrix of) predicted values
    """
    if X_test is None:
        X_test = self._X_test
    return self.reg.predict(X_test)

Utility classes and functions related to gresit.

Copyright (c) 2025 Robert Bosch GmbH

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program. If not, see https://www.gnu.org/licenses/.

DAG

Bases: GRAPH

General class for dealing with directed acyclic graph i.e.

graphs that are directed and must not contain any cycles.

Source code in gresit/graphs.py

class DAG(GRAPH):
    """General class for dealing with directed acyclic graph i.e.

    graphs that are directed and must not contain any cycles.
    """

    def __init__(
        self,
        nodes: list[str] | None = None,
        edges: list[tuple[str, str]] | None = None,
    ) -> None:
        """DAG constructor.

        Args:
            nodes (list[str] | None, optional): Nodes. Defaults to None.
            edges (list[tuple[str,str]] | None, optional): Edges. Defaults to None.
        """
        if nodes is None:
            nodes = []
        if edges is None:
            edges = []

        self._nodes: set[str] = set(nodes)
        self._edges: set[tuple[str, str]] = set()
        self._parents: defaultdict[str, set[str]] = defaultdict(set)
        self._children: defaultdict[str, set[str]] = defaultdict(set)
        self._random_state: np.random.Generator = np.random.default_rng(seed=2023)

        for edge in edges:
            self._add_edge(*edge)

    def _add_node(self, node: str) -> None:
        self._nodes.add(node)

    def _add_edge(self, i: str, j: str) -> None:
        self._nodes.add(i)
        self._nodes.add(j)
        self._edges.add((i, j))

        # Check if graph is acyclic
        if not self.is_acyclic():
            raise ValueError(
                "The edge set you provided \
                induces one or more cycles.\
                Check your input!"
            )

        self._children[i].add(j)
        self._parents[j].add(i)

    @property
    def random_state(self) -> np.random.Generator:
        """Current random state.

        Returns:
            np.random.Generator: Generator object.
        """
        return self._random_state

    @random_state.setter
    def random_state(self, r: np.random.Generator) -> None:
        if not isinstance(r, np.random.Generator):
            raise AssertionError("Specify numpy random number generator object!")
        self._random_state = r

    def add_edge(self, edge: tuple[str, str]) -> None:
        """Add edge to DAG.

        Args:
            edge (tuple[str, str]): Edge to add
        """
        self._add_edge(*edge)

    def add_node(self, node: str) -> None:
        """Add node to DAG.

        Args:
            node (str): node to add
        """
        self._add_node(node)

    def add_edges_from(self, edges: list[tuple[str, str]]) -> None:
        """Add multiple edges to DAG.

        Args:
            edges (list[tuple[str, str]]): Edges to add
        """
        for edge in edges:
            self.add_edge(edge=edge)

    def add_nodes_from(self, nodes: list[str]) -> None:
        """Add multiple nodes to DAG.

        Args:
            nodes (list[str]): nodes to add
        """
        for node in nodes:
            self.add_node(node)

    def children(self, of_node: str) -> list[str]:
        """Gives all children of node `node`.

        Args:
            of_node (str): node in current DAG.

        Returns:
            list: of children.
        """
        if of_node in self._children.keys():
            return list(self._children[of_node])
        else:
            return []

    def parents(self, of_node: str) -> list[str]:
        """Gives all parents of node `node`.

        Args:
            of_node (str): node in current DAG.

        Returns:
            list: of parents.
        """
        if of_node in self._parents.keys():
            return list(self._parents[of_node])
        else:
            return []

    def induced_subgraph(self, nodes: list[str]) -> DAG:
        """Returns the induced subgraph on the nodes in `nodes`.

        Args:
            nodes (list[str]): List of nodes.

        Returns:
            DAG: Induced subgraph.
        """
        edges = [(i, j) for i, j in self.edges if i in nodes and j in nodes]
        return DAG(nodes=nodes, edges=edges)

    def is_adjacent(self, i: str, j: str) -> bool:
        """Return True if the graph contains an directed edge between i and j.

        Args:
            i (str): node i.
            j (str): node j.

        Returns:
            bool: True if i->j or i<-j
        """
        return (j, i) in self.edges or (i, j) in self.edges

    def is_clique(self, potential_clique: set[str]) -> bool:
        """Check every pair of node X potential_clique is adjacent."""
        return all(self.is_adjacent(i, j) for i, j in combinations(potential_clique, 2))

    def is_acyclic(self) -> bool:
        """Check if the graph is acyclic.

        Returns:
            bool: True if graph is acyclic.
        """
        nx_dag = self.to_networkx()
        acyclic: bool = nx.is_directed_acyclic_graph(nx_dag)
        return acyclic

    @classmethod
    def from_pandas_adjacency(cls, pd_amat: pd.DataFrame, *args: Any, **kwargs: Any) -> DAG:
        """Build DAG from a Pandas adjacency matrix.

        Args:
            pd_amat (pd.DataFrame): input adjacency matrix.
            args (Any): Additional arguments.
            kwargs (Any): Additional arguments.

        Returns:
            DAG
        """
        assert pd_amat.shape[0] == pd_amat.shape[1]
        nodes = pd_amat.columns

        all_connections = []
        start, end = np.where(pd_amat != 0)
        for idx, _ in enumerate(start):
            all_connections.append((pd_amat.columns[start[idx]], pd_amat.columns[end[idx]]))

        temp = [set(i) for i in all_connections]
        temp2 = [arc for arc in all_connections if temp.count(set(arc)) > 1]

        dir_edges = [edge for edge in all_connections if edge not in temp2]

        return DAG(nodes=nodes, edges=dir_edges)

    def remove_edge(self, i: str, j: str) -> None:
        """Removes edge in question.

        Args:
            i (str): tail
            j (str): head

        Raises:
            AssertionError: if edge does not exist
        """
        if (i, j) not in self.edges:
            raise AssertionError("Edge does not exist in current DAG")

        self._edges.discard((i, j))
        self._children[i].discard(j)
        self._parents[j].discard(i)

    def remove_node(self, node: str) -> None:
        """Remove a node from the graph."""
        self._nodes.remove(node)

        self._edges = {(i, j) for i, j in self._edges if node not in (i, j)}

        for child in self._children[node]:
            self._parents[child].remove(node)

        for parent in self._parents[node]:
            self._children[parent].remove(node)

        self._parents.pop(node, "I was never here")
        self._children.pop(node, "I was never here")

    @property
    def adjacency_matrix(self) -> pd.DataFrame:
        """Returns adjacency matrix.

        The i,jth entry being one indicates that there is an edge
        from i to j. A zero indicates that there is no edge.

        Returns:
            pd.DataFrame: adjacency matrix
        """
        amat = pd.DataFrame(
            np.zeros([self.num_nodes, self.num_nodes]),
            index=self.nodes,
            columns=self.nodes,
        )
        for edge in self.edges:
            amat.loc[edge] = 1
        return amat

    def vstructs(self) -> set[tuple[str, str]]:
        """Retrieve v-structures.

        Returns:
            set: set of all v-structures
        """
        vstructures = set()
        for node in self._nodes:
            for p1, p2 in combinations(self._parents[node], 2):
                if p1 not in self._parents[p2] and p2 not in self._parents[p1]:
                    vstructures.add((p1, node))
                    vstructures.add((p2, node))
        return vstructures

    def copy(self) -> DAG:
        """Return a copy of the graph."""
        return DAG(nodes=list(self._nodes), edges=list(self._edges))

    def show(self) -> None:
        """Plot DAG."""
        graph = self.to_networkx()
        pos = nx.circular_layout(graph)
        nx.draw(graph, pos=pos, with_labels=True)

    def to_networkx(self) -> nx.DiGraph:
        """Convert to networkx graph.

        Returns:
            nx.MultiDiGraph: Graph with directed and undirected edges.
        """
        nx_dag = nx.DiGraph()
        nx_dag.add_nodes_from(self.nodes)
        nx_dag.add_edges_from(self.edges)

        return nx_dag

    @property
    def nodes(self) -> list[str]:
        """Get all nods in current DAG.

        Returns:
            list: list of nodes.
        """
        return sorted(list(self._nodes))

    @property
    def num_nodes(self) -> int:
        """Number of nodes in current DAG.

        Returns:
            int: Number of nodes
        """
        return len(self._nodes)

    @property
    def num_edges(self) -> int:
        """Number of directed edges in current DAG.

        Returns:
            int: Number of directed edges
        """
        return len(self._edges)

    @property
    def sparsity(self) -> float:
        """Sparsity of the graph.

        Returns:
            float: in [0,1]
        """
        s = self.num_nodes
        return self.num_edges / s / (s - 1) * 2

    @property
    def edges(self) -> list[tuple[str, str]]:
        """Gives all directed edges in current DAG.

        Returns:
            list[tuple[str,str]]: List of directed edges.
        """
        return list(self._edges)

    @property
    def causal_order(self) -> list[str]:
        """Returns the causal order of the current graph.

        Note that this order is in general not unique.

        Returns:
            list[str]: Causal order
        """
        return list(nx.lexicographical_topological_sort(self.to_networkx()))

    @property
    def sink_nodes(self) -> list[str]:
        """Returns all sink nodes, i.e.

        nodes with no descendents in particular no children.

        Returns:
            list[str]: list of sink nodes.
        """
        return [
            s
            for b, s in zip([self.children(of_node=node) == [] for node in self.nodes], self.nodes)
            if b
        ]

    @property
    def source_nodes(self) -> list[str]:
        """Returns all source nodes, i.e.

        nodes with no ancesters in particular no parents.

        Returns:
            list[str]: list of sink nodes.
        """
        return [
            s
            for b, s in zip([self.parents(of_node=node) == [] for node in self.nodes], self.nodes)
            if b
        ]

    @property
    def max_in_degree(self) -> int:
        """Maximum in-degree of the graph.

        Returns:
            int: Maximum in-degree
        """
        return max(len(self._parents[node]) for node in self._nodes)

    @property
    def max_out_degree(self) -> int:
        """Maximum out-degree of the graph.

        Returns:
            int: Maximum out-degree
        """
        return max(len(self._children[node]) for node in self._nodes)

    @classmethod
    def from_nx(cls, nx_dag: nx.DiGraph, *args: Any, **kwargs: Any) -> DAG:
        """Convert to DAG from nx.DiGraph.

        Args:
            nx_dag (nx.DiGraph): DAG in question.
            args (Any): additional arguments
            kwargs (Any): additional arguments

        Raises:
            TypeError: If DAG is not nx.DiGraph

        Returns:
            DAG
        """
        if not isinstance(nx_dag, nx.DiGraph):
            raise TypeError("DAG must be of type nx.DiGraph")
        return DAG(nodes=list(nx_dag.nodes), edges=list(nx_dag.edges))

    def to_cpdag(self) -> PDAG:
        """Convert DAG to CPDAG.

        Returns:
            PDAG: CPDAG representing the MEC.
        """
        return dag2cpdag(dag=self.to_networkx())

adjacency_matrix property

Returns adjacency matrix.

The i,jth entry being one indicates that there is an edge from i to j. A zero indicates that there is no edge.

Returns:

Type	Description
DataFrame	pd.DataFrame: adjacency matrix

causal_order property

Returns the causal order of the current graph.

Note that this order is in general not unique.

Returns:

Type	Description
list[str]	list[str]: Causal order

edges property

Gives all directed edges in current DAG.

Returns:

Type	Description
list[tuple[str, str]]	list[tuple[str,str]]: List of directed edges.

max_in_degree property

Maximum in-degree of the graph.

Returns:

Name	Type	Description
int	int	Maximum in-degree

max_out_degree property

Maximum out-degree of the graph.

Returns:

Name	Type	Description
int	int	Maximum out-degree

nodes property

Get all nods in current DAG.

Returns:

Name	Type	Description
list	list[str]	list of nodes.

num_edges property

Number of directed edges in current DAG.

Returns:

Name	Type	Description
int	int	Number of directed edges

num_nodes property

Number of nodes in current DAG.

Returns:

Name	Type	Description
int	int	Number of nodes

random_state property writable

Current random state.

Returns:

Type	Description
Generator	np.random.Generator: Generator object.

sink_nodes property

Returns all sink nodes, i.e.

nodes with no descendents in particular no children.

Returns:

Type	Description
list[str]	list[str]: list of sink nodes.

source_nodes property

Returns all source nodes, i.e.

nodes with no ancesters in particular no parents.

Returns:

Type	Description
list[str]	list[str]: list of sink nodes.

sparsity property

Sparsity of the graph.

Returns:

Name	Type	Description
float	float	in [0,1]

__init__(nodes=None, edges=None)

DAG constructor.

Parameters:

Name	Type	Description	Default
nodes	list[str] | None	Nodes. Defaults to None.	None
edges	list[tuple[str, str]] | None	Edges. Defaults to None.	None

Source code in gresit/graphs.py

def __init__(
    self,
    nodes: list[str] | None = None,
    edges: list[tuple[str, str]] | None = None,
) -> None:
    """DAG constructor.

    Args:
        nodes (list[str] | None, optional): Nodes. Defaults to None.
        edges (list[tuple[str,str]] | None, optional): Edges. Defaults to None.
    """
    if nodes is None:
        nodes = []
    if edges is None:
        edges = []

    self._nodes: set[str] = set(nodes)
    self._edges: set[tuple[str, str]] = set()
    self._parents: defaultdict[str, set[str]] = defaultdict(set)
    self._children: defaultdict[str, set[str]] = defaultdict(set)
    self._random_state: np.random.Generator = np.random.default_rng(seed=2023)

    for edge in edges:
        self._add_edge(*edge)

add_edge(edge)

Add edge to DAG.

Parameters:

Name	Type	Description	Default
edge	tuple[str, str]	Edge to add	required

Source code in gresit/graphs.py

def add_edge(self, edge: tuple[str, str]) -> None:
    """Add edge to DAG.

    Args:
        edge (tuple[str, str]): Edge to add
    """
    self._add_edge(*edge)

add_edges_from(edges)

Add multiple edges to DAG.

Parameters:

Name	Type	Description	Default
edges	list[tuple[str, str]]	Edges to add	required

Source code in gresit/graphs.py

def add_edges_from(self, edges: list[tuple[str, str]]) -> None:
    """Add multiple edges to DAG.

    Args:
        edges (list[tuple[str, str]]): Edges to add
    """
    for edge in edges:
        self.add_edge(edge=edge)

add_node(node)

Add node to DAG.

Parameters:

Name	Type	Description	Default
node	str	node to add	required

Source code in gresit/graphs.py

def add_node(self, node: str) -> None:
    """Add node to DAG.

    Args:
        node (str): node to add
    """
    self._add_node(node)

add_nodes_from(nodes)

Add multiple nodes to DAG.

Parameters:

Name	Type	Description	Default
nodes	list[str]	nodes to add	required

Source code in gresit/graphs.py

def add_nodes_from(self, nodes: list[str]) -> None:
    """Add multiple nodes to DAG.

    Args:
        nodes (list[str]): nodes to add
    """
    for node in nodes:
        self.add_node(node)

children(of_node)

Gives all children of node node.

Parameters:

Name	Type	Description	Default
of_node	str	node in current DAG.	required

Returns:

Name	Type	Description
list	list[str]	of children.

Source code in gresit/graphs.py

def children(self, of_node: str) -> list[str]:
    """Gives all children of node `node`.

    Args:
        of_node (str): node in current DAG.

    Returns:
        list: of children.
    """
    if of_node in self._children.keys():
        return list(self._children[of_node])
    else:
        return []

copy()

Return a copy of the graph.

Source code in gresit/graphs.py

def copy(self) -> DAG:
    """Return a copy of the graph."""
    return DAG(nodes=list(self._nodes), edges=list(self._edges))

from_nx(nx_dag, *args, **kwargs) classmethod

Convert to DAG from nx.DiGraph.

Parameters:

Name	Type	Description	Default
nx_dag	DiGraph	DAG in question.	required
args	Any	additional arguments	()
kwargs	Any	additional arguments	{}

Raises:

Type	Description
TypeError	If DAG is not nx.DiGraph

Returns:

Type	Description
DAG	DAG

Source code in gresit/graphs.py

@classmethod
def from_nx(cls, nx_dag: nx.DiGraph, *args: Any, **kwargs: Any) -> DAG:
    """Convert to DAG from nx.DiGraph.

    Args:
        nx_dag (nx.DiGraph): DAG in question.
        args (Any): additional arguments
        kwargs (Any): additional arguments

    Raises:
        TypeError: If DAG is not nx.DiGraph

    Returns:
        DAG
    """
    if not isinstance(nx_dag, nx.DiGraph):
        raise TypeError("DAG must be of type nx.DiGraph")
    return DAG(nodes=list(nx_dag.nodes), edges=list(nx_dag.edges))

from_pandas_adjacency(pd_amat, *args, **kwargs) classmethod

Build DAG from a Pandas adjacency matrix.

Parameters:

Name	Type	Description	Default
pd_amat	DataFrame	input adjacency matrix.	required
args	Any	Additional arguments.	()
kwargs	Any	Additional arguments.	{}

Returns:

Type	Description
DAG	DAG

Source code in gresit/graphs.py

@classmethod
def from_pandas_adjacency(cls, pd_amat: pd.DataFrame, *args: Any, **kwargs: Any) -> DAG:
    """Build DAG from a Pandas adjacency matrix.

    Args:
        pd_amat (pd.DataFrame): input adjacency matrix.
        args (Any): Additional arguments.
        kwargs (Any): Additional arguments.

    Returns:
        DAG
    """
    assert pd_amat.shape[0] == pd_amat.shape[1]
    nodes = pd_amat.columns

    all_connections = []
    start, end = np.where(pd_amat != 0)
    for idx, _ in enumerate(start):
        all_connections.append((pd_amat.columns[start[idx]], pd_amat.columns[end[idx]]))

    temp = [set(i) for i in all_connections]
    temp2 = [arc for arc in all_connections if temp.count(set(arc)) > 1]

    dir_edges = [edge for edge in all_connections if edge not in temp2]

    return DAG(nodes=nodes, edges=dir_edges)

induced_subgraph(nodes)

Returns the induced subgraph on the nodes in nodes.

Parameters:

Name	Type	Description	Default
nodes	list[str]	List of nodes.	required

Returns:

Name	Type	Description
DAG	DAG	Induced subgraph.

Source code in gresit/graphs.py

def induced_subgraph(self, nodes: list[str]) -> DAG:
    """Returns the induced subgraph on the nodes in `nodes`.

    Args:
        nodes (list[str]): List of nodes.

    Returns:
        DAG: Induced subgraph.
    """
    edges = [(i, j) for i, j in self.edges if i in nodes and j in nodes]
    return DAG(nodes=nodes, edges=edges)

is_acyclic()

Check if the graph is acyclic.

Returns:

Name	Type	Description
bool	bool	True if graph is acyclic.

Source code in gresit/graphs.py

def is_acyclic(self) -> bool:
    """Check if the graph is acyclic.

    Returns:
        bool: True if graph is acyclic.
    """
    nx_dag = self.to_networkx()
    acyclic: bool = nx.is_directed_acyclic_graph(nx_dag)
    return acyclic

is_adjacent(i, j)

Return True if the graph contains an directed edge between i and j.

Parameters:

Name	Type	Description	Default
i	str	node i.	required
j	str	node j.	required

Returns:

Name	Type	Description
bool	bool	True if i->j or i<-j

Source code in gresit/graphs.py

def is_adjacent(self, i: str, j: str) -> bool:
    """Return True if the graph contains an directed edge between i and j.

    Args:
        i (str): node i.
        j (str): node j.

    Returns:
        bool: True if i->j or i<-j
    """
    return (j, i) in self.edges or (i, j) in self.edges

is_clique(potential_clique)

Check every pair of node X potential_clique is adjacent.

Source code in gresit/graphs.py

def is_clique(self, potential_clique: set[str]) -> bool:
    """Check every pair of node X potential_clique is adjacent."""
    return all(self.is_adjacent(i, j) for i, j in combinations(potential_clique, 2))

parents(of_node)

Gives all parents of node node.

Parameters:

Name	Type	Description	Default
of_node	str	node in current DAG.	required

Returns:

Name	Type	Description
list	list[str]	of parents.

Source code in gresit/graphs.py

def parents(self, of_node: str) -> list[str]:
    """Gives all parents of node `node`.

    Args:
        of_node (str): node in current DAG.

    Returns:
        list: of parents.
    """
    if of_node in self._parents.keys():
        return list(self._parents[of_node])
    else:
        return []

remove_edge(i, j)

Removes edge in question.

Parameters:

Name	Type	Description	Default
i	str	tail	required
j	str	head	required

Raises:

Type	Description
AssertionError	if edge does not exist

Source code in gresit/graphs.py

def remove_edge(self, i: str, j: str) -> None:
    """Removes edge in question.

    Args:
        i (str): tail
        j (str): head

    Raises:
        AssertionError: if edge does not exist
    """
    if (i, j) not in self.edges:
        raise AssertionError("Edge does not exist in current DAG")

    self._edges.discard((i, j))
    self._children[i].discard(j)
    self._parents[j].discard(i)

remove_node(node)

Remove a node from the graph.

Source code in gresit/graphs.py

def remove_node(self, node: str) -> None:
    """Remove a node from the graph."""
    self._nodes.remove(node)

    self._edges = {(i, j) for i, j in self._edges if node not in (i, j)}

    for child in self._children[node]:
        self._parents[child].remove(node)

    for parent in self._parents[node]:
        self._children[parent].remove(node)

    self._parents.pop(node, "I was never here")
    self._children.pop(node, "I was never here")

show()

Plot DAG.

Source code in gresit/graphs.py

def show(self) -> None:
    """Plot DAG."""
    graph = self.to_networkx()
    pos = nx.circular_layout(graph)
    nx.draw(graph, pos=pos, with_labels=True)

to_cpdag()

Convert DAG to CPDAG.

Returns:

Name	Type	Description
PDAG	PDAG	CPDAG representing the MEC.

Source code in gresit/graphs.py

def to_cpdag(self) -> PDAG:
    """Convert DAG to CPDAG.

    Returns:
        PDAG: CPDAG representing the MEC.
    """
    return dag2cpdag(dag=self.to_networkx())

to_networkx()

Convert to networkx graph.

Returns:

Type	Description
DiGraph	nx.MultiDiGraph: Graph with directed and undirected edges.

Source code in gresit/graphs.py

def to_networkx(self) -> nx.DiGraph:
    """Convert to networkx graph.

    Returns:
        nx.MultiDiGraph: Graph with directed and undirected edges.
    """
    nx_dag = nx.DiGraph()
    nx_dag.add_nodes_from(self.nodes)
    nx_dag.add_edges_from(self.edges)

    return nx_dag

vstructs()

Retrieve v-structures.

Returns:

Name	Type	Description
set	set[tuple[str, str]]	set of all v-structures

Source code in gresit/graphs.py

def vstructs(self) -> set[tuple[str, str]]:
    """Retrieve v-structures.

    Returns:
        set: set of all v-structures
    """
    vstructures = set()
    for node in self._nodes:
        for p1, p2 in combinations(self._parents[node], 2):
            if p1 not in self._parents[p2] and p2 not in self._parents[p1]:
                vstructures.add((p1, node))
                vstructures.add((p2, node))
    return vstructures

GRAPH

Abstract base class for all Graphs in current project.

Source code in gresit/graphs.py

class GRAPH(metaclass=ABCMeta):
    """Abstract base class for all Graphs in current project."""

    def __init__(self) -> None:
        """Init ABC."""
        pass

    @property
    @abstractmethod
    def adjacency_matrix(self) -> pd.DataFrame:
        """Return adjacency matrix.

        Raises:
            AssertionError: _description_
            AssertionError: _description_
            ValueError: _description_
            AssertionError: _description_
            AssertionError: _description_
            TypeError: _description_

        Returns:
            pd.DataFrame: Adjacency matrix of underlying graph.
        """

    @property
    @abstractmethod
    def causal_order(self) -> list[str] | None:
        """Return causal order.

        Raises:
            AssertionError: _description_
            AssertionError: _description_
            ValueError: _description_
            AssertionError: _description_
            AssertionError: _description_
            TypeError: _description_

        Returns:
            list[str] | None: Causal order of underlying graph.
                None if not a DAG.
        """

adjacency_matrix abstractmethod property

Return adjacency matrix.

Raises:

Type	Description
AssertionError	description
AssertionError	description
ValueError	description
AssertionError	description
AssertionError	description
TypeError	description

Returns:

Type	Description
DataFrame	pd.DataFrame: Adjacency matrix of underlying graph.

causal_order abstractmethod property

Return causal order.

Raises:

Type	Description
AssertionError	description
AssertionError	description
ValueError	description
AssertionError	description
AssertionError	description
TypeError	description

Returns:

Type	Description
list[str] | None	list[str] | None: Causal order of underlying graph. None if not a DAG.

__init__()

Init ABC.

Source code in gresit/graphs.py

def __init__(self) -> None:
    """Init ABC."""
    pass

LayeredDAG

Bases: DAG

Class to construct Layered DAGs.

Layered DAGs L are DAGs where the Nodes V follow some natural layering. In other words, no edge can ever point into any of the earlier layers.

Source code in gresit/graphs.py

class LayeredDAG(DAG):
    """Class to construct Layered DAGs.

    Layered DAGs `L` are DAGs where the Nodes `V` follow some natural layering.
    In other words, no edge can ever point into any of the earlier layers.
    """

    def __init__(
        self,
        nodes: list[str] | None = None,
        edges: list[tuple[str, str]] | None = None,
        layering: dict[str, list[str]] | None = None,
    ) -> None:
        """Layered DAG constructor.

        Args:
            nodes (list[str] | None, optional): Nodes of LDAG. Defaults to None.
            edges (list[tuple[str, str]] | None, optional): Edges of LDAG. Defaults to None.
            layering (dict[str, list[str]] | None, optional): Layering. Defaults to None.
        """
        self._layering = layering
        super().__init__(nodes=nodes, edges=edges)

    def _add_edge(self, i: str, j: str) -> None:
        if not self.layering:
            raise ValueError("Layering must be provided before adding edges.")

        self._nodes.add(i)
        self._nodes.add(j)
        self._edges.add((i, j))

        # Check if graph is acyclic
        if not self.is_acyclic():
            raise ValueError(
                "The edge set you provided \
                induces one or more cycles.\
                Check your input!"
            )

        # Check if edge is allowed due to layering
        if not self._is_allowed(edge=(i, j)):
            raise ValueError(
                "The edge set you provided \
                does not agree with the layering.\
                Check your input!"
            )

        self._children[i].add(j)
        self._parents[j].add(i)

    @property
    def layering(self) -> dict[str, list[str]] | None:
        """Current layering dict.

        Returns:
            dict[str, list[str]]: Layering
        """
        return self._layering

    @layering.setter
    def layering(self, la: np.random.Generator) -> None:
        if not isinstance(la, dict):
            raise AssertionError("Layering must be a dictionary!")
        self._layering = la

    def _is_allowed(self, edge: tuple[str, str]) -> bool:
        if not self.layering:
            raise ValueError("Layering must be provided before adding edges.")
        i, j = edge
        layers = list(self.layering.keys())
        for layer, nodes in self.layering.items():
            if i in nodes:
                i_layer = layers.index(layer)
            if j in nodes:
                j_layer = layers.index(layer)
        return i_layer <= j_layer

    def layer_induced_subgraph(self, nodes: list[str]) -> DAG:
        """Returns the induced subgraph on the nodes in `nodes`.

        Args:
            nodes (list[str]): List of nodes.

        Returns:
            DAG: Induced subgraph.
        """
        if self.layering is not None and not any(
            [nodes == layer for layer in self.layering.values()]
        ):
            raise ValueError("Nodes you provide must correspond to a layer.")
        edges = [(i, j) for i, j in self.edges if i in nodes and j in nodes]
        return DAG(nodes=nodes, edges=edges)

    @classmethod
    def from_pandas_adjacency(
        cls, pd_amat: pd.DataFrame, layering: dict[str, list[str]]
    ) -> LayeredDAG:
        """Build LayeredDAG from a Pandas adjacency matrix.

        Args:
            pd_amat (pd.DataFrame): input adjacency matrix.
            layering (dict[str, list[str]]): layering of nodes.

        Returns:
            LayeredDAG
        """
        assert pd_amat.shape[0] == pd_amat.shape[1]
        nodes = pd_amat.columns

        all_connections = []
        start, end = np.where(pd_amat != 0)
        for idx, _ in enumerate(start):
            all_connections.append((pd_amat.columns[start[idx]], pd_amat.columns[end[idx]]))

        temp = [set(i) for i in all_connections]
        temp2 = [arc for arc in all_connections if temp.count(set(arc)) > 1]

        dir_edges = [edge for edge in all_connections if edge not in temp2]

        return LayeredDAG(nodes=nodes, edges=dir_edges, layering=layering)

    def copy(self) -> LayeredDAG:
        """Return a copy of the graph."""
        return LayeredDAG(nodes=list(self._nodes), edges=list(self._edges), layering=self.layering)

    @classmethod
    def from_nx(cls, nx_dag: nx.DiGraph, layering: dict[str, list[str]]) -> LayeredDAG:
        """Convert to DAG from nx.DiGraph.

        Args:
            nx_dag (nx.DiGraph): DAG in question.
            layering (dict[str, list[str]]): layering of nodes.

        Raises:
            TypeError: If DAG is not nx.DiGraph

        Returns:
            LayeredDAG
        """
        if not isinstance(nx_dag, nx.DiGraph):
            raise TypeError("DAG must be of type nx.DiGraph")
        return LayeredDAG(nodes=list(nx_dag.nodes), edges=list(nx_dag.edges), layering=layering)

layering property writable

Current layering dict.

Returns:

Type	Description
dict[str, list[str]] | None	dict[str, list[str]]: Layering

__init__(nodes=None, edges=None, layering=None)

Layered DAG constructor.

Parameters:

Name	Type	Description	Default
nodes	list[str] | None	Nodes of LDAG. Defaults to None.	None
edges	list[tuple[str, str]] | None	Edges of LDAG. Defaults to None.	None
layering	dict[str, list[str]] | None	Layering. Defaults to None.	None

Source code in gresit/graphs.py

def __init__(
    self,
    nodes: list[str] | None = None,
    edges: list[tuple[str, str]] | None = None,
    layering: dict[str, list[str]] | None = None,
) -> None:
    """Layered DAG constructor.

    Args:
        nodes (list[str] | None, optional): Nodes of LDAG. Defaults to None.
        edges (list[tuple[str, str]] | None, optional): Edges of LDAG. Defaults to None.
        layering (dict[str, list[str]] | None, optional): Layering. Defaults to None.
    """
    self._layering = layering
    super().__init__(nodes=nodes, edges=edges)

copy()

Return a copy of the graph.

Source code in gresit/graphs.py

def copy(self) -> LayeredDAG:
    """Return a copy of the graph."""
    return LayeredDAG(nodes=list(self._nodes), edges=list(self._edges), layering=self.layering)

from_nx(nx_dag, layering) classmethod

Convert to DAG from nx.DiGraph.

Parameters:

Name	Type	Description	Default
nx_dag	DiGraph	DAG in question.	required
layering	dict[str, list[str]]	layering of nodes.	required

Raises:

Type	Description
TypeError	If DAG is not nx.DiGraph

Returns:

Type	Description
LayeredDAG	LayeredDAG

Source code in gresit/graphs.py

@classmethod
def from_nx(cls, nx_dag: nx.DiGraph, layering: dict[str, list[str]]) -> LayeredDAG:
    """Convert to DAG from nx.DiGraph.

    Args:
        nx_dag (nx.DiGraph): DAG in question.
        layering (dict[str, list[str]]): layering of nodes.

    Raises:
        TypeError: If DAG is not nx.DiGraph

    Returns:
        LayeredDAG
    """
    if not isinstance(nx_dag, nx.DiGraph):
        raise TypeError("DAG must be of type nx.DiGraph")
    return LayeredDAG(nodes=list(nx_dag.nodes), edges=list(nx_dag.edges), layering=layering)

from_pandas_adjacency(pd_amat, layering) classmethod

Build LayeredDAG from a Pandas adjacency matrix.

Parameters:

Name	Type	Description	Default
pd_amat	DataFrame	input adjacency matrix.	required
layering	dict[str, list[str]]	layering of nodes.	required

Returns:

Type	Description
LayeredDAG	LayeredDAG

Source code in gresit/graphs.py

@classmethod
def from_pandas_adjacency(
    cls, pd_amat: pd.DataFrame, layering: dict[str, list[str]]
) -> LayeredDAG:
    """Build LayeredDAG from a Pandas adjacency matrix.

    Args:
        pd_amat (pd.DataFrame): input adjacency matrix.
        layering (dict[str, list[str]]): layering of nodes.

    Returns:
        LayeredDAG
    """
    assert pd_amat.shape[0] == pd_amat.shape[1]
    nodes = pd_amat.columns

    all_connections = []
    start, end = np.where(pd_amat != 0)
    for idx, _ in enumerate(start):
        all_connections.append((pd_amat.columns[start[idx]], pd_amat.columns[end[idx]]))

    temp = [set(i) for i in all_connections]
    temp2 = [arc for arc in all_connections if temp.count(set(arc)) > 1]

    dir_edges = [edge for edge in all_connections if edge not in temp2]

    return LayeredDAG(nodes=nodes, edges=dir_edges, layering=layering)

layer_induced_subgraph(nodes)

Returns the induced subgraph on the nodes in nodes.

Parameters:

Name	Type	Description	Default
nodes	list[str]	List of nodes.	required

Returns:

Name	Type	Description
DAG	DAG	Induced subgraph.

Source code in gresit/graphs.py

def layer_induced_subgraph(self, nodes: list[str]) -> DAG:
    """Returns the induced subgraph on the nodes in `nodes`.

    Args:
        nodes (list[str]): List of nodes.

    Returns:
        DAG: Induced subgraph.
    """
    if self.layering is not None and not any(
        [nodes == layer for layer in self.layering.values()]
    ):
        raise ValueError("Nodes you provide must correspond to a layer.")
    edges = [(i, j) for i, j in self.edges if i in nodes and j in nodes]
    return DAG(nodes=nodes, edges=edges)

PDAG

Bases: GRAPH

Class for dealing with partially directed graph i.e.

graphs that contain both directed and undirected edges.

Source code in gresit/graphs.py

class PDAG(GRAPH):
    """Class for dealing with partially directed graph i.e.

    graphs that contain both directed and undirected edges.
    """

    def __init__(
        self,
        nodes: list[str] | None = None,
        dir_edges: list[tuple[str, str]] | None = None,
        undir_edges: list[tuple[str, str]] | None = None,
    ) -> None:
        """PDAG constructor.

        Args:
            nodes (list[str] | None, optional): Nodes in the PDAG. Defaults to None.
            dir_edges (list[tuple[str,str]] | None, optional): directed edges. Defaults to None.
            undir_edges (list[tuple[str,str]] | None, optional): undirected edges. Defaults to None.
        """
        if nodes is None:
            nodes = []
        if dir_edges is None:
            dir_edges = []
        if undir_edges is None:
            undir_edges = []

        self._nodes = set(nodes)
        self._undir_edges: set[tuple[str, str]] = set()
        self._dir_edges: set[tuple[str, str]] = set()
        self._parents: defaultdict[str, set[str]] = defaultdict(set)
        self._children: defaultdict[str, set[str]] = defaultdict(set)
        self._neighbors: defaultdict[str, set[str]] = defaultdict(set)
        self._undirected_neighbors: defaultdict[str, set[str]] = defaultdict(set)

        for dir_edge in dir_edges:
            self._add_dir_edge(*dir_edge)
        for unir_edge in undir_edges:
            self._add_undir_edge(*unir_edge)

    def _add_dir_edge(self, i: str, j: str) -> None:
        self._nodes.add(i)
        self._nodes.add(j)
        self._dir_edges.add((i, j))

        self._neighbors[i].add(j)
        self._neighbors[j].add(i)

        self._children[i].add(j)
        self._parents[j].add(i)

    def _add_undir_edge(self, i: str, j: str) -> None:
        self._nodes.add(i)
        self._nodes.add(j)
        self._undir_edges.add((i, j))

        self._neighbors[i].add(j)
        self._neighbors[j].add(i)

        self._undirected_neighbors[i].add(j)
        self._undirected_neighbors[j].add(i)

    def children(self, node: str) -> set[str]:
        """Gives all children of node `node`.

        Args:
            node (str): node in current PDAG.

        Returns:
            set: set of children.
        """
        if node in self._children.keys():
            return self._children[node]
        else:
            return set()

    def parents(self, node: str) -> set[str]:
        """Gives all parents of node `node`.

        Args:
            node (str): node in current PDAG.

        Returns:
            set: set of parents.
        """
        if node in self._parents.keys():
            return self._parents[node]
        else:
            return set()

    def neighbors(self, node: str) -> set[str]:
        """Gives all neighbors of node `node`.

        Args:
            node (str): node in current PDAG.

        Returns:
            set: set of neighbors.
        """
        if node in self._neighbors.keys():
            return self._neighbors[node]
        else:
            return set()

    def undir_neighbors(self, node: str) -> set[str]:
        """Gives all undirected neighbors of node `node`.

        Args:
            node (str): node in current PDAG.

        Returns:
            set: set of undirected neighbors.
        """
        if node in self._undirected_neighbors.keys():
            return self._undirected_neighbors[node]
        else:
            return set()

    def is_adjacent(self, i: str, j: str) -> bool:
        """Return True if the graph contains an directed or undirected edge between i and j.

        Args:
            i (str): node i.
            j (str): node j.

        Returns:
            bool: True if i-j or i->j or i<-j
        """
        return any(
            (
                (j, i) in self.dir_edges or (j, i) in self.undir_edges,
                (i, j) in self.dir_edges or (i, j) in self.undir_edges,
            )
        )

    def is_clique(self, potential_clique: set[str]) -> bool:
        """Check every pair of node X potential_clique is adjacent."""
        return all(self.is_adjacent(i, j) for i, j in combinations(potential_clique, 2))

    @classmethod
    def from_pandas_adjacency(cls, pd_amat: pd.DataFrame) -> PDAG:
        """Build PDAG from a Pandas adjacency matrix.

        Args:
            pd_amat (pd.DataFrame): input adjacency matrix.

        Returns:
            PDAG
        """
        assert pd_amat.shape[0] == pd_amat.shape[1]
        nodes = pd_amat.columns

        all_connections = []
        start, end = np.where(pd_amat != 0)
        for idx, _ in enumerate(start):
            all_connections.append((pd_amat.columns[start[idx]], pd_amat.columns[end[idx]]))

        temp = [set(i) for i in all_connections]
        temp2 = [arc for arc in all_connections if temp.count(set(arc)) > 1]
        undir_edges = [tuple(item) for item in set(frozenset(item) for item in temp2)]

        dir_edges = [edge for edge in all_connections if edge not in temp2]

        return PDAG(nodes=nodes, dir_edges=dir_edges, undir_edges=undir_edges)

    def remove_edge(self, i: str, j: str) -> None:
        """Removes edge in question.

        Args:
            i (str): tail
            j (str): head

        Raises:
            AssertionError: if edge does not exist
        """
        if (i, j) not in self.dir_edges and (i, j) not in self.undir_edges:
            raise AssertionError("Edge does not exist in current PDAG")

        self._undir_edges.discard((i, j))
        self._dir_edges.discard((i, j))
        self._children[i].discard(j)
        self._parents[j].discard(i)
        self._neighbors[i].discard(j)
        self._neighbors[j].discard(i)
        self._undirected_neighbors[i].discard(j)
        self._undirected_neighbors[j].discard(i)

    def undir_to_dir_edge(self, tail: str, head: str) -> None:
        """Takes a undirected edge and turns it into a directed one.

        tail indicates the starting node of the edge and head the end node, i.e.
        tail -> head.

        Args:
            tail (str): starting node
            head (str): end node

        Raises:
            AssertionError: if edge does not exist or is not undirected.
        """
        if (tail, head) not in self.undir_edges and (
            head,
            tail,
        ) not in self.undir_edges:
            raise AssertionError("Edge seems not to be undirected or even there at all.")
        self._undir_edges.discard((tail, head))
        self._undir_edges.discard((head, tail))
        self._neighbors[tail].discard(head)
        self._neighbors[head].discard(tail)
        self._undirected_neighbors[tail].discard(head)
        self._undirected_neighbors[head].discard(tail)

        self._add_dir_edge(i=tail, j=head)

    def remove_node(self, node: str) -> None:
        """Remove a node from the graph.

        Args:
            node (str): node to remove
        """
        self._nodes.remove(node)

        self._dir_edges = {(i, j) for i, j in self._dir_edges if node not in (i, j)}

        self._undir_edges = {(i, j) for i, j in self._undir_edges if node not in (i, j)}

        for child in self._children[node]:
            self._parents[child].remove(node)
            self._neighbors[child].remove(node)

        for parent in self._parents[node]:
            self._children[parent].remove(node)
            self._neighbors[parent].remove(node)

        for u_nbr in self._undirected_neighbors[node]:
            self._undirected_neighbors[u_nbr].remove(node)
            self._neighbors[u_nbr].remove(node)

        self._parents.pop(node, "I was never here")
        self._children.pop(node, "I was never here")
        self._neighbors.pop(node, "I was never here")
        self._undirected_neighbors.pop(node, "I was never here")

    def to_dag(self) -> nx.DiGraph:
        r"""Algorithm as described in Chickering (2002).

            1. From PDAG P create DAG G containing all directed edges from P
            2. Repeat the following: Select node v in P s.t.
                i. v has no outgoing edges (children) i.e. \\(ch(v) = \\emptyset \\)

                ii. \\(neigh(v) \\neq \\emptyset\\)
                    Then \\( (pa(v) \\cup (neigh(v) \\) form a clique.
                    For each v that is in a clique and is part of an undirected edge in P
                    i.e. w - v, insert a directed edge w -> v in G.
                    Remove v and all incident edges from P and continue with next node.
                    Until all nodes have been deleted from P.

        Returns:
            nx.DiGraph: DAG that belongs to the MEC implied by the PDAG
        """
        pdag = self.copy()

        dag = nx.DiGraph()
        dag.add_nodes_from(pdag.nodes)
        dag.add_edges_from(pdag.dir_edges)

        if pdag.num_undir_edges == 0:
            return dag
        else:
            while pdag.num_nodes > 0:
                # find node with (1) no directed outgoing edges and
                #                (2) the set of undirected neighbors is either empty or
                #                    undirected neighbors + parents of X are a clique
                found = False
                for node in pdag.nodes:
                    children = pdag.children(node)
                    neighbors = pdag.neighbors(node)
                    # pdag._undirected_neighbors[node]
                    parents = pdag.parents(node)
                    potential_clique_members = neighbors.union(parents)

                    is_clique = pdag.is_clique(potential_clique_members)

                    if not children and (not neighbors or is_clique):
                        found = True
                        # add all edges of node as outgoing edges to dag
                        for edge in pdag.undir_edges:
                            if node in edge:
                                incident_node = set(edge) - {node}
                                dag.add_edge(*incident_node, node)

                        pdag.remove_node(node)
                        break

                if not found:
                    logger.warning("PDAG not extendible: Random DAG on skeleton drawn.")

                    dag = nx.from_pandas_adjacency(self._amat_to_dag(), create_using=nx.DiGraph)

                    break

            return dag

    @property
    def adjacency_matrix(self) -> pd.DataFrame:
        """Returns adjacency matrix.

        The i,jth entry being one indicates that there is an edge
        from i to j. A zero indicates that there is no edge.

        Returns:
            pd.DataFrame: adjacency matrix
        """
        amat = pd.DataFrame(
            np.zeros([self.num_nodes, self.num_nodes]),
            index=self.nodes,
            columns=self.nodes,
        )
        for edge in self.dir_edges:
            amat.loc[edge] = 1
        for edge in self.undir_edges:
            amat.loc[edge] = amat.loc[edge[::-1]] = 1
        return amat

    @property
    def causal_order(self) -> None:
        """Causal order is None.

        This is because PDAGs only allow for a partial causal order.

        Returns:
            None: None
        """
        return None

    def _amat_to_dag(self) -> pd.DataFrame:
        """Transform the adjacency matrix of an PDAG to the adjacency matrix.

            of SOME DAG in the Markov equivalence class.

        Returns:
            pd.DataFrame: DAG, a member of the MEC.
        """
        pdag_amat = self.adjacency_matrix.to_numpy()

        p = pdag_amat.shape[0]
        ## amat to skel
        skel = pdag_amat + pdag_amat.T
        skel[np.where(skel > 1)] = 1
        ## permute skel
        permute_ord = np.random.choice(a=p, size=p, replace=False)
        skel = skel[:, permute_ord][permute_ord]

        ## skel to dag
        for i in range(1, p):
            for j in range(0, i + 1):
                if skel[i, j] == 1:
                    skel[i, j] = 0

        ## inverse permutation
        i_ord = np.sort(permute_ord)
        skel = skel[:, i_ord][i_ord]
        return pd.DataFrame(
            skel,
            index=self.adjacency_matrix.index,
            columns=self.adjacency_matrix.columns,
        )

    def vstructs(self) -> set[tuple[str, str]]:
        """Retrieve v-structures.

        Returns:
            set: set of all v-structures
        """
        vstructures = set()
        for node in self._nodes:
            for p1, p2 in combinations(self._parents[node], 2):
                if p1 not in self._parents[p2] and p2 not in self._parents[p1]:
                    vstructures.add((p1, node))
                    vstructures.add((p2, node))
        return vstructures

    def copy(self) -> PDAG:
        """Return a copy of the graph."""
        return PDAG(
            nodes=list(self._nodes),
            dir_edges=list(self._dir_edges),
            undir_edges=list(self._undir_edges),
        )

    def show(self) -> None:
        """Plot PDAG."""
        graph = self.to_networkx()
        pos = nx.circular_layout(graph)
        nx.draw(graph, pos=pos, with_labels=True)

    def to_networkx(self) -> nx.MultiDiGraph:
        """Convert to networkx graph.

        Returns:
            nx.MultiDiGraph: Graph with directed and undirected edges.
        """
        nx_pdag = nx.MultiDiGraph()
        nx_pdag.add_nodes_from(self.nodes)
        nx_pdag.add_edges_from(self.dir_edges)
        for edge in self.undir_edges:
            nx_pdag.add_edge(*edge)
            nx_pdag.add_edge(*edge[::-1])

        return nx_pdag

    def _meek_mec_enumeration(self, pdag: PDAG, dag_list: list[DAG]) -> None:
        """Apply Meek's MEC enumeration algorithm.

        Args:
            pdag (PDAG): partially directed graph in question.
            dag_list (list): list of currently found DAGs.

        References:
            Wienöbst, Marcel, et al. "Efficient enumeration of Markov equivalent DAGs."
            Proceedings of the AAAI Conference on Artificial Intelligence.
            Vol. 37. No. 10. 2023.
        """
        g_copy = pdag.copy()
        g_copy = self._apply_meek_rules(g_copy)  # Apply Meek rules

        undir_edges = g_copy.undir_edges
        if undir_edges:
            i, j = undir_edges[0]  # Take first undirected edge

        if not g_copy.undir_edges:
            # makes sure that flaoting nodes are preserved
            new_member = DAG()
            new_member.add_nodes_from(g_copy.nodes)
            new_member.add_edges_from(g_copy.dir_edges)
            dag_list.append(new_member)
            return  # Add DAG to current list

        # Recursion first orientation:
        g_copy.undir_to_dir_edge(i, j)
        self._meek_mec_enumeration(pdag=g_copy, dag_list=dag_list)
        g_copy.remove_edge(i, j)

        # Recursion second orientation
        g_copy._add_dir_edge(j, i)
        self._meek_mec_enumeration(pdag=g_copy, dag_list=dag_list)

    def to_allDAGs(self) -> list[DAG]:
        """Recursion algorithm which recursively applies the following steps.

            1. Orient the first undirected edge found.
            2. Apply Meek rules.
            3. Recurse with each direction of the oriented edge.
        This corresponds to Algorithm 2 in Wienöbst et al. (2023).

        References:
            Wienöbst, Marcel, et al. "Efficient enumeration of Markov equivalent DAGs."
            Proceedings of the AAAI Conference on Artificial Intelligence.
            Vol. 37. No. 10. 2023.
        """
        all_dags: list[DAG] = []
        self._meek_mec_enumeration(pdag=self, dag_list=all_dags)
        return all_dags

    # use Meek's cpdag2alldag
    def _apply_meek_rules(self, G: PDAG) -> PDAG:
        """Apply all four Meek rules to a PDAG turning it into a CPDAG.

        Args:
            G (PDAG): PDAG to complete

        Returns:
            PDAG: completed PDAG.
        """
        # Apply Meek Rules
        cpdag = G.copy()
        cpdag = rule_1(pdag=cpdag)
        cpdag = rule_2(pdag=cpdag)
        cpdag = rule_3(pdag=cpdag)
        cpdag = rule_4(pdag=cpdag)
        return cpdag

    def to_random_dag(self) -> DAG:
        """Provides a random DAG residing in the MEC.

        Returns:
            nx.DiGraph: random DAG living in MEC
        """
        to_dag_candidate = self.copy()

        while to_dag_candidate.num_undir_edges > 0:
            chosen_edge = to_dag_candidate.undir_edges[
                np.random.choice(to_dag_candidate.num_undir_edges)
            ]
            choose_orientation = [chosen_edge, chosen_edge[::-1]]
            node_i, node_j = choose_orientation[np.random.choice(len(choose_orientation))]

            to_dag_candidate.undir_to_dir_edge(tail=node_i, head=node_j)
            to_dag_candidate = to_dag_candidate._apply_meek_rules(G=to_dag_candidate)

        return DAG.from_pandas_adjacency(to_dag_candidate.adjacency_matrix)

    @property
    def nodes(self) -> list[str]:
        """Get all nods in current PDAG.

        Returns:
            list: list of nodes.
        """
        return sorted(list(self._nodes))

    @property
    def num_nodes(self) -> int:
        """Number of nodes in current PDAG.

        Returns:
            int: Number of nodes
        """
        return len(self._nodes)

    @property
    def num_undir_edges(self) -> int:
        """Number of undirected edges in current PDAG.

        Returns:
            int: Number of undirected edges
        """
        return len(self._undir_edges)

    @property
    def num_dir_edges(self) -> int:
        """Number of directed edges in current PDAG.

        Returns:
            int: Number of directed edges
        """
        return len(self._dir_edges)

    @property
    def num_adjacencies(self) -> int:
        """Number of adjacent nodes in current PDAG.

        Returns:
            int: Number of adjacent nodes
        """
        return self.num_undir_edges + self.num_dir_edges

    @property
    def undir_edges(self) -> list[tuple[str, str]]:
        """Gives all undirected edges in current PDAG.

        Returns:
            list[tuple[str,str]]: List of undirected edges.
        """
        return list(self._undir_edges)

    @property
    def dir_edges(self) -> list[tuple[str, str]]:
        """Gives all directed edges in current PDAG.

        Returns:
            list[tuple[str,str]]: List of directed edges.
        """
        return list(self._dir_edges)

adjacency_matrix property

Returns adjacency matrix.

The i,jth entry being one indicates that there is an edge from i to j. A zero indicates that there is no edge.

Returns:

Type	Description
DataFrame	pd.DataFrame: adjacency matrix

causal_order property

Causal order is None.

This is because PDAGs only allow for a partial causal order.

Returns:

Name	Type	Description
None	None	None

dir_edges property

Gives all directed edges in current PDAG.

Returns:

Type	Description
list[tuple[str, str]]	list[tuple[str,str]]: List of directed edges.

nodes property

Get all nods in current PDAG.

Returns:

Name	Type	Description
list	list[str]	list of nodes.

num_adjacencies property

Number of adjacent nodes in current PDAG.

Returns:

Name	Type	Description
int	int	Number of adjacent nodes

num_dir_edges property

Number of directed edges in current PDAG.

Returns:

Name	Type	Description
int	int	Number of directed edges

num_nodes property

Number of nodes in current PDAG.

Returns:

Name	Type	Description
int	int	Number of nodes

num_undir_edges property

Number of undirected edges in current PDAG.

Returns:

Name	Type	Description
int	int	Number of undirected edges

undir_edges property

Gives all undirected edges in current PDAG.

Returns:

Type	Description
list[tuple[str, str]]	list[tuple[str,str]]: List of undirected edges.

__init__(nodes=None, dir_edges=None, undir_edges=None)

PDAG constructor.

Parameters:

Name	Type	Description	Default
nodes	list[str] | None	Nodes in the PDAG. Defaults to None.	None
dir_edges	list[tuple[str, str]] | None	directed edges. Defaults to None.	None
undir_edges	list[tuple[str, str]] | None	undirected edges. Defaults to None.	None

Source code in gresit/graphs.py

def __init__(
    self,
    nodes: list[str] | None = None,
    dir_edges: list[tuple[str, str]] | None = None,
    undir_edges: list[tuple[str, str]] | None = None,
) -> None:
    """PDAG constructor.

    Args:
        nodes (list[str] | None, optional): Nodes in the PDAG. Defaults to None.
        dir_edges (list[tuple[str,str]] | None, optional): directed edges. Defaults to None.
        undir_edges (list[tuple[str,str]] | None, optional): undirected edges. Defaults to None.
    """
    if nodes is None:
        nodes = []
    if dir_edges is None:
        dir_edges = []
    if undir_edges is None:
        undir_edges = []

    self._nodes = set(nodes)
    self._undir_edges: set[tuple[str, str]] = set()
    self._dir_edges: set[tuple[str, str]] = set()
    self._parents: defaultdict[str, set[str]] = defaultdict(set)
    self._children: defaultdict[str, set[str]] = defaultdict(set)
    self._neighbors: defaultdict[str, set[str]] = defaultdict(set)
    self._undirected_neighbors: defaultdict[str, set[str]] = defaultdict(set)

    for dir_edge in dir_edges:
        self._add_dir_edge(*dir_edge)
    for unir_edge in undir_edges:
        self._add_undir_edge(*unir_edge)

children(node)

Gives all children of node node.

Parameters:

Name	Type	Description	Default
node	str	node in current PDAG.	required

Returns:

Name	Type	Description
set	set[str]	set of children.

Source code in gresit/graphs.py

def children(self, node: str) -> set[str]:
    """Gives all children of node `node`.

    Args:
        node (str): node in current PDAG.

    Returns:
        set: set of children.
    """
    if node in self._children.keys():
        return self._children[node]
    else:
        return set()

copy()

Return a copy of the graph.

Source code in gresit/graphs.py

def copy(self) -> PDAG:
    """Return a copy of the graph."""
    return PDAG(
        nodes=list(self._nodes),
        dir_edges=list(self._dir_edges),
        undir_edges=list(self._undir_edges),
    )

from_pandas_adjacency(pd_amat) classmethod

Build PDAG from a Pandas adjacency matrix.

Parameters:

Name	Type	Description	Default
pd_amat	DataFrame	input adjacency matrix.	required

Returns:

Type	Description
PDAG	PDAG

Source code in gresit/graphs.py

@classmethod
def from_pandas_adjacency(cls, pd_amat: pd.DataFrame) -> PDAG:
    """Build PDAG from a Pandas adjacency matrix.

    Args:
        pd_amat (pd.DataFrame): input adjacency matrix.

    Returns:
        PDAG
    """
    assert pd_amat.shape[0] == pd_amat.shape[1]
    nodes = pd_amat.columns

    all_connections = []
    start, end = np.where(pd_amat != 0)
    for idx, _ in enumerate(start):
        all_connections.append((pd_amat.columns[start[idx]], pd_amat.columns[end[idx]]))

    temp = [set(i) for i in all_connections]
    temp2 = [arc for arc in all_connections if temp.count(set(arc)) > 1]
    undir_edges = [tuple(item) for item in set(frozenset(item) for item in temp2)]

    dir_edges = [edge for edge in all_connections if edge not in temp2]

    return PDAG(nodes=nodes, dir_edges=dir_edges, undir_edges=undir_edges)

is_adjacent(i, j)

Return True if the graph contains an directed or undirected edge between i and j.

Parameters:

Name	Type	Description	Default
i	str	node i.	required
j	str	node j.	required

Returns:

Name	Type	Description
bool	bool	True if i-j or i->j or i<-j

Source code in gresit/graphs.py

def is_adjacent(self, i: str, j: str) -> bool:
    """Return True if the graph contains an directed or undirected edge between i and j.

    Args:
        i (str): node i.
        j (str): node j.

    Returns:
        bool: True if i-j or i->j or i<-j
    """
    return any(
        (
            (j, i) in self.dir_edges or (j, i) in self.undir_edges,
            (i, j) in self.dir_edges or (i, j) in self.undir_edges,
        )
    )

is_clique(potential_clique)

Check every pair of node X potential_clique is adjacent.

Source code in gresit/graphs.py

def is_clique(self, potential_clique: set[str]) -> bool:
    """Check every pair of node X potential_clique is adjacent."""
    return all(self.is_adjacent(i, j) for i, j in combinations(potential_clique, 2))

neighbors(node)

Gives all neighbors of node node.

Parameters:

Name	Type	Description	Default
node	str	node in current PDAG.	required

Returns:

Name	Type	Description
set	set[str]	set of neighbors.

Source code in gresit/graphs.py

def neighbors(self, node: str) -> set[str]:
    """Gives all neighbors of node `node`.

    Args:
        node (str): node in current PDAG.

    Returns:
        set: set of neighbors.
    """
    if node in self._neighbors.keys():
        return self._neighbors[node]
    else:
        return set()

parents(node)

Gives all parents of node node.

Parameters:

Name	Type	Description	Default
node	str	node in current PDAG.	required

Returns:

Name	Type	Description
set	set[str]	set of parents.

Source code in gresit/graphs.py

def parents(self, node: str) -> set[str]:
    """Gives all parents of node `node`.

    Args:
        node (str): node in current PDAG.

    Returns:
        set: set of parents.
    """
    if node in self._parents.keys():
        return self._parents[node]
    else:
        return set()

remove_edge(i, j)

Removes edge in question.

Parameters:

Name	Type	Description	Default
i	str	tail	required
j	str	head	required

Raises:

Type	Description
AssertionError	if edge does not exist

Source code in gresit/graphs.py

def remove_edge(self, i: str, j: str) -> None:
    """Removes edge in question.

    Args:
        i (str): tail
        j (str): head

    Raises:
        AssertionError: if edge does not exist
    """
    if (i, j) not in self.dir_edges and (i, j) not in self.undir_edges:
        raise AssertionError("Edge does not exist in current PDAG")

    self._undir_edges.discard((i, j))
    self._dir_edges.discard((i, j))
    self._children[i].discard(j)
    self._parents[j].discard(i)
    self._neighbors[i].discard(j)
    self._neighbors[j].discard(i)
    self._undirected_neighbors[i].discard(j)
    self._undirected_neighbors[j].discard(i)

remove_node(node)

Remove a node from the graph.

Parameters:

Name	Type	Description	Default
node	str	node to remove	required

Source code in gresit/graphs.py

def remove_node(self, node: str) -> None:
    """Remove a node from the graph.

    Args:
        node (str): node to remove
    """
    self._nodes.remove(node)

    self._dir_edges = {(i, j) for i, j in self._dir_edges if node not in (i, j)}

    self._undir_edges = {(i, j) for i, j in self._undir_edges if node not in (i, j)}

    for child in self._children[node]:
        self._parents[child].remove(node)
        self._neighbors[child].remove(node)

    for parent in self._parents[node]:
        self._children[parent].remove(node)
        self._neighbors[parent].remove(node)

    for u_nbr in self._undirected_neighbors[node]:
        self._undirected_neighbors[u_nbr].remove(node)
        self._neighbors[u_nbr].remove(node)

    self._parents.pop(node, "I was never here")
    self._children.pop(node, "I was never here")
    self._neighbors.pop(node, "I was never here")
    self._undirected_neighbors.pop(node, "I was never here")

show()

Plot PDAG.

Source code in gresit/graphs.py

def show(self) -> None:
    """Plot PDAG."""
    graph = self.to_networkx()
    pos = nx.circular_layout(graph)
    nx.draw(graph, pos=pos, with_labels=True)

to_allDAGs()

Recursion algorithm which recursively applies the following steps.

1. Orient the first undirected edge found.
2. Apply Meek rules.
3. Recurse with each direction of the oriented edge.

This corresponds to Algorithm 2 in Wienöbst et al. (2023).

References

Wienöbst, Marcel, et al. "Efficient enumeration of Markov equivalent DAGs." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. No. 10. 2023.

Source code in gresit/graphs.py

def to_allDAGs(self) -> list[DAG]:
    """Recursion algorithm which recursively applies the following steps.

        1. Orient the first undirected edge found.
        2. Apply Meek rules.
        3. Recurse with each direction of the oriented edge.
    This corresponds to Algorithm 2 in Wienöbst et al. (2023).

    References:
        Wienöbst, Marcel, et al. "Efficient enumeration of Markov equivalent DAGs."
        Proceedings of the AAAI Conference on Artificial Intelligence.
        Vol. 37. No. 10. 2023.
    """
    all_dags: list[DAG] = []
    self._meek_mec_enumeration(pdag=self, dag_list=all_dags)
    return all_dags

to_dag()

Algorithm as described in Chickering (2002).

1. From PDAG P create DAG G containing all directed edges from P
2. Repeat the following: Select node v in P s.t.
    i. v has no outgoing edges (children) i.e. \\(ch(v) = \\emptyset \\)

    ii. \\(neigh(v) \\neq \\emptyset\\)
        Then \\( (pa(v) \\cup (neigh(v) \\) form a clique.
        For each v that is in a clique and is part of an undirected edge in P
        i.e. w - v, insert a directed edge w -> v in G.
        Remove v and all incident edges from P and continue with next node.
        Until all nodes have been deleted from P.

Returns:

Type	Description
DiGraph	nx.DiGraph: DAG that belongs to the MEC implied by the PDAG

Source code in gresit/graphs.py

def to_dag(self) -> nx.DiGraph:
    r"""Algorithm as described in Chickering (2002).

        1. From PDAG P create DAG G containing all directed edges from P
        2. Repeat the following: Select node v in P s.t.
            i. v has no outgoing edges (children) i.e. \\(ch(v) = \\emptyset \\)

            ii. \\(neigh(v) \\neq \\emptyset\\)
                Then \\( (pa(v) \\cup (neigh(v) \\) form a clique.
                For each v that is in a clique and is part of an undirected edge in P
                i.e. w - v, insert a directed edge w -> v in G.
                Remove v and all incident edges from P and continue with next node.
                Until all nodes have been deleted from P.

    Returns:
        nx.DiGraph: DAG that belongs to the MEC implied by the PDAG
    """
    pdag = self.copy()

    dag = nx.DiGraph()
    dag.add_nodes_from(pdag.nodes)
    dag.add_edges_from(pdag.dir_edges)

    if pdag.num_undir_edges == 0:
        return dag
    else:
        while pdag.num_nodes > 0:
            # find node with (1) no directed outgoing edges and
            #                (2) the set of undirected neighbors is either empty or
            #                    undirected neighbors + parents of X are a clique
            found = False
            for node in pdag.nodes:
                children = pdag.children(node)
                neighbors = pdag.neighbors(node)
                # pdag._undirected_neighbors[node]
                parents = pdag.parents(node)
                potential_clique_members = neighbors.union(parents)

                is_clique = pdag.is_clique(potential_clique_members)

                if not children and (not neighbors or is_clique):
                    found = True
                    # add all edges of node as outgoing edges to dag
                    for edge in pdag.undir_edges:
                        if node in edge:
                            incident_node = set(edge) - {node}
                            dag.add_edge(*incident_node, node)

                    pdag.remove_node(node)
                    break

            if not found:
                logger.warning("PDAG not extendible: Random DAG on skeleton drawn.")

                dag = nx.from_pandas_adjacency(self._amat_to_dag(), create_using=nx.DiGraph)

                break

        return dag

to_networkx()

Convert to networkx graph.

Returns:

Type	Description
MultiDiGraph	nx.MultiDiGraph: Graph with directed and undirected edges.

Source code in gresit/graphs.py

def to_networkx(self) -> nx.MultiDiGraph:
    """Convert to networkx graph.

    Returns:
        nx.MultiDiGraph: Graph with directed and undirected edges.
    """
    nx_pdag = nx.MultiDiGraph()
    nx_pdag.add_nodes_from(self.nodes)
    nx_pdag.add_edges_from(self.dir_edges)
    for edge in self.undir_edges:
        nx_pdag.add_edge(*edge)
        nx_pdag.add_edge(*edge[::-1])

    return nx_pdag

to_random_dag()

Provides a random DAG residing in the MEC.

Returns:

Type	Description
DAG	nx.DiGraph: random DAG living in MEC

Source code in gresit/graphs.py

def to_random_dag(self) -> DAG:
    """Provides a random DAG residing in the MEC.

    Returns:
        nx.DiGraph: random DAG living in MEC
    """
    to_dag_candidate = self.copy()

    while to_dag_candidate.num_undir_edges > 0:
        chosen_edge = to_dag_candidate.undir_edges[
            np.random.choice(to_dag_candidate.num_undir_edges)
        ]
        choose_orientation = [chosen_edge, chosen_edge[::-1]]
        node_i, node_j = choose_orientation[np.random.choice(len(choose_orientation))]

        to_dag_candidate.undir_to_dir_edge(tail=node_i, head=node_j)
        to_dag_candidate = to_dag_candidate._apply_meek_rules(G=to_dag_candidate)

    return DAG.from_pandas_adjacency(to_dag_candidate.adjacency_matrix)

undir_neighbors(node)

Gives all undirected neighbors of node node.

Parameters:

Name	Type	Description	Default
node	str	node in current PDAG.	required

Returns:

Name	Type	Description
set	set[str]	set of undirected neighbors.

Source code in gresit/graphs.py

def undir_neighbors(self, node: str) -> set[str]:
    """Gives all undirected neighbors of node `node`.

    Args:
        node (str): node in current PDAG.

    Returns:
        set: set of undirected neighbors.
    """
    if node in self._undirected_neighbors.keys():
        return self._undirected_neighbors[node]
    else:
        return set()

undir_to_dir_edge(tail, head)

Takes a undirected edge and turns it into a directed one.

tail indicates the starting node of the edge and head the end node, i.e. tail -> head.

Parameters:

Name	Type	Description	Default
tail	str	starting node	required
head	str	end node	required

Raises:

Type	Description
AssertionError	if edge does not exist or is not undirected.

Source code in gresit/graphs.py

def undir_to_dir_edge(self, tail: str, head: str) -> None:
    """Takes a undirected edge and turns it into a directed one.

    tail indicates the starting node of the edge and head the end node, i.e.
    tail -> head.

    Args:
        tail (str): starting node
        head (str): end node

    Raises:
        AssertionError: if edge does not exist or is not undirected.
    """
    if (tail, head) not in self.undir_edges and (
        head,
        tail,
    ) not in self.undir_edges:
        raise AssertionError("Edge seems not to be undirected or even there at all.")
    self._undir_edges.discard((tail, head))
    self._undir_edges.discard((head, tail))
    self._neighbors[tail].discard(head)
    self._neighbors[head].discard(tail)
    self._undirected_neighbors[tail].discard(head)
    self._undirected_neighbors[head].discard(tail)

    self._add_dir_edge(i=tail, j=head)

vstructs()

Retrieve v-structures.

Returns:

Name	Type	Description
set	set[tuple[str, str]]	set of all v-structures

Source code in gresit/graphs.py

def vstructs(self) -> set[tuple[str, str]]:
    """Retrieve v-structures.

    Returns:
        set: set of all v-structures
    """
    vstructures = set()
    for node in self._nodes:
        for p1, p2 in combinations(self._parents[node], 2):
            if p1 not in self._parents[p2] and p2 not in self._parents[p1]:
                vstructures.add((p1, node))
                vstructures.add((p2, node))
    return vstructures

dag2cpdag(dag)

Convertes a DAG into its unique CPDAG.

Parameters:

Name	Type	Description	Default
dag	DiGraph	DAG the CPDAG corresponds to.	required

Returns:

Name	Type	Description
PDAG	PDAG	unique CPDAG

Source code in gresit/graphs.py

def dag2cpdag(dag: nx.DiGraph) -> PDAG:
    """Convertes a DAG into its unique CPDAG.

    Args:
        dag (nx.DiGraph): DAG the CPDAG corresponds to.

    Returns:
        PDAG: unique CPDAG
    """
    copy_dag = dag.copy()
    # Skeleton
    skeleton = nx.to_pandas_adjacency(copy_dag.to_undirected())
    # v-Structures
    vstructures = vstructs(dag=copy_dag)

    for edge in vstructures:  # orient v-structures
        skeleton.loc[edge[::-1]] = 0

    pdag_init = PDAG.from_pandas_adjacency(skeleton)

    # Apply Meek Rules
    cpdag = rule_1(pdag=pdag_init)
    cpdag = rule_2(pdag=cpdag)
    cpdag = rule_3(pdag=cpdag)
    cpdag = rule_4(pdag=cpdag)

    return cpdag

rule_1(pdag)

Applies first Meek rule.

Given the following pattern X -> Y - Z. Orient Y - Z to Y -> Z if X and Z are non-adjacent (otherwise a new v-structure arises).

Parameters:

Name	Type	Description	Default
pdag	PDAG	PDAG before application of rule.	required

Returns:

Name	Type	Description
PDAG	PDAG	PDAG after application of rule.

Source code in gresit/graphs.py

def rule_1(pdag: PDAG) -> PDAG:
    """Applies first Meek rule.

    Given the following pattern X -> Y - Z. Orient Y - Z to Y -> Z
    if X and Z are non-adjacent (otherwise a new v-structure arises).

    Args:
        pdag (PDAG): PDAG before application of rule.

    Returns:
        PDAG: PDAG after application of rule.
    """
    copy_pdag = pdag.copy()
    for edge in copy_pdag.undir_edges:
        reverse_edge = edge[::-1]
        test_edges = [edge, reverse_edge]
        for tail, head in test_edges:
            orient = False
            undir_parents = copy_pdag.parents(tail)
            if undir_parents:
                for parent in undir_parents:
                    if not copy_pdag.is_adjacent(parent, head):
                        orient = True
            if orient:
                copy_pdag.undir_to_dir_edge(tail=tail, head=head)
                break
    return copy_pdag

rule_2(pdag)

Applies the second Meek rule.

Given the following directed triple X -> Y -> Z where X - Z are indeed adjacent. Orient X - Z to X -> Z otherwise a cycle arises.

Parameters:

Name	Type	Description	Default
pdag	PDAG	PDAG before application of rule.	required

Returns:

Name	Type	Description
PDAG	PDAG	PDAG after application of rule.

Source code in gresit/graphs.py

def rule_2(pdag: PDAG) -> PDAG:
    """Applies the second Meek rule.

    Given the following directed triple
    X -> Y -> Z where X - Z are indeed adjacent.
    Orient X - Z to X -> Z otherwise a cycle arises.

    Args:
        pdag (PDAG): PDAG before application of rule.

    Returns:
        PDAG: PDAG after application of rule.
    """
    copy_pdag = pdag.copy()
    for edge in copy_pdag.undir_edges:
        reverse_edge = edge[::-1]
        test_edges = [edge, reverse_edge]
        for tail, head in test_edges:
            orient = False
            undir_children = copy_pdag.children(tail)
            if undir_children:
                for child in undir_children:
                    if head in copy_pdag.children(child):
                        orient = True
            if orient:
                copy_pdag.undir_to_dir_edge(tail=tail, head=head)
                break
    return copy_pdag

rule_3(pdag)

Apply 3rd Meek rule.

Orient X - Z to X -> Z, whenever there are two triples X - Y1 -> Z and X - Y2 -> Z such that Y1 and Y2 are non-adjacent.

Parameters:

Name	Type	Description	Default
pdag	PDAG	PDAG before application of rule.	required

Returns:

Name	Type	Description
PDAG	PDAG	PDAG after application of rule.

Source code in gresit/graphs.py

def rule_3(pdag: PDAG) -> PDAG:
    """Apply 3rd Meek rule.

    Orient X - Z to X -> Z, whenever there are two triples
    X - Y1 -> Z and X - Y2 -> Z such that Y1 and Y2 are non-adjacent.

    Args:
        pdag (PDAG): PDAG before application of rule.

    Returns:
        PDAG: PDAG after application of rule.
    """
    copy_pdag = pdag.copy()
    for edge in copy_pdag.undir_edges:
        reverse_edge = edge[::-1]
        test_edges = [edge, reverse_edge]
        for tail, head in test_edges:
            # if true that tail - node1 -> head and tail - node2 -> head
            # while {node1 U node2} = 0 then orient tail -> head
            orient = False
            num_neighbors = 2
            if len(copy_pdag.undir_neighbors(tail)) >= num_neighbors:
                undir_n = copy_pdag.undir_neighbors(tail)
                selection = [
                    (node1, node2)
                    for node1, node2 in combinations(undir_n, 2)
                    if not copy_pdag.is_adjacent(node1, node2)
                ]
                if selection:
                    for node1, node2 in selection:
                        if head in copy_pdag.parents(node1).intersection(copy_pdag.parents(node2)):
                            orient = True
            if orient:
                copy_pdag.undir_to_dir_edge(tail=tail, head=head)
                break
    return pdag

rule_4(pdag)

Apply 4th Meek rule.

Orient X - Y1 to X -> Y1, whenever there are two triples with X - Z and X - Y1 <- Z and X - Y2 -> Z such that Y1 and Y2 are non-adjacent.

Parameters:

Name	Type	Description	Default
pdag	PDAG	PDAG before application of rule.	required

Returns:

Name	Type	Description
PDAG	PDAG	PDAG after application of rule.

Source code in gresit/graphs.py

def rule_4(pdag: PDAG) -> PDAG:
    """Apply 4th Meek rule.

    Orient X - Y1 to X -> Y1, whenever there are
    two triples with X - Z and X - Y1 <- Z and X - Y2 -> Z
    such that Y1 and Y2 are non-adjacent.

    Args:
        pdag (PDAG): PDAG before application of rule.

    Returns:
        PDAG: PDAG after application of rule.
    """
    copy_pdag = pdag.copy()
    for edge in copy_pdag.undir_edges:
        reverse_edge = edge[::-1]
        test_edges = [edge, reverse_edge]
        for tail, head in test_edges:
            orient = False
            if len(copy_pdag.undir_neighbors(tail)) > 0:
                undirected_n = copy_pdag.undir_neighbors(tail)
                for undir_n in undirected_n:
                    if tail in copy_pdag.children(undir_n):
                        children_select = list(copy_pdag.children(undir_n))
                        if children_select:
                            for parent in children_select:
                                if head in copy_pdag.children(parent):
                                    orient = True
            if orient:
                copy_pdag.undir_to_dir_edge(tail=tail, head=head)
                break
    return pdag

vstructs(dag)

Retrieve all v-structures in a DAG.

Parameters:

Name	Type	Description	Default
dag	DiGraph	DAG in question	required

Returns:

Name	Type	Description
set	set[tuple[str, str]]	Set of all v-structures.

Source code in gresit/graphs.py

def vstructs(dag: nx.DiGraph) -> set[tuple[str, str]]:
    """Retrieve all v-structures in a DAG.

    Args:
        dag (nx.DiGraph): DAG in question

    Returns:
        set: Set of all v-structures.
    """
    vstructures = set()
    for node in dag.nodes():
        for p1, p2 in combinations(list(dag.predecessors(node)), 2):  # get all parents of node
            if not dag.has_edge(p1, p2) and not dag.has_edge(p2, p1):
                vstructures.add((p1, node))
                vstructures.add((p2, node))
    return vstructures