From 3d370ec06be948b5f1aca25b4887c4d4e75096ee Mon Sep 17 00:00:00 2001
From: Jan Bachmann <jan@mannbach.de>
Date: Wed, 11 Sep 2024 20:07:54 +0300
Subject: [PATCH] Introduce DHModel

---
 docs/source/models/directed.rst |  4 ++
 netin/models/__init__.py        |  1 +
 netin/models/dh_model.py        | 74 +++++++++++++++++++++++++++++++++
 netin/models/tests/test_dh.py   | 35 ++++++++++++++++
 4 files changed, 114 insertions(+)
 create mode 100644 netin/models/dh_model.py
 create mode 100644 netin/models/tests/test_dh.py

diff --git a/docs/source/models/directed.rst b/docs/source/models/directed.rst
index 87a661bb..5e77b675 100644
--- a/docs/source/models/directed.rst
+++ b/docs/source/models/directed.rst
@@ -1,6 +1,10 @@
 Directed Graph Models
 =====================
 
+.. autoclass:: netin.models.DHModel
+    :members:
+    :inherited-members:
+
 .. autoclass:: netin.models.DPAModel
     :members:
     :inherited-members:
diff --git a/netin/models/__init__.py b/netin/models/__init__.py
index 5dc79224..70c32c4a 100644
--- a/netin/models/__init__.py
+++ b/netin/models/__init__.py
@@ -8,5 +8,6 @@
 from .pah_model import PAHModel
 from .patch_model import PATCHModel, CompoundLFM
 
+from .dh_model import DHModel
 from .dpa_model import DPAModel
 from .dpah_model import DPAHModel
diff --git a/netin/models/dh_model.py b/netin/models/dh_model.py
new file mode 100644
index 00000000..8a4425e4
--- /dev/null
+++ b/netin/models/dh_model.py
@@ -0,0 +1,74 @@
+from typing import Union
+
+import numpy as np
+
+from .directed_model import DirectedModel
+from ..utils.constants import CLASS_ATTRIBUTE
+from ..link_formation_mechanisms.two_class_homophily import TwoClassHomophily
+
+class DHModel(DirectedModel):
+    """The DHmodel is a directed model that joins new nodes to
+    the existing nodes with a probability proportional to the
+    group assignment of the source and target nodes.
+
+    Parameters
+    ----------
+    N : int
+        Number of nodes to add.
+    f_m : float
+        The minority group fraction.
+    d : float
+        The target network density to reach.
+    plo_M : float
+        The power law exponent for the majority activity.
+    plo_m : float
+        The power law exponent for the minority activity.
+    h_m : float
+        The homophily of the minority nodes.
+    h_M : float
+        The homophily of the majority nodes.
+    seed : Union[int, np.random.Generator, None], optional
+        Randomization seed or generator, by default None.
+        If None, each run will be different.
+    """
+    SHORT = "DH"
+
+    h_m: float
+    h_M: float
+
+    h: TwoClassHomophily
+
+    def __init__(
+            self, *args,
+            N: int, f_m: float, d: float,
+            plo_M: float, plo_m: float,
+            h_m: float, h_M: float,
+            seed: Union[int, np.random.Generator, None] = None, **kwargs):
+        super().__init__(*args, N=N, f_m=f_m, d=d, plo_M=plo_M, plo_m=plo_m, seed=seed, **kwargs)
+        self.h_m = h_m
+        self.h_M = h_M
+
+    def _initialize_lfms(self):
+        """Initializes the link formation mechanisms for the DHModel.
+        """
+        super()._initialize_lfms()
+        self.h = TwoClassHomophily.from_two_class_homophily(
+            node_class_values=self.graph.get_node_class(CLASS_ATTRIBUTE),
+            homophily=(self.h_m, self.h_M))
+
+    def compute_target_probabilities(self, source: int)\
+        -> np.ndarray:
+        """Computes the target probabilities based on :class:`TwoClassHomophily`.
+
+        Parameters
+        ----------
+        source : int
+            The source node.
+
+        Returns
+        -------
+        np.ndarray
+            The target node probabilities.
+        """
+        return super().compute_target_probabilities(source)\
+            * self.h.get_target_probabilities(source)
diff --git a/netin/models/tests/test_dh.py b/netin/models/tests/test_dh.py
new file mode 100644
index 00000000..7d48afde
--- /dev/null
+++ b/netin/models/tests/test_dh.py
@@ -0,0 +1,35 @@
+import numpy as np
+
+from netin.models import DHModel
+
+class TestDPAHModel(object):
+    @staticmethod
+    def _create_model(
+        N=1000, d=0.005, f_m=0.1,
+        plo_M=2.0, plo_m=2.0,
+        h_M=0.2, h_m=0.9, seed=1234) -> DHModel:
+        return DHModel(
+            N=N, d=d, f_m=f_m, plo_M=plo_M, plo_m=plo_m, h_M=h_M, h_m=h_m, seed=seed)
+
+    def test_simulation(self):
+        model = TestDPAHModel._create_model()
+        model.simulate()
+        graph = model.graph
+
+        assert len(graph) == model.N
+        assert graph.is_directed()
+        n_edges = graph.number_of_edges()
+        assert np.isclose(
+            n_edges / (model.N * (model.N - 1)),
+            model.d,
+            atol=1e-5)
+
+    def test_preload_graph(self):
+        pass
+
+    def test_no_invalid_links(self):
+        model = TestDPAHModel._create_model()
+        model.simulate()
+        graph = model.graph
+        for node in graph.nodes():
+            assert not graph.has_edge(node, node)