Repository: AngryMaciek/angry-moran-simulator

Issue:
Python modularisation
double-check the modularisation setup in the `init`.

Discussion:
codecov-io: # [Codecov](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24?src=pr&el=h1) Report
> Merging [#24](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24?src=pr&el=desc) (e294e6b) into [master](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/commit/a065091015628bd568f9168b3abf3d8c84167be7?el=desc) (a065091) will **not change** coverage.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24/graphs/tree.svg?width=650&height=150&src=pr&token=V9IFEOWN71)](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##            master       #24    +/-   ##
==========================================
  Coverage   100.00%   100.00%            
==========================================
  Files            1         5     +4     
  Lines           10       870   +860     
==========================================
+ Hits            10       870   +860     
```


| [Impacted Files](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [moranpycess/\_\_init\_\_.py](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24/diff?src=pr&el=tree#diff-bW9yYW5weWNlc3MvX19pbml0X18ucHk=) | `100.00% <ø> (ø)` | |
| [moranpycess/MoranProcess.py](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24/diff?src=pr&el=tree#diff-bW9yYW5weWNlc3MvTW9yYW5Qcm9jZXNzLnB5) | `100.00% <100.00%> (ø)` | |
| [moranpycess/MoranProcess2D.py](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24/diff?src=pr&el=tree#diff-bW9yYW5weWNlc3MvTW9yYW5Qcm9jZXNzMkQucHk=) | `100.00% <100.00%> (ø)` | |
| [moranpycess/MoranProcess3D.py](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24/diff?src=pr&el=tree#diff-bW9yYW5weWNlc3MvTW9yYW5Qcm9jZXNzM0QucHk=) | `100.00% <100.00%> (ø)` | |
| [moranpycess/Individual.py](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24/diff?src=pr&el=tree#diff-bW9yYW5weWNlc3MvSW5kaXZpZHVhbC5weQ==) | `100.00% <0.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24?src=pr&el=footer). Last update [a065091...e294e6b](https://codecov.io/gh/AngryMaciek/angry-moran-simulator/pull/24?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index fbc60d7..db7d90e 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -51,7 +51,7 @@ jobs:
         run: |
           flake8 --max-line-length=88 --ignore F401 moranpycess/__init__.py
           flake8 --max-line-length=88 moranpycess/Individual.py
-          flake8 --max-line-length=95 --ignore F401,E231,W503,E741 moranpycess/MoranProcess.py
+          flake8 --max-line-length=101 --ignore F401,E231,W503,E741 moranpycess/MoranProcess.py
           flake8 --max-line-length=101 --ignore F401,E231,W503,E741 moranpycess/MoranProcess2D.py
           flake8 --max-line-length=101 --ignore F401,E231,W503,E741 moranpycess/MoranProcess3D.py
           flake8 --max-line-length=88 --ignore F401,E402 tests/unit/context.py
diff --git a/documentation.md b/documentation.md
index 4426052..e48d0aa 100644
--- a/documentation.md
+++ b/documentation.md
@@ -77,7 +77,7 @@ ## General Moran Process
 * per-sub-population Death Fitness of an individual from a given sub-popualtion
 * Entropy of the distribution of Strategies in the whole population
 
-Additionally to the *MoranProcess* class the user is equipped with several plotting functions to visualise results of the simulation:
+The *MoranProcess* class is equipped with several plotting methods to visualise results of the simulation:
 * `PlotSize`
 * `PlotAvgBirthPayoff`
 * `PlotAvgDeathPayoff`
@@ -87,13 +87,11 @@ ## General Moran Process
 
 Each of which with the same signature:
 ```python
-def FUNCTION(mp, df, path):
+def FUNCTION(self, df, path):
 ```
 
 With the following arguments:
 ```python
-mp # instance of the MoranProcess
-
 df # simulation results - pandas dataframe returned by the method .simulate()
 
 path # path for the output plot in png format
@@ -101,12 +99,12 @@ ## General Moran Process
 
 Following the previous simulation one may generate the plots with:
 ```python
-moranpycess.PlotSize(mp, df, "Size.png")
-moranpycess.PlotAvgBirthPayoff(mp, df, "AvgBirthPayoff.png")
-moranpycess.PlotAvgDeathPayoff(mp, df, "AvgDeathPayoff.png")
-moranpycess.PlotBirthFitness(mp, df, "BirthFitness.png")
-moranpycess.PlotDeathFitness(mp, df, "DeathFitness.png")
-moranpycess.PlotEntropy(mp, df, "Entropy.png")
+mp.PlotSize(df, "Size.png")
+mp.PlotAvgBirthPayoff(df, "AvgBirthPayoff.png")
+mp.PlotAvgDeathPayoff(df, "AvgDeathPayoff.png")
+mp.PlotBirthFitness(df, "BirthFitness.png")
+mp.PlotDeathFitness(df, "DeathFitness.png")
+mp.PlotEntropy(df, "Entropy.png")
 ```
 
 ## Moran Model based on 2D neighbourhood
@@ -172,28 +170,26 @@ ## Moran Model based on 2D neighbourhood
 * per-sub-population sub-population's size
 * Entropy of the distribution of Strategies in the whole population
 
-Additionally to the *MoranProcess2D* class the user is equipped with three plotting functions to visualise results of the simulation:
+The *MoranProcess2D* class is equipped with three plotting methods to visualise results of the simulation:
 * `PlotSize2D`
 * `PlotEntropy2D`
 * `PlotPopulationSnapshot2D`
 
 With `PlotSize2D` and `PlotEntropy2D` having the same signatures as their previous analogues. The latter, `PlotPopulationSnapshot2D`, may produce a heatmap-like snapshot of a population at it's current state:
 ```python
-def PlotPopulationSnapshot2D(mp, path):
+def PlotPopulationSnapshot2D(self, path):
 ```
 
 With the following arguments:
 ```python
-mp # instance of the MoranProcess
-
 path # path for the output plot in png format
 ```
 
 Following the previous simulation one may generate the plots with:
 ```python
-moranpycess.PlotSize2D(mp, df, "Size2D.png")
-moranpycess.PlotEntropy2D(mp, df, "Entropy2D.png")
-moranpycess.PlotPopulationSnapshot2D(mp, "PopulationSnapshot2D.png")
+mp.PlotSize2D(df, "Size2D.png")
+mp.PlotEntropy2D(df, "Entropy2D.png")
+mp.PlotPopulationSnapshot2D("PopulationSnapshot2D.png")
 ```
 
 ## Moran Model based on 3D neighbourhood
@@ -255,15 +251,15 @@ ## Moran Model based on 3D neighbourhood
 * per-sub-population sub-population's size
 * Entropy of the distribution of Strategies in the whole population
 
-Additionally to the *MoranProcess3D* class the user is equipped with two plotting functions to visualise results of the simulation:
+The *MoranProcess3D* class is equipped with two plotting methods to visualise results of the simulation:
 * `PlotSize3D`
 * `PlotEntropy3D`
 
 The functions have the same signatures as their previous analogues.  
 Following the previous simulation one may generate the plots with:
 ```python
-moranpycess.PlotSize3D(mp, df, "Size3D.png")
-moranpycess.PlotEntropy3D(mp, df, "Entropy3D.png")
+mp.PlotSize3D(df, "Size3D.png")
+mp.PlotEntropy3D(df, "Entropy3D.png")
 ```
 
 ## Use cases
diff --git a/moranpycess/MoranProcess.py b/moranpycess/MoranProcess.py
index 57e5f68..6e96db2 100644
--- a/moranpycess/MoranProcess.py
+++ b/moranpycess/MoranProcess.py
@@ -425,114 +425,108 @@ def UpdateEntropy(self):
             if fraction != 0.0:
                 self.Entropy -= fraction * np.log2(fraction)
 
-
-def PlotSize(mp, df, path):
-    """Plot the sub-populations' sizes after a simulation of a given Moran Process."""
-    plt.figure(figsize=(14, 6))
-    ax = plt.gca()
-    ax.tick_params(width=1)
-    for axis in ["top", "bottom", "left", "right"]:
-        ax.spines[axis].set_linewidth(1)
-    cmap = plt.get_cmap("coolwarm")
-    columns = [l + "__size" for l in mp.init_label_list]
-    df_copy = df[columns].copy()
-    df_copy.columns = mp.init_label_list
-    df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
-    population_size = len(mp.population)
-    ax.set_ylim([0, population_size])
-    plt.xlabel("Generation", size=14)
-    plt.ylabel("# Individuals", size=14)
-    ax.tick_params(axis="both", which="major", labelsize=12)
-    ax.legend(loc=4, fontsize=20)
-    plt.savefig(fname=path, dpi=300)
-
-
-def PlotAvgBirthPayoff(mp, df, path):
-    """Plot the sub-populations' AvgBirthPayoff after a simulation of a given Moran Process."""
-    plt.figure(figsize=(14, 6))
-    ax = plt.gca()
-    ax.tick_params(width=1)
-    for axis in ["top", "bottom", "left", "right"]:
-        ax.spines[axis].set_linewidth(1)
-    cmap = plt.get_cmap("coolwarm")
-    columns = [l + "__AvgBirthPayoff" for l in mp.init_label_list]
-    df_copy = df[columns].copy()
-    df_copy.columns = mp.init_label_list
-    df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
-    plt.xlabel("Generation", size=14)
-    plt.ylabel("Average Birth Payoff", size=14)
-    ax.tick_params(axis="both", which="major", labelsize=12)
-    ax.legend(loc=4, fontsize=20)
-    plt.savefig(fname=path, dpi=300)
-
-
-def PlotAvgDeathPayoff(mp, df, path):
-    """Plot the sub-populations' AvgDeathPayoff after a simulation of a given Moran Process."""
-    plt.figure(figsize=(14, 6))
-    ax = plt.gca()
-    ax.tick_params(width=1)
-    for axis in ["top", "bottom", "left", "right"]:
-        ax.spines[axis].set_linewidth(1)
-    cmap = plt.get_cmap("coolwarm")
-    columns = [l + "__AvgDeathPayoff" for l in mp.init_label_list]
-    df_copy = df[columns].copy()
-    df_copy.columns = mp.init_label_list
-    df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
-    plt.xlabel("Generation", size=14)
-    plt.ylabel("Average Death Payoff", size=14)
-    ax.tick_params(axis="both", which="major", labelsize=12)
-    ax.legend(loc=4, fontsize=20)
-    plt.savefig(fname=path, dpi=300)
-
-
-def PlotBirthFitness(mp, df, path):
-    """Plot the sub-populations' BirthFitness after a simulation of a given Moran Process."""
-    plt.figure(figsize=(14, 6))
-    ax = plt.gca()
-    ax.tick_params(width=1)
-    for axis in ["top", "bottom", "left", "right"]:
-        ax.spines[axis].set_linewidth(1)
-    cmap = plt.get_cmap("coolwarm")
-    columns = [l + "__BirthFitness" for l in mp.init_label_list]
-    df_copy = df[columns].copy()
-    df_copy.columns = mp.init_label_list
-    df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
-    plt.xlabel("Generation", size=14)
-    plt.ylabel("Birth Fitness", size=14)
-    ax.tick_params(axis="both", which="major", labelsize=12)
-    ax.legend(loc=4, fontsize=20)
-    plt.savefig(fname=path, dpi=300)
-
-
-def PlotDeathFitness(mp, df, path):
-    """Plot the sub-populations' DeathFitness after a simulation of a given Moran Process."""
-    plt.figure(figsize=(14, 6))
-    ax = plt.gca()
-    ax.tick_params(width=1)
-    for axis in ["top", "bottom", "left", "right"]:
-        ax.spines[axis].set_linewidth(1)
-    cmap = plt.get_cmap("coolwarm")
-    columns = [l + "__DeathFitness" for l in mp.init_label_list]
-    df_copy = df[columns].copy()
-    df_copy.columns = mp.init_label_list
-    df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
-    plt.xlabel("Generation", size=14)
-    plt.ylabel("Death Fitness", size=14)
-    ax.tick_params(axis="both", which="major", labelsize=12)
-    ax.legend(loc=4, fontsize=20)
-    plt.savefig(fname=path, dpi=300)
-
-
-def PlotEntropy(mp, df, path):
-    """Plot the whole populations entropy after a simulation of a given Moran Process."""
-    plt.figure(figsize=(14, 6))
-    ax = plt.gca()
-    ax.tick_params(width=1)
-    for axis in ["top", "bottom", "left", "right"]:
-        ax.spines[axis].set_linewidth(1)
-    df["Entropy"].plot(color="black", linewidth=1.5, ax=ax, label="Entropy")
-    plt.xlabel("Generation", size=14)
-    plt.ylabel("", size=14)
-    ax.tick_params(axis="both", which="major", labelsize=12)
-    ax.legend(loc=4, fontsize=20)
-    plt.savefig(fname=path, dpi=300)
+    def PlotSize(self, df, path):
+        """Plot the sub-populations' sizes after a simulation of a given Moran Process."""
+        plt.figure(figsize=(14, 6))
+        ax = plt.gca()
+        ax.tick_params(width=1)
+        for axis in ["top", "bottom", "left", "right"]:
+            ax.spines[axis].set_linewidth(1)
+        cmap = plt.get_cmap("coolwarm")
+        columns = [l + "__size" for l in self.init_label_list]
+        df_copy = df[columns].copy()
+        df_copy.columns = self.init_label_list
+        df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
+        population_size = len(self.population)
+        ax.set_ylim([0, population_size])
+        plt.xlabel("Generation", size=14)
+        plt.ylabel("# Individuals", size=14)
+        ax.tick_params(axis="both", which="major", labelsize=12)
+        ax.legend(loc=4, fontsize=20)
+        plt.savefig(fname=path, dpi=300)
+
+    def PlotAvgBirthPayoff(self, df, path):
+        """Plot the sub-populations' AvgBirthPayoff after a simulation of a given Moran Process."""
+        plt.figure(figsize=(14, 6))
+        ax = plt.gca()
+        ax.tick_params(width=1)
+        for axis in ["top", "bottom", "left", "right"]:
+            ax.spines[axis].set_linewidth(1)
+        cmap = plt.get_cmap("coolwarm")
+        columns = [l + "__AvgBirthPayoff" for l in self.init_label_list]
+        df_copy = df[columns].copy()
+        df_copy.columns = self.init_label_list
+        df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
+        plt.xlabel("Generation", size=14)
+        plt.ylabel("Average Birth Payoff", size=14)
+        ax.tick_params(axis="both", which="major", labelsize=12)
+        ax.legend(loc=4, fontsize=20)
+        plt.savefig(fname=path, dpi=300)
+
+    def PlotAvgDeathPayoff(self, df, path):
+        """Plot the sub-populations' AvgDeathPayoff after a simulation of a given Moran Process."""
+        plt.figure(figsize=(14, 6))
+        ax = plt.gca()
+        ax.tick_params(width=1)
+        for axis in ["top", "bottom", "left", "right"]:
+            ax.spines[axis].set_linewidth(1)
+        cmap = plt.get_cmap("coolwarm")
+        columns = [l + "__AvgDeathPayoff" for l in self.init_label_list]
+        df_copy = df[columns].copy()
+        df_copy.columns = self.init_label_list
+        df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
+        plt.xlabel("Generation", size=14)
+        plt.ylabel("Average Death Payoff", size=14)
+        ax.tick_params(axis="both", which="major", labelsize=12)
+        ax.legend(loc=4, fontsize=20)
+        plt.savefig(fname=path, dpi=300)
+
+    def PlotBirthFitness(self, df, path):
+        """Plot the sub-populations' BirthFitness after a simulation of a given Moran Process."""
+        plt.figure(figsize=(14, 6))
+        ax = plt.gca()
+        ax.tick_params(width=1)
+        for axis in ["top", "bottom", "left", "right"]:
+            ax.spines[axis].set_linewidth(1)
+        cmap = plt.get_cmap("coolwarm")
+        columns = [l + "__BirthFitness" for l in self.init_label_list]
+        df_copy = df[columns].copy()
+        df_copy.columns = self.init_label_list
+        df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
+        plt.xlabel("Generation", size=14)
+        plt.ylabel("Birth Fitness", size=14)
+        ax.tick_params(axis="both", which="major", labelsize=12)
+        ax.legend(loc=4, fontsize=20)
+        plt.savefig(fname=path, dpi=300)
+
+    def PlotDeathFitness(self, df, path):
+        """Plot the sub-populations' DeathFitness after a simulation of a given Moran Process."""
+        plt.figure(figsize=(14, 6))
+        ax = plt.gca()
+        ax.tick_params(width=1)
+        for axis in ["top", "bottom", "left", "right"]:
+            ax.spines[axis].set_linewidth(1)
+        cmap = plt.get_cmap("coolwarm")
+        columns = [l + "__DeathFitness" for l in self.init_label_list]
+        df_copy = df[columns].copy()
+        df_copy.columns = self.init_label_list
+        df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
+        plt.xlabel("Generation", size=14)
+        plt.ylabel("Death Fitness", size=14)
+        ax.tick_params(axis="both", which="major", labelsize=12)
+        ax.legend(loc=4, fontsize=20)
+        plt.savefig(fname=path, dpi=300)
+
+    def PlotEntropy(self, df, path):
+        """Plot the whole populations entropy after a simulation of a given Moran Process."""
+        plt.figure(figsize=(14, 6))
+        ax = plt.gca()
+        ax.tick_params(width=1)
+        for axis in ["top", "bottom", "left", "right"]:
+            ax.spines[axis].set_linewidth(1)
+        df["Entropy"].plot(color="black", linewidth=1.5, ax=ax, label="Entropy")
+        plt.xlabel("Generation", size=14)
+        plt.ylabel("", size=14)
+        ax.tick_params(axis="both", which="major", labelsize=12)
+        ax.legend(loc=4, fontsize=20)
+        plt.savefig(fname=path, dpi=300)
diff --git a/moranpycess/MoranProcess2D.py b/moranpycess/MoranProcess2D.py
index 8cda90d..6012f53 100644
--- a/moranpycess/MoranProcess2D.py
+++ b/moranpycess/MoranProcess2D.py
@@ -476,68 +476,68 @@ def simulate(self, generations):
 
         return log_df
 
-
-def PlotSize2D(mp, df, path):
-    """Plot the sub-populations' sizes after a simulation of a given 2D Moran Process."""
-    plt.figure(figsize=(14, 6))
-    ax = plt.gca()
-    ax.tick_params(width=1)
-    for axis in ["top", "bottom", "left", "right"]:
-        ax.spines[axis].set_linewidth(1)
-    cmap = plt.get_cmap("coolwarm")
-    columns = [l + "__size" for l in mp.init_label_list]
-    df_copy = df[columns].copy()
-    df_copy.columns = mp.init_label_list
-    df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
-    population_size = mp.population.size
-    ax.set_ylim([0, population_size])
-    plt.xlabel("Generation", size=14)
-    plt.ylabel("# Individuals", size=14)
-    ax.tick_params(axis="both", which="major", labelsize=12)
-    ax.legend(loc=4, fontsize=20)
-    plt.savefig(fname=path, dpi=300)
-
-
-def PlotEntropy2D(mp, df, path):
-    """Plot the whole populations entropy after a simulation of a given 2D Moran Process."""
-    plt.figure(figsize=(14, 6))
-    ax = plt.gca()
-    ax.tick_params(width=1)
-    for axis in ["top", "bottom", "left", "right"]:
-        ax.spines[axis].set_linewidth(1)
-    df["Entropy"].plot(color="black", linewidth=1.5, ax=ax, label="Entropy")
-    plt.xlabel("Generation", size=14)
-    plt.ylabel("", size=14)
-    ax.tick_params(axis="both", which="major", labelsize=12)
-    ax.legend(loc=4, fontsize=20)
-    plt.savefig(fname=path, dpi=300)
-
-
-def PlotPopulationSnapshot2D(mp, path):
-    """Plot the whole populations entropy after a simulation of a given 2D Moran Process."""
-    plt.figure(figsize=(10, 10))
-    ax = plt.gca()
-    ax.tick_params(width=1)
-    for axis in ["top", "bottom", "left", "right"]:
-        ax.spines[axis].set_linewidth(1)
-    plot_grid = mp.curr_grid.copy()
-    ticks_labels = []
-    for label_index in range(len(mp.init_label_list)):
-        plot_grid[plot_grid == mp.init_label_list[label_index]] = label_index
-        ticks_labels.append(mp.init_label_list[label_index])
-    plot_grid = plot_grid.astype(float)
-    cmap = plt.get_cmap(
-        "coolwarm", np.max(plot_grid) - np.min(plot_grid) + 1
-    )  # get discrete colormap
-    mat = plt.matshow(
-        plot_grid, cmap=cmap, vmin=np.min(plot_grid) - 0.5, vmax=np.max(plot_grid) + 0.5
-    )  # set limits .5 outside true range
-    cax = plt.colorbar(
-        mat, ticks=np.arange(np.min(plot_grid), np.max(plot_grid) + 1)
-    )  # tell the colorbar to tick at integers
-    cax.set_ticklabels(ticks_labels)
-    plt.ylabel("")
-    plt.yticks([])
-    plt.xlabel("")
-    plt.xticks([])
-    plt.savefig(fname=path, dpi=300)
+    def PlotSize2D(self, df, path):
+        """Plot the sub-populations' sizes after a simulation of a given 2D Moran Process."""
+        plt.figure(figsize=(14, 6))
+        ax = plt.gca()
+        ax.tick_params(width=1)
+        for axis in ["top", "bottom", "left", "right"]:
+            ax.spines[axis].set_linewidth(1)
+        cmap = plt.get_cmap("coolwarm")
+        columns = [l + "__size" for l in self.init_label_list]
+        df_copy = df[columns].copy()
+        df_copy.columns = self.init_label_list
+        df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
+        population_size = self.population.size
+        ax.set_ylim([0, population_size])
+        plt.xlabel("Generation", size=14)
+        plt.ylabel("# Individuals", size=14)
+        ax.tick_params(axis="both", which="major", labelsize=12)
+        ax.legend(loc=4, fontsize=20)
+        plt.savefig(fname=path, dpi=300)
+
+    def PlotEntropy2D(self, df, path):
+        """Plot the whole populations entropy after a simulation of a given 2D Moran Process."""
+        plt.figure(figsize=(14, 6))
+        ax = plt.gca()
+        ax.tick_params(width=1)
+        for axis in ["top", "bottom", "left", "right"]:
+            ax.spines[axis].set_linewidth(1)
+        df["Entropy"].plot(color="black", linewidth=1.5, ax=ax, label="Entropy")
+        plt.xlabel("Generation", size=14)
+        plt.ylabel("", size=14)
+        ax.tick_params(axis="both", which="major", labelsize=12)
+        ax.legend(loc=4, fontsize=20)
+        plt.savefig(fname=path, dpi=300)
+
+    def PlotPopulationSnapshot2D(self, path):
+        """Plot the whole populations entropy after a simulation of a given 2D Moran Process."""
+        plt.figure(figsize=(10, 10))
+        ax = plt.gca()
+        ax.tick_params(width=1)
+        for axis in ["top", "bottom", "left", "right"]:
+            ax.spines[axis].set_linewidth(1)
+        plot_grid = self.curr_grid.copy()
+        ticks_labels = []
+        for label_index in range(len(self.init_label_list)):
+            plot_grid[plot_grid == self.init_label_list[label_index]] = label_index
+            ticks_labels.append(self.init_label_list[label_index])
+        plot_grid = plot_grid.astype(float)
+        cmap = plt.get_cmap(
+            "coolwarm", np.max(plot_grid) - np.min(plot_grid) + 1
+        )  # get discrete colormap
+        mat = plt.matshow(
+            plot_grid,
+            cmap=cmap,
+            vmin=np.min(plot_grid) - 0.5,
+            vmax=np.max(plot_grid) + 0.5,
+        )  # set limits .5 outside true range
+        cax = plt.colorbar(
+            mat, ticks=np.arange(np.min(plot_grid), np.max(plot_grid) + 1)
+        )  # tell the colorbar to tick at integers
+        cax.set_ticklabels(ticks_labels)
+        plt.ylabel("")
+        plt.yticks([])
+        plt.xlabel("")
+        plt.xticks([])
+        plt.savefig(fname=path, dpi=300)
diff --git a/moranpycess/MoranProcess3D.py b/moranpycess/MoranProcess3D.py
index 8754103..2036c52 100644
--- a/moranpycess/MoranProcess3D.py
+++ b/moranpycess/MoranProcess3D.py
@@ -591,38 +591,36 @@ def simulate(self, generations):
 
         return log_df
 
-
-def PlotSize3D(mp, df, path):
-    """Plot the sub-populations' sizes after a simulation of a given 3D Moran Process."""
-    plt.figure(figsize=(14, 6))
-    ax = plt.gca()
-    ax.tick_params(width=1)
-    for axis in ["top", "bottom", "left", "right"]:
-        ax.spines[axis].set_linewidth(1)
-    cmap = plt.get_cmap("coolwarm")
-    columns = [l + "__size" for l in mp.init_label_list]
-    df_copy = df[columns].copy()
-    df_copy.columns = mp.init_label_list
-    df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
-    population_size = mp.population.size
-    ax.set_ylim([0, population_size])
-    plt.xlabel("Generation", size=14)
-    plt.ylabel("# Individuals", size=14)
-    ax.tick_params(axis="both", which="major", labelsize=12)
-    ax.legend(loc=4, fontsize=20)
-    plt.savefig(fname=path, dpi=300)
-
-
-def PlotEntropy3D(mp, df, path):
-    """Plot the whole populations entropy after a simulation of a given 3D Moran Process."""
-    plt.figure(figsize=(14, 6))
-    ax = plt.gca()
-    ax.tick_params(width=1)
-    for axis in ["top", "bottom", "left", "right"]:
-        ax.spines[axis].set_linewidth(1)
-    df["Entropy"].plot(color="black", linewidth=1.5, ax=ax, label="Entropy")
-    plt.xlabel("Generation", size=14)
-    plt.ylabel("", size=14)
-    ax.tick_params(axis="both", which="major", labelsize=12)
-    ax.legend(loc=4, fontsize=20)
-    plt.savefig(fname=path, dpi=300)
+    def PlotSize3D(self, df, path):
+        """Plot the sub-populations' sizes after a simulation of a given 3D Moran Process."""
+        plt.figure(figsize=(14, 6))
+        ax = plt.gca()
+        ax.tick_params(width=1)
+        for axis in ["top", "bottom", "left", "right"]:
+            ax.spines[axis].set_linewidth(1)
+        cmap = plt.get_cmap("coolwarm")
+        columns = [l + "__size" for l in self.init_label_list]
+        df_copy = df[columns].copy()
+        df_copy.columns = self.init_label_list
+        df_copy.plot(linewidth=1.5, ax=ax, cmap=cmap)
+        population_size = self.population.size
+        ax.set_ylim([0, population_size])
+        plt.xlabel("Generation", size=14)
+        plt.ylabel("# Individuals", size=14)
+        ax.tick_params(axis="both", which="major", labelsize=12)
+        ax.legend(loc=4, fontsize=20)
+        plt.savefig(fname=path, dpi=300)
+
+    def PlotEntropy3D(self, df, path):
+        """Plot the whole populations entropy after a simulation of a given 3D Moran Process."""
+        plt.figure(figsize=(14, 6))
+        ax = plt.gca()
+        ax.tick_params(width=1)
+        for axis in ["top", "bottom", "left", "right"]:
+            ax.spines[axis].set_linewidth(1)
+        df["Entropy"].plot(color="black", linewidth=1.5, ax=ax, label="Entropy")
+        plt.xlabel("Generation", size=14)
+        plt.ylabel("", size=14)
+        ax.tick_params(axis="both", which="major", labelsize=12)
+        ax.legend(loc=4, fontsize=20)
+        plt.savefig(fname=path, dpi=300)
diff --git a/moranpycess/__init__.py b/moranpycess/__init__.py
index 05705e5..a1dcf59 100644
--- a/moranpycess/__init__.py
+++ b/moranpycess/__init__.py
@@ -16,11 +16,5 @@
 # imports
 from .Individual import Individual
 from .MoranProcess import MoranProcess
-from .MoranProcess import PlotSize
-from .MoranProcess import PlotAvgBirthPayoff, PlotAvgDeathPayoff
-from .MoranProcess import PlotBirthFitness, PlotDeathFitness
-from .MoranProcess import PlotEntropy
 from .MoranProcess2D import MoranProcess2D
-from .MoranProcess2D import PlotSize2D, PlotEntropy2D, PlotPopulationSnapshot2D
 from .MoranProcess3D import MoranProcess3D
-from .MoranProcess3D import PlotSize3D, PlotEntropy3D


Test patch:
diff --git a/tests/unit/MoranProcess.py b/tests/unit/MoranProcess.py
index 4f55bb8..7c8acd1 100644
--- a/tests/unit/MoranProcess.py
+++ b/tests/unit/MoranProcess.py
@@ -290,20 +290,12 @@ def test_plots(self):
         random.seed(0)
         simulation = mp.simulate(generations=25000)
         # test the plotting:
-        moranpycess.PlotSize(mp, simulation, "./tests/output/PD_size.png")
-        moranpycess.PlotAvgBirthPayoff(
-            mp, simulation, "./tests/output/PD_AvgBirthPayoff.png"
-        )
-        moranpycess.PlotAvgDeathPayoff(
-            mp, simulation, "./tests/output/PD_AvgDeathPayoff.png"
-        )
-        moranpycess.PlotBirthFitness(
-            mp, simulation, "./tests/output/PD_BirthFitness.png"
-        )
-        moranpycess.PlotDeathFitness(
-            mp, simulation, "./tests/output/PD_DeathFitness.png"
-        )
-        moranpycess.PlotEntropy(mp, simulation, "./tests/output/PD_Entropy.png")
+        mp.PlotSize(simulation, "./tests/output/PD_size.png")
+        mp.PlotAvgBirthPayoff(simulation, "./tests/output/PD_AvgBirthPayoff.png")
+        mp.PlotAvgDeathPayoff(simulation, "./tests/output/PD_AvgDeathPayoff.png")
+        mp.PlotBirthFitness(simulation, "./tests/output/PD_BirthFitness.png")
+        mp.PlotDeathFitness(simulation, "./tests/output/PD_DeathFitness.png")
+        mp.PlotEntropy(simulation, "./tests/output/PD_Entropy.png")
         assert True  # mark that no error was raised before
 
     def test_MoranProcessWithTransitionMatrix(self):
diff --git a/tests/unit/MoranProcess2D.py b/tests/unit/MoranProcess2D.py
index b9f002a..8b7cfbf 100644
--- a/tests/unit/MoranProcess2D.py
+++ b/tests/unit/MoranProcess2D.py
@@ -358,9 +358,9 @@ def test_plots2D(self):
         random.seed(0)
         simulation = mp.simulate(generations=10)
         # test the plotting:
-        moranpycess.PlotSize2D(mp, simulation, "./tests/output/2D_size.png")
-        moranpycess.PlotEntropy2D(mp, simulation, "./tests/output/2D_Entropy.png")
-        moranpycess.PlotPopulationSnapshot2D(mp, "./tests/output/2D_snapshot.png")
+        mp.PlotSize2D(simulation, "./tests/output/2D_size.png")
+        mp.PlotEntropy2D(simulation, "./tests/output/2D_Entropy.png")
+        mp.PlotPopulationSnapshot2D("./tests/output/2D_snapshot.png")
         assert True  # mark that no error was raised before
 
     def test_MoranProcess2DWithTransitionMatrix(self):
diff --git a/tests/unit/MoranProcess3D.py b/tests/unit/MoranProcess3D.py
index 001a7e7..caa5705 100644
--- a/tests/unit/MoranProcess3D.py
+++ b/tests/unit/MoranProcess3D.py
@@ -555,8 +555,8 @@ def test_plots3D(self):
         random.seed(0)
         simulation = mp.simulate(generations=10)
         # test the plotting:
-        moranpycess.PlotSize3D(mp, simulation, "./tests/output/3D_size.png")
-        moranpycess.PlotEntropy3D(mp, simulation, "./tests/output/3D_Entropy.png")
+        mp.PlotSize3D(simulation, "./tests/output/3D_size.png")
+        mp.PlotEntropy3D(simulation, "./tests/output/3D_Entropy.png")
         assert True  # mark that no error was raised before
 
     def test_MoranProcess3DWithTransitionMatrix(self):
diff --git a/tests/unit/context.py b/tests/unit/context.py
index aa2e642..f84416e 100644
--- a/tests/unit/context.py
+++ b/tests/unit/context.py
@@ -1,6 +1,7 @@
 import os
 import sys
 
+# python path modification to resolve the package properly during testing
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
 
 import moranpycess