InfluenceDiffusion 0.0.18__tar.gz → 0.0.19__tar.gz

{influencediffusion-0.0.18 → influencediffusion-0.0.19}/InfluenceDiffusion/Graph.py

@@ -1,12 +1,12 @@
 import numpy as np
-from typing import List, Tuple, Union, Dict
+from typing import List, Tuple, Union, Dict, Optional
 import networkx as nx
 
 __all__ = ["Graph"]
 
 
 class Graph(nx.DiGraph):
-    def __init__(self, incoming_graph_data=None, weights: Union[List, np.array] = None, **attr):
+    def __init__(self, incoming_graph_data=None, weights: Optional[Union[List, np.ndarray]] = None, **attr):
         super().__init__(incoming_graph_data, **attr)
         self.edge_array = np.array(self.edges)
         if weights is None:
@@ -34,7 +34,7 @@ class Graph(nx.DiGraph):
     def _get_edge_weight_attributes(self):
         return np.array([edge[2] for edge in self.edges.data("weight", default=1)])
 
-    def set_weights(self, weights: Union[List, np.array]) -> "Graph":
+    def set_weights(self, weights: Union[List, np.ndarray]) -> "Graph":
         """
         Set weights for the edges in the graph.
 
@@ -58,7 +58,7 @@ class Graph(nx.DiGraph):
         self.weights = self._get_edge_weight_attributes()
         return self
 
-    def get_edges(self, as_array: bool = False) -> Union[np.array, List[Tuple[int, int]]]:
+    def get_edges(self, as_array: bool = False) -> Union[np.ndarray, List[Tuple[int, int]]]:
         """
         Retrieve the edges of the graph.
 
@@ -131,7 +131,7 @@ class Graph(nx.DiGraph):
         """
         return len(self.nodes)
 
-    def get_children_mask(self, vertex: int) -> np.array:
+    def get_children_mask(self, vertex: int) -> np.ndarray:
         """
         Create a mask for edges originating from a specified vertex.
 
@@ -147,7 +147,7 @@ class Graph(nx.DiGraph):
         """
         return self.edge_array[:, 0] == vertex
 
-    def get_children(self, vertex: int, out_type: Union[set, list, np.array] = set) -> Union[set, list, np.array]:
+    def get_children(self, vertex: int) -> np.ndarray:
         """
         Get the children (outgoing edges) of a specified vertex.
 
@@ -155,17 +155,15 @@ class Graph(nx.DiGraph):
         ----------
         vertex : int
             The vertex for which to find children.
-        out_type : Union[set, list, np.array], optional
-            The type of output to return (default is set).
 
         Returns
         -------
         Union[set, list, np.array]
             The children of the vertex in the specified output type.
         """
-        return out_type(self.edge_array[:, 1][self.get_children_mask(vertex)])
+        return self.edge_array[:, 1][self.get_children_mask(vertex)]
 
-    def get_parents_mask(self, vertex: int) -> np.array:
+    def get_parents_mask(self, vertex: int) -> np.ndarray:
         """
         Create a mask for edges leading to a specified vertex.
 
@@ -181,7 +179,7 @@ class Graph(nx.DiGraph):
         """
         return self.edge_array[:, 1] == vertex
 
-    def get_parents(self, vertex: int, out_type: Union[set, list, np.array] = set) -> Union[set, list, np.array]:
+    def get_parents(self, vertex: int) -> np.ndarray:
         """
         Get the parents (incoming edges) of a specified vertex.
 
@@ -189,15 +187,13 @@ class Graph(nx.DiGraph):
         ----------
         vertex : int
             The vertex for which to find parents.
-        out_type : Union[set, list, np.array], optional
-            The type of output to return (default is set).
 
         Returns
         -------
         Union[set, list, np.array]
             The parents of the vertex in the specified output type.
         """
-        return out_type(self.edge_array[:, 0][self.get_parents_mask(vertex)])
+        return self.edge_array[:, 0][self.get_parents_mask(vertex)]
 
     def get_vertex_2_indegree_dict(self, weighted: bool = False) -> Dict[int, Union[float, int]]:
         """
@@ -215,7 +211,7 @@ class Graph(nx.DiGraph):
         """
         return {vertex: self.get_indegree(vertex, weighted=weighted) for vertex in self.get_vertices()}
 
-    def get_outdegree(self, vertex: int, weighted: bool = False) -> Union[float, int]:
+    def get_outdegree(self, vertex: int, weighted: bool = False) -> float:
         """
         Get the outdegree of a specified vertex.
 
@@ -228,13 +224,13 @@ class Graph(nx.DiGraph):
 
         Returns
         -------
-        Union[float, int]
+        float
             The outdegree of the vertex.
         """
         children_weights = self.weights[self.edge_array[:, 0] == vertex]
-        return np.sum(children_weights) if weighted else len(children_weights)
+        return np.sum(children_weights, dtype=float) if weighted else len(children_weights)
 
-    def get_all_outdegrees(self, weighted: bool = False) -> np.array:
+    def get_all_outdegrees(self, weighted: bool = False) -> np.ndarray:
         """
         Get outdegrees for all vertices in the graph.
 
@@ -250,7 +246,7 @@ class Graph(nx.DiGraph):
         """
         return np.array([self.get_outdegree(v, weighted=weighted) for v in self.get_vertices()])
 
-    def get_avg_outdegree(self, weighted: bool = False) -> Union[float, int]:
+    def get_avg_outdegree(self, weighted: bool = False) -> float:
         """
         Calculate the average outdegree of all vertices.
 
@@ -261,10 +257,10 @@ class Graph(nx.DiGraph):
 
         Returns
         -------
-        Union[float, int]
+        float
             The average outdegree of the vertices.
         """
-        return np.mean(self.get_all_outdegrees(weighted=weighted))
+        return np.mean(self.get_all_outdegrees(weighted=weighted), dtype=float)
 
     def get_max_outdegree(self, weighted: bool = False) -> Union[float, int]:
         """
@@ -282,7 +278,7 @@ class Graph(nx.DiGraph):
         """
         return np.max(self.get_all_outdegrees(weighted))
 
-    def get_indegree(self, vertex: int, weighted: bool = False) -> Union[float, int]:
+    def get_indegree(self, vertex: int, weighted: bool = False) -> float:
         """
         Get the indegree of a specified vertex.
 
@@ -299,7 +295,7 @@ class Graph(nx.DiGraph):
             The indegree of the vertex.
         """
         parent_weights = self.weights[self.edge_array[:, 1] == vertex]
-        return np.sum(parent_weights) if weighted else len(parent_weights)
+        return np.sum(parent_weights, dtype=float) if weighted else len(parent_weights)
     
     def get_indegrees_dict(self, weighted: bool = False) -> Dict[int, Union[int, float]]:
         """
@@ -317,7 +313,7 @@ class Graph(nx.DiGraph):
         """
         return {v: self.get_indegree(v, weighted=weighted) for v in self.get_vertices()}
 
-    def get_all_indegrees(self, weighted: bool = False) -> np.array:
+    def get_all_indegrees(self, weighted: bool = False) -> np.ndarray:
         """
         Get indegrees for all vertices in the graph.
 
@@ -364,9 +360,9 @@ class Graph(nx.DiGraph):
         float
             The average indegree of the vertices.
         """
-        return np.mean(self.get_all_indegrees(weighted))
+        return np.mean(self.get_all_indegrees(weighted), dtype=float)
 
-    def get_edges_mask_from_set_to_vertex(self, vertex_set: set, vertex: int) -> np.array:
+    def get_edges_mask_from_set_to_vertex(self, vertex_set: set, vertex: int) -> np.ndarray:
         """
         Create a mask for edges leading from a set of vertices to a specified vertex.
 

{influencediffusion-0.0.18 → influencediffusion-0.0.19}/InfluenceDiffusion/Inference.py

@@ -1,7 +1,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from typing import Dict, Callable
+from typing import Dict, Callable, Optional
 from functools import partial
 from scipy.stats import norm
 from scipy.stats._distn_infrastructure import rv_frozen
@@ -12,7 +12,7 @@ __all__ = ["GLTInferenceModule"]
 
 
 class GLTInferenceModule:
-    def __init__(self, estimator: GLTWeightEstimator, vertex_2_jax_cdf: Dict[int, Callable] = None):
+    def __init__(self, estimator: GLTWeightEstimator, vertex_2_jax_cdf: Optional[Dict[int, Callable]] = None):
         self.estimator = estimator
         if vertex_2_jax_cdf is None:
             self.vertex_2_jax_cdf = {vertex: self._make_jax_cdf(distrib)

{influencediffusion-0.0.18 → influencediffusion-0.0.19}/InfluenceDiffusion/Trace.py

@@ -34,7 +34,7 @@ class Traces(list):
         The graph associated with the traces.
     """
 
-    def __new__(cls, graph: Graph, traces: List[Union[Trace, Tuple[Set]]]) -> Traces:
+    def __new__(cls, graph: Graph, traces: List[Union[Trace, Tuple[Set, ...]]]) -> Traces:
         return super(Traces, cls).__new__(cls, traces)
 
     def __init__(self, graph: Graph, traces: List[Union[Trace, Tuple[Set]]]) -> None:
@@ -42,7 +42,7 @@ class Traces(list):
         super().__init__(traces)
         self.graph = graph
 
-    def append(self, obj: Union[Trace, Tuple[Set]]) -> None:
+    def append(self, obj: Union[Trace, Tuple[Set, ...]]) -> None:
         """Append a new trace or tuple of sets to the list of traces.
 
         Parameters
@@ -68,10 +68,10 @@ class Trace(tuple):
         Whether to check the feasibility of the trace during initialization.
     """
 
-    def __new__(cls, graph: Graph, trace: Tuple[Set], check_feasibility: bool = True) -> Trace:
+    def __new__(cls, graph: Graph, trace: Tuple[Set, ...], check_feasibility: bool = True) -> Trace:
         return super(Trace, cls).__new__(cls, trace)
 
-    def __init__(self, graph: Graph, trace: Tuple[Set], check_feasibility: bool = True) -> None:
+    def __init__(self, graph: Graph, trace: Tuple[Set, ...], check_feasibility: bool = True) -> None:
         self.graph = graph
         self.check_feasibility = check_feasibility
         
@@ -407,7 +407,7 @@ class Trace(tuple):
             A set of edges where activation attempts were made.
         """
         attempt_edges_mask = np.array([self.was_attempt_through_edge(edge) for edge in self.graph.edge_array])
-        return self.graph.edge_array[attempt_edges_mask]
+        return set(self.graph.edge_array[attempt_edges_mask])
 
     def get_edges_between_subseq_active_nodes(self) -> Set[Tuple[int, int]]:
         """Get edges that connect subsequent active nodes.

{influencediffusion-0.0.18 → influencediffusion-0.0.19}/InfluenceDiffusion/estimation_models/BaseWeightEstimator.py

@@ -1,10 +1,11 @@
-from typing import List, Union, Dict
+from typing import Iterable, List, Optional, Tuple, Union, Dict
 import os
 import pickle
 import numpy as np
 from joblib import Parallel, delayed
 from scipy.stats import uniform
 from scipy.stats._distn_infrastructure import rv_frozen
+from abc import ABC, abstractmethod
 
 from ..Trace import Traces, PseudoTraces
 from ..Graph import Graph
@@ -13,14 +14,14 @@ from ..utils import multiple_union, random_vector_inside_simplex
 __all__ = ["BaseWeightEstimator", "BaseGLTWEightEstimator", "BaseICWEightEstimator"]
 
 
-class BaseWeightEstimator:
+class BaseWeightEstimator(ABC):
     """Base class for weight estimation in influence models.
 
     Attributes
     ----------
     graph : Graph
         The graph representing the network.
-    n_jobs : Optional[int]
+    n_jobs : int, optional
         Number of jobs for parallel processing.
     informative_vertices : set
         Set of informative vertices in the graph.
@@ -36,7 +37,7 @@ class BaseWeightEstimator:
         Masks for active parents at time t.
     """
 
-    def __init__(self, graph: Graph, n_jobs: int = None):
+    def __init__(self, graph: Graph, n_jobs: Optional[int] = None):
         self.graph = graph
         self.n_jobs = n_jobs
         self.informative_vertices = set()
@@ -65,7 +66,7 @@ class BaseWeightEstimator:
             active_vertices = np.array(list(trace.get_all_activated_vertices()))
             failed_vertices = np.array(list(trace.get_all_failed_vertices()))
             for vertex in failed_vertices:
-                parents = self.graph.get_parents(vertex, out_type=np.array)
+                parents = self.graph.get_parents(vertex)
                 active_parents_mask = np.in1d(parents, active_vertices)
                 failed_vertices_masks[vertex].append(active_parents_mask)
             
@@ -87,7 +88,7 @@ class BaseWeightEstimator:
             cum_trace_list = [set()] + list(trace.cum_union())
             for new_vertices_tp1, vertices_t, vertices_tm1 in zip(trace[1:], cum_trace_list[1:-1], cum_trace_list[:-2]):
                 for vertex in new_vertices_tp1:
-                    parents = self.graph.get_parents(vertex, out_type=np.array)
+                    parents = self.graph.get_parents(vertex)
                     active_parents_mask_t = np.in1d(parents, np.array(list(vertices_t)))
                     active_parents_mask_tm1 = np.in1d(parents, np.array(list(vertices_tm1)))
                     vertex_2_active_parent_mask_tm1[vertex].append(active_parents_mask_tm1)
@@ -145,7 +146,7 @@ class BaseWeightEstimator:
 
         for vertex, vertex_pseudo_traces in traces.items():
             for vertices_tm1, new_vertices_t in vertex_pseudo_traces:
-                parents = self.graph.get_parents(vertex, out_type=np.array)
+                parents = self.graph.get_parents(vertex)
                 active_parents_mask_tm1 = np.isin(parents, np.array(list(vertices_tm1)))
                 if new_vertices_t:
                     vertices_t = vertices_tm1 | new_vertices_t
@@ -180,14 +181,14 @@ class BaseWeightEstimator:
         else:
             raise NotImplementedError("traces can only be of type Traces or PseudoTraces")
 
-    def _preprocess_traces(self, traces: Union[Traces, PseudoTraces], masks_path: str = None) -> None:
+    def _preprocess_traces(self, traces: Union[Traces, PseudoTraces], masks_path: Optional[str] = None) -> None:
         """Preprocess traces by validating and computing masks.
 
         Parameters
         ----------
         traces : Union[Traces, PseudoTraces]
             The traces to preprocess.
-        masks_path : Optional[str]
+        masks_path : str, optional
             Path to load masks from, if applicable.
         """
         traces = self._validate_traces(traces)
@@ -231,12 +232,12 @@ class BaseWeightEstimator:
         """
         assert len(init_weights) == self.graph.count_edges()
 
-    def _init_weights(self, init_weights: Union[List, np.ndarray]) -> None:
+    def _init_weights(self, init_weights: Optional[Union[List, np.ndarray]]) -> None:
         """Initialize weights based on provided initial weights or generate random weights.
 
         Parameters
         ----------
-        init_weights : Optional[Union[List, np.ndarray]]
+        init_weights : Union[List, np.ndarray], optional
             Initial weights to set. If None, random weights are generated.
         """
         if init_weights is not None:
@@ -245,6 +246,7 @@ class BaseWeightEstimator:
         else:
             self.weights_ = self._generate_random_weights()
 
+    @abstractmethod
     def _generate_random_weights(self) -> np.ndarray:
         """Generate random weights for edges.
 
@@ -253,46 +255,25 @@ class BaseWeightEstimator:
         np.ndarray
             Randomly generated weights.
         """
-        raise NotImplementedError("Subclasses should implement this method.")
-
-    def _optimize_vertex_parent_params(self, vertex: int, **optim_kwargs) -> np.ndarray:
-        """Optimize parameters for a specific vertex.
+        pass
+    
+    @abstractmethod
+    def _fit_vertex_parent_params(self, vertex: int, **optim_kwargs) -> Tuple:
+        """Fit parameters for a specific vertex.
 
         Parameters
         ----------
         vertex : int
-            The vertex to optimize parameters for.
-        optim_kwargs : dict
-            Additional optimization parameters.
+            The vertex to fit parameters for.
 
         Returns
         -------
-        np.ndarray
-            Optimized parameters for the vertex.
+        Tuple
+            Fitted parameters for the vertex.
         """
-        raise NotImplementedError("Subclasses should implement this method.")
+        pass
 
-    def _optimize_informative_vertices_parent_params(self, verbose: bool = False, **optim_kwargs) -> List[np.ndarray]:
-        """Optimize parameters for all informative vertices in parallel.
-
-        Parameters
-        ----------
-        verbose : bool, optional
-            If True, progress messages are printed.
-        optim_kwargs : dict
-            Additional optimization parameters.
-
-        Returns
-        -------
-        List[np.ndarray]
-            List of optimized parameters for each informative vertex.
-        """
-        informative_vertices_params = Parallel(n_jobs=self.n_jobs, verbose=verbose)(
-            delayed(self._optimize_vertex_parent_params)(vertex, **optim_kwargs)
-            for vertex in self.informative_vertices)
-        return informative_vertices_params
-
-    def _set_informative_vertices_parent_params(self, informative_vertices_params: List[np.ndarray]) -> None:
+    def _set_informative_vertices_parent_params(self, informative_vertices_params: Iterable) -> None:
         """Set the optimized parameters for informative vertices.
 
         Parameters
@@ -304,51 +285,51 @@ class BaseWeightEstimator:
             self.weights_[self.graph.get_parents_mask(vertex)] = parent_weights
 
     def _pre_fit(self, traces: Union[Traces, PseudoTraces],
-                 init_weights: Union[List, np.ndarray] = None,
-                 masks_path: str = None) -> None:
+                 init_weights: Optional[Union[List, np.ndarray]] = None,
+                 masks_path: Optional[str] = None) -> None:
         """Prepare for fitting by preprocessing traces and initializing weights.
 
         Parameters
         ----------
         traces : Union[Traces, PseudoTraces]
             The traces to analyze.
-        init_weights : Optional[Union[List, np.ndarray]]
+        init_weights : Union[List, np.ndarray], optional
             Initial weights to set.
-        masks_path : Optional[str]
+        masks_path : str, optional
             Path to load masks from.
         """
         self._preprocess_traces(traces, masks_path=masks_path)
         self._init_weights(init_weights)
 
     def fit(self, traces: Union[Traces, PseudoTraces], 
-            init_weights: Union[List, np.ndarray] = None,
+            init_weights: Optional[Union[List, np.ndarray]] = None,
             verbose: bool = False,
-            masks_path: str = None,
-            **optim_kwargs) -> np.ndarray:
+            masks_path: Optional[str] = None,
+            **optim_kwargs) -> "BaseWeightEstimator":
         """Fit the model to the given traces.
 
         Parameters
         ----------
         traces : Union[Traces, PseudoTraces]
             The traces to analyze.
-        init_weights : Optional[Union[List, np.ndarray]]
+        init_weights : Union[List, np.ndarray], optional
             Initial weights to set.
         verbose : bool, optional
             If True, progress messages are printed.
-        masks_path : Optional[str]
+        masks_path : str, optional
             Path to load masks from.
-        optim_kwargs : dict
-            Additional optimization parameters.
 
         Returns
         -------
-        np.ndarray
-            The final estimated weights.
+        self : object
+            The fitted estimator instance.
         """
         self._pre_fit(traces, init_weights=init_weights, masks_path=masks_path)
-        informative_vertices_params = self._optimize_informative_vertices_parent_params(verbose=verbose, **optim_kwargs)
+        informative_vertices_params = Parallel(n_jobs=self.n_jobs, verbose=verbose)(
+            delayed(self._fit_vertex_parent_params)(vertex, **optim_kwargs)
+            for vertex in self.informative_vertices)
         self._set_informative_vertices_parent_params(informative_vertices_params)
-        return self.weights_
+        return self
 
 
 class BaseGLTWEightEstimator(BaseWeightEstimator):
@@ -360,13 +341,32 @@ class BaseGLTWEightEstimator(BaseWeightEstimator):
         Distribution mapping for vertices.
     """
 
-    def __init__(self, graph: Graph, n_jobs: int = None,
-                 vertex_2_distrib: Union[Dict[int, rv_frozen], rv_frozen] = None):
+    def __init__(self, graph: Graph, n_jobs: Optional[int] = None,
+                 vertex_2_distrib: Optional[Union[Dict[int, rv_frozen], rv_frozen]] = None):
         super().__init__(graph, n_jobs)
+        self._validate_vertex_distributions(vertex_2_distrib)
+
+    def _validate_vertex_distributions(self, vertex_2_distrib: Optional[Union[Dict[int, rv_frozen], rv_frozen]] = None,
+                                   ) -> None:
+        """Initialize vertex distributions. 
+        If vertex_2_distrib is None, a uniform distribution is assigned to each vertex.
+        If vertex_2_distrib is a single distribution, it is assigned to all vertices.
+        If vertex_2_distrib is a dictionary, it must have an entry for each vertex
+
+        Parameters
+        ----------
+        vertex_2_distrib : Union[Dict[int, rv_frozen], rv_frozen], optional
+            Distribution mapping for vertices.
+        """
         if vertex_2_distrib is None:
-            vertex_2_distrib = {v: uniform(0, 1) for v in graph.get_vertices()}
+            vertex_2_distrib = {v: uniform(0, 1) for v in self.graph.get_vertices()}
         elif isinstance(vertex_2_distrib, rv_frozen):
-            vertex_2_distrib = {v: vertex_2_distrib for v in graph.get_vertices()}
+            vertex_2_distrib = {v: vertex_2_distrib for v in self.graph.get_vertices()}
+        elif isinstance(vertex_2_distrib, Dict):
+            assert len(vertex_2_distrib) == self.graph.number_of_nodes(), \
+                  "vertex_2_distrib must have an entry for each vertex in the graph."
+        else:
+            raise ValueError("vertex_2_distrib must be either a rv_frozen distribution or a dictionary")
         self.vertex_2_distrib = vertex_2_distrib
 
     def _generate_random_weights(self) -> np.ndarray:
@@ -380,7 +380,7 @@ class BaseGLTWEightEstimator(BaseWeightEstimator):
         weights = np.zeros(self.graph.count_edges())
         for vertex in self.informative_vertices:
             parent_mask = self.graph.get_parents_mask(vertex)
-            num_parents = self.graph.get_indegree(vertex, weighted=False)
+            num_parents = int(self.graph.get_indegree(vertex, weighted=False))
             support_ub = self.vertex_2_distrib[vertex].support()[1]
             support_ub = 1 if support_ub == np.inf else support_ub
             weights[parent_mask] = random_vector_inside_simplex(num_parents, ub=support_ub)
@@ -434,4 +434,6 @@ class BaseICWEightEstimator(BaseWeightEstimator):
         AssertionError
             If initial weights are not within the range [0, 1].
         """
-        assert np.all((init_weights >= 0) & (init_weights <= 1))
+        if init_weights is not None:
+            init_weights = np.array(init_weights)
+            assert np.all((init_weights >= 0) & (init_weights <= 1))

{influencediffusion-0.0.18 → influencediffusion-0.0.19}/InfluenceDiffusion/estimation_models/EMEstimation.py

@@ -1,5 +1,5 @@
 import numpy as np
-from typing import Union
+from typing import Optional, Union
 
 from .BaseWeightEstimator import BaseWeightEstimator, BaseGLTWEightEstimator, BaseICWEightEstimator
 from ..Trace import Traces, PseudoTraces
@@ -22,14 +22,14 @@ class BaseWeightEstimatorEM(BaseWeightEstimator):
         Optimizes the parent weights for the specified vertex using the EM algorithm.
     """
 
-    def _em_step(self, vertex: int, vertex_parent_weights: np.array) -> np.array:
+    def _em_step(self, vertex: int, vertex_parent_weights: np.ndarray) -> np.ndarray:
         """Perform a single EM step for the specified vertex.
 
         Parameters
         ----------
         vertex : int
             The vertex for which to perform the EM step.
-        vertex_parent_weights : np.array
+        vertex_parent_weights : np.ndarray
             The current weights of the parent vertices.
 
         Returns
@@ -39,7 +39,7 @@ class BaseWeightEstimatorEM(BaseWeightEstimator):
         """
         raise NotImplementedError()
 
-    def _optimize_vertex_parent_params(self, vertex: int, max_iter: int = 50, tol: float = 1e-3) -> np.array:
+    def _optimize_vertex_parent_params(self, vertex: int, max_iter: int = 50, tol: float = 1e-3) -> np.ndarray:
         """Optimize the parent weights for the specified vertex using the EM algorithm.
 
         Parameters
@@ -99,7 +99,7 @@ class ICWeightEstimatorEM(BaseICWEightEstimator, BaseWeightEstimatorEM):
                 n_traces_parents_activated += self._failed_vertices_masks[vertex].sum(0)
             self._vertex_2_num_traces_parents_activated[vertex] = n_traces_parents_activated
 
-    def _em_step(self, vertex: int, vertex_parent_weights: np.array) -> np.array:
+    def _em_step(self, vertex: int, vertex_parent_weights: np.ndarray) -> np.ndarray:
         """Perform a single EM step for the specified vertex.
 
         This method computes the updated weights based on the current parent weights.
@@ -133,7 +133,7 @@ class ICWeightEstimatorEM(BaseICWEightEstimator, BaseWeightEstimatorEM):
         upd_weights[seen_parents_mask] /= n_traces_parents_activated[seen_parents_mask]
         return np.clip(upd_weights, 0, 1)
 
-    def _preprocess_traces(self, traces: Union[Traces, PseudoTraces], masks_path: str = None) -> None:
+    def _preprocess_traces(self, traces: Union[Traces, PseudoTraces], masks_path: Optional[str] = None) -> None:
         """Preprocesses the input traces and prepares necessary data for EM steps.
 
         Parameters
@@ -146,7 +146,7 @@ class ICWeightEstimatorEM(BaseICWEightEstimator, BaseWeightEstimatorEM):
         super()._preprocess_traces(traces=traces, masks_path=masks_path)
         self._precompute_num_traces_parents_activated()
 
-    def _generate_random_weights(self) -> np.array:
+    def _generate_random_weights(self) -> np.ndarray:
         """Generates random initial weights for the parent vertices.
 
         Returns
@@ -173,7 +173,7 @@ class LTWeightEstimatorEM(BaseGLTWEightEstimator, BaseWeightEstimatorEM):
         Generates random initial weights for the parent vertices.
     """
 
-    def __init__(self, graph: Graph, n_jobs: int = None):
+    def __init__(self, graph: Graph, n_jobs: Optional[int] = None):
         """Initialize the Linear Threshold Weight Estimator.
 
         Parameters
@@ -185,19 +185,19 @@ class LTWeightEstimatorEM(BaseGLTWEightEstimator, BaseWeightEstimatorEM):
         """
         super().__init__(graph, n_jobs=n_jobs)
 
-    def _em_step(self, vertex: int, vertex_parent_weights: np.array) -> np.array:
+    def _em_step(self, vertex: int, vertex_parent_weights: np.ndarray) -> np.ndarray:
         """Perform a single EM step for the specified vertex.
 
         Parameters
         ----------
         vertex : int
             The vertex for which to perform the EM step.
-        vertex_parent_weights : np.array
+        vertex_parent_weights : np.ndarray
             The current weights of the parent vertices.
 
         Returns
         -------
-        np.array
+        np.ndarray
             The updated weights for the parent vertices.
         """
         if vertex in self._vertex_2_active_parent_mask_t:
@@ -222,12 +222,12 @@ class LTWeightEstimatorEM(BaseGLTWEightEstimator, BaseWeightEstimatorEM):
         upd_weights = H_uvs / (H_empty_v + H_uvs.sum())
         return upd_weights
 
-    def _generate_random_weights(self) -> np.array:
+    def _generate_random_weights(self) -> np.ndarray:
         """Generates random initial weights for the parent vertices.
 
         Returns
         -------
-        np.array
+        np.ndarray
             An array of random weights for the parent vertices.
         """
         weights = np.zeros(self.graph.count_edges())