Merge pull request #245 from CiwPython/impovements_from_PRs

Impovements from PRs

ciw/arrival_node.py

@@ -2,7 +2,7 @@
 from .individual import Individual
 
 
-class ArrivalNode(object):
+class ArrivalNode:
     """Class for the arrival node of the network.
     
     See Also

ciw/disciplines.py

@@ -1,25 +1,49 @@
+from typing import List
+
+from ciw.individual import Individual
 from ciw.auxiliary import random_choice
 
 
-def FIFO(individuals):
+def FIFO(individuals: List[Individual]) -> Individual:
     """
-    FIFO: First in first out / First come first served
-    Returns the individual at the head of the queue
+    FIFO: First In, First Out (FIFO)
+
+    Returns the individual at the head of the queue.
+
+    Parameters:
+    - individuals (List[Individual]): List of individuals in the queue.
+
+    Returns:
+    - Individual: The individual at the head of the queue.
     """
     return individuals[0]
 
 
-def SIRO(individuals):
+def SIRO(individuals: List[Individual]) -> Individual:
     """
-    SIRO: Service in random order
-    Returns a random individual from the queue
+    SIRO: Service In Random Order (SIRO)
+
+    Returns a random individual from the queue.
+
+    Parameters:
+    - individuals (List[Individual]): List of individuals in the queue.
+
+    Returns:
+    - Individual: A randomly selected individual from the queue.
     """
     return random_choice(individuals)
 
 
-def LIFO(individuals):
+def LIFO(individuals: List[Individual]) -> Individual:
     """
-    LIFO: Last in first out / Last come first served
-    Returns the individual who joined the queue most recently
+    LIFO: Last In, First Out (LIFO)
+
+    Returns the individual who joined the queue most recently.
+
+    Parameters:
+    - individuals (List[Individual]): List of individuals in the queue.
+
+    Returns:
+    - Individual: The individual who joined the queue most recently.
     """
-    return individuals[-1]
+    return individuals[-1]

ciw/dists/distributions.py

@@ -1,7 +1,7 @@
-from ciw.auxiliary import *
-from itertools import cycle
-import numpy as np
+'''Distributions available in Ciw.'''
+
 import copy
+from itertools import cycle
 from operator import add, mul, sub, truediv
 from random import (
     expovariate,
@@ -11,7 +11,12 @@
     lognormvariate,
     weibullvariate,
 )
+from typing import List, NoReturn
 
+import numpy as np
+
+from ciw.auxiliary import *
+from ciw.individual import Individual
 
 class Distribution(object):
     """
@@ -99,7 +104,7 @@ def __init__(self, lower, upper):
         self.upper = upper
 
     def __repr__(self):
-        return "Uniform: {0}, {1}".format(self.lower, self.upper)
+        return f"Uniform(lower={self.lower}, upper={self.upper})"
 
     def sample(self, t=None, ind=None):
         return uniform(self.lower, self.upper)
@@ -121,7 +126,7 @@ def __init__(self, value):
         self.value = value
 
     def __repr__(self):
-        return "Deterministic: {0}".format(self.value)
+        return f"Deterministic(value={self.value})"
 
     def sample(self, t=None, ind=None):
         return self.value
@@ -151,7 +156,7 @@ def __init__(self, lower, mode, upper):
         self.upper = upper
 
     def __repr__(self):
-        return "Triangular: {0}, {1}, {2}".format(self.lower, self.mode, self.upper)
+        return f"Triangular(lower={self.lower}, mode={self.mode}, upper={self.upper})"
 
     def sample(self, t=None, ind=None):
         return triangular(self.lower, self.upper, self.mode)
@@ -173,7 +178,7 @@ def __init__(self, rate):
         self.rate = rate
 
     def __repr__(self):
-        return "Exponential: {0}".format(self.rate)
+        return f"Exponential(rate={self.rate})"
 
     def sample(self, t=None, ind=None):
         return expovariate(self.rate)
@@ -193,7 +198,7 @@ def __init__(self, shape, scale):
         self.scale = scale
 
     def __repr__(self):
-        return "Gamma: {0}, {1}".format(self.shape, self.scale)
+        return f"Gamma(shape={self.shape}, scale={self.scale})"
 
     def sample(self, t=None, ind=None):
         return gammavariate(self.shape, self.scale)
@@ -213,7 +218,7 @@ def __init__(self, mean, sd):
         self.sd = sd
 
     def __repr__(self):
-        return "Normal: {0}, {1}".format(self.mean, self.sd)
+        return f"Normal(mean={self.mean}, sd={self.sd})"
 
     def sample(self, t=None, ind=None):
         return truncated_normal(self.mean, self.sd)
@@ -233,7 +238,7 @@ def __init__(self, mean, sd):
         self.sd = sd
 
     def __repr__(self):
-        return "Lognormal: {0}, {1}".format(self.mean, self.sd)
+        return f"Lognormal(mean={self.mean}, sd={self.sd})"
 
     def sample(self, t=None, ind=None):
         return lognormvariate(self.mean, self.sd)
@@ -253,7 +258,7 @@ def __init__(self, scale, shape):
         self.shape = shape
 
     def __repr__(self):
-        return "Weibull: {0}, {1}".format(self.scale, self.shape)
+        return f"Weibull(shape={self.scale}, scale={self.shape})"
 
     def sample(self, t=None, ind=None):
         return weibullvariate(self.scale, self.shape)
@@ -298,7 +303,10 @@ def __init__(self, sequence):
         self.generator = cycle(self.sequence)
 
     def __repr__(self):
-        return "Sequential"
+        if len(self.sequence) <= 3:
+            return f"Sequential({self.sequence})"
+        else:
+            return f"Sequential([{self.sequence[0]}, ..., {self.sequence[-1]}])"
 
     def sample(self, t=None, ind=None):
         return next(self.generator)
@@ -324,7 +332,7 @@ def __init__(self, values, probs):
         self.probs = probs
 
     def __repr__(self):
-        return "Pmf"
+        return f"Pmf(values={self.values}, probs={self.probs})"
 
     def sample(self, t=None, ind=None):
         return random_choice(self.values, self.probs)
@@ -420,7 +428,7 @@ def __init__(self, rate, num_phases):
         super().__init__(initial_state, absorbing_matrix)
 
     def __repr__(self):
-        return f"Erlang: {self.rate}, {self.num_phases}"
+        return f"Erlang(rate={self.rate}, k={self.num_phases})"
 
 
 class HyperExponential(PhaseType):
@@ -611,7 +619,7 @@ def sample(self, t=None, ind=None):
         return ciw.rng.poisson(lam=self.rate)
 
     def __repr__(self):
-        return f"Poisson: {self.rate}"
+        return f"Poisson(rate={self.rate})"
 
 
 class Geometric(Distribution):
@@ -634,7 +642,7 @@ def sample(self, t=None, ind=None):
         return ciw.rng.geometric(p=self.prob)
 
     def __repr__(self):
-        return f"Geometric: {self.prob}"
+        return f"Geometric(prob={self.prob})"
 
 
 class Binomial(Distribution):
@@ -663,4 +671,76 @@ def sample(self, t=None, ind=None):
         return ciw.rng.binomial(n=self.n, p=self.prob)
 
     def __repr__(self):
-        return f"Binomial: {self.n}, {self.prob}"
+        return f"Binomial(n={self.n}, prob={self.prob})"
+
+
+class MixtureDistribution(Distribution):
+    """
+    A mixture distribution combining multiple probability distributions.
+
+    Parameters
+    ----------
+    dists : List[Distribution]
+        A list of probability distributions to be combined in the mixture.
+    probs : List[float]
+        A list of weights corresponding to the importance of each distribution in the mixture.
+        The weights must sum to 1.
+
+    Attributes
+    ----------
+    probs : List[float]
+        List of weights assigned to each distribution in the mixture.
+    dists : List[Distribution]
+        List of probability distributions in the mixture.
+
+    Methods
+    -------
+    sample(t: float, inds: List[Individual] = None) -> float:
+        Generate a random sample from the mixture distribution.
+
+    Notes
+    -----
+    The weights in `probs` should sum to 1, indicating the relative importance of each distribution
+    in the mixture. The distributions in `dists` should be instances of `ciw.dists.Distribution`.
+    """
+
+    def __init__(self, dists: List[Distribution], probs: List[float]) -> NoReturn:
+        """
+        Initialize the MixtureDistribution.
+
+        Parameters
+        ----------
+        dists : List[Distribution]
+            A list of probability distributions to be combined in the mixture.
+        probs : List[float]
+            A list of weights corresponding to the importance of each distribution in the mixture.
+            The weights must sum to 1.
+        """
+        self.probs = probs
+        self.dists = dists
+
+    def sample(self, t: float = None, inds: List[Individual] = None) -> float:
+        """
+        Generate a random sample from the mixture distribution.
+
+        Parameters
+        ----------
+        t : float
+            The time parameter for the sample generation.
+        inds : List[Individual], optional
+            List of individuals associated with the sample, if applicable.
+
+        Returns
+        -------
+        float
+            A random sample from the mixture distribution.
+        """
+        chosen_dist = random.choices(
+            population=self.dists,
+            weights=self.probs,
+            k=1)[0]
+
+        return chosen_dist.sample(t, inds)
+
+    def __repr__(self):
+        return "MixtureDistribution"

ciw/import_params.py

@@ -43,23 +43,23 @@ def create_network(
         "service_distributions": service_distributions,
     }
 
-    if baulking_functions != None:
+    if baulking_functions is not None:
         params["baulking_functions"] = baulking_functions
-    if class_change_matrices != None:
+    if class_change_matrices is not None:
         params["class_change_matrices"] = class_change_matrices
     if class_change_time_distributions is not None:
         params["class_change_time_distributions"] = class_change_time_distributions
-    if priority_classes != None:
+    if priority_classes is not None:
         params["priority_classes"] = priority_classes
-    if queue_capacities != None:
+    if queue_capacities is not None:
         params["queue_capacities"] = queue_capacities
-    if routing != None:
+    if routing is not None:
         params["routing"] = routing
-    if batching_distributions != None:
+    if batching_distributions is not None:
         params["batching_distributions"] = batching_distributions
-    if ps_thresholds != None:
+    if ps_thresholds is not None:
         params["ps_thresholds"] = ps_thresholds
-    if server_priority_functions != None:
+    if server_priority_functions is not None:
         params["server_priority_functions"] = server_priority_functions
     if reneging_time_distributions is not None:
         params["reneging_time_distributions"] = reneging_time_distributions