Repository: g-walley/cegpy (language: python)

Issue:
Broken CI graphviz calls
CI Is broken due to graphviz calls in the Tests. Calls should be commented out, left there for when developers wish to see the graphs.
Python notebook - displaying graphs
Makes filename optional and enables displaying the graphs in Jupyter notebooks without saving to a file.

Fixes a related unit test.

Resolves #72 


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml
deleted file mode 100644
index fcc9ba1..0000000
--- a/.github/workflows/docker-image.yml
+++ /dev/null
@@ -1,18 +0,0 @@
-name: Docker Image CI
-
-on:
-  push:
-    branches: [ main, dev]
-  pull_request:
-    branches: [ main ]
-
-jobs:
-
-  build:
-
-    runs-on: ubuntu-latest
-
-    steps:
-    - uses: actions/checkout@v2
-    - name: Build the Docker image
-      run: docker build . --file Dockerfile --tag my-image-name:$(date +%s)
diff --git a/.github/workflows/python-app.yaml b/.github/workflows/python-app.yaml
index 2540b3a..dde2e9e 100644
--- a/.github/workflows/python-app.yaml
+++ b/.github/workflows/python-app.yaml
@@ -7,7 +7,7 @@ on:
   push:
     branches: [ main, dev ]
   pull_request:
-    branches: [ main ]
+    branches: [ main, dev ]
 
 jobs:
   build:
@@ -26,12 +26,6 @@ jobs:
         pip install wheel
         pip install flake8 pytest pandas pydotplus pathlib IPython openpyxl colorutils networkx
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-    - name: Lint with flake8
-      run: |
-        # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
     - name: Test with pytest
       run: |
-        pytest tests
+        pytest src/tests
diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml
deleted file mode 100644
index 2e77ea0..0000000
--- a/.github/workflows/python-app.yml
+++ /dev/null
@@ -1,37 +0,0 @@
-# This workflow will install Python dependencies, run tests and lint with a single version of Python
-# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
-
-name: Python application
-
-on:
-  push:
-    branches: [ ** ]
-  pull_request:
-    branches: [ ** ]
-
-jobs:
-  build:
-
-    runs-on: ubuntu-latest
-
-    steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.9
-      uses: actions/setup-python@v2
-      with:
-        python-version: 3.9
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install wheel
-        pip install flake8 pytest pandas pydotplus pathlib IPython openpyxl colorutils networkx
-        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-    - name: Lint with flake8
-      run: |
-        # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
-    - name: Test with pytest
-      run: |
-        pytest tests
diff --git a/.gitignore b/.gitignore
index 303f936..8c6c7fb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -21,16 +21,15 @@ st_figure.png
 tic-tac-toe.py
 pdp_experiment.py
 gv_playground.gv
-.gitignore
 data/breast-cancer.xlsx
 data/Falls_Data_test.xlsx
 data/medical_dm_modified_test.xlsx
 data/medical_dm_posterior_prior.xlsx
-.gitignore
 data/voting.xlsx
 scripts/docker/run_example_docker.sh
 data/tic-tac-toe.xlsx
 *.pdf
 .tox/
 .coverage
-.coverage.*
\ No newline at end of file
+.coverage.*
+cegpy-binder
\ No newline at end of file
diff --git a/CHANGELOG b/CHANGELOG
index f5841e4..826c108 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -7,4 +7,10 @@
 
 -- 0.0.1
 - Event Trees can be generated from compatible datasets
-- Event Tree can be printed to a .png file for human readability 
+- Event Tree can be printed to a .png file for human readability
+
+-- 0.2
+- Var order can now be set at creation of event tree.
+- Stage colours can now be set manually
+- DOT graph has been made accessable for customisation of style.
+- Minor improvements to AHC.
diff --git a/data/Asym.csv b/data/Asym.csv
new file mode 100644
index 0000000..a3930d7
--- /dev/null
+++ b/data/Asym.csv
@@ -0,0 +1,1001 @@
+X1,X2,X3,X4
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,0
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,0,1
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,0
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,0,1,1
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,0,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+0,1,1,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,0
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,0,1
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,0,1,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,0
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,0,1
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,0
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
+1,1,1,1
diff --git a/docs/make.bat b/docs/make.bat
index 6fcf05b..747ffb7 100644
--- a/docs/make.bat
+++ b/docs/make.bat
@@ -10,8 +10,6 @@ if "%SPHINXBUILD%" == "" (
 set SOURCEDIR=source
 set BUILDDIR=build
 
-if "%1" == "" goto help
-
 %SPHINXBUILD% >NUL 2>NUL
 if errorlevel 9009 (
 	echo.
@@ -25,6 +23,8 @@ if errorlevel 9009 (
 	exit /b 1
 )
 
+if "%1" == "" goto help
+
 %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
 goto end
 
diff --git a/docs/source/api.rst b/docs/source/api.rst
new file mode 100644
index 0000000..dffcb1d
--- /dev/null
+++ b/docs/source/api.rst
@@ -0,0 +1,10 @@
+API
+===
+
+There are four main classes to interact with in cegpy. These are:
+
+.. toctree::
+    event_tree
+    staged_tree
+    chain_event_graph
+    chain_event_graph_reducer
diff --git a/docs/source/chain_event_graph.rst b/docs/source/chain_event_graph.rst
new file mode 100644
index 0000000..5b3e996
--- /dev/null
+++ b/docs/source/chain_event_graph.rst
@@ -0,0 +1,11 @@
+ChainEventGraph
+===============
+.. autoclass:: cegpy.ChainEventGraph
+    :members:
+        sink,
+        root,
+        stages,
+        path_list,
+        generate,
+        dot_graph,
+        create_figure
diff --git a/docs/source/chain_event_graph_reducer.rst b/docs/source/chain_event_graph_reducer.rst
new file mode 100644
index 0000000..d5b3c8e
--- /dev/null
+++ b/docs/source/chain_event_graph_reducer.rst
@@ -0,0 +1,27 @@
+ChainEventGraphReducer
+======================
+.. autoclass:: cegpy.ChainEventGraphReducer
+    :members:
+        certain_edges,
+        uncertain_edges,
+        certain_nodes,
+        uncertain_nodes,
+        paths,
+        graph,
+        clear_all_evidence,
+        add_certain_edge,
+        remove_certain_edge,
+        add_certain_edge_list,
+        remove_certain_edge_list,
+        add_uncertain_edge_set,
+        remove_uncertain_edge_set,
+        add_uncertain_edge_set_list,
+        remove_uncertain_edge_set_list,
+        add_certain_node,
+        remove_certain_node,
+        add_certain_node_set,
+        remove_certain_node_set,
+        add_uncertain_node_set,
+        remove_uncertain_node_set,
+        add_uncertain_node_set_list,
+        remove_uncertain_node_set_list,
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 07fbba7..9af9856 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -1,56 +1,33 @@
 # Configuration file for the Sphinx documentation builder.
 #
-# This file only contains a selection of the most common options. For a full
-# list see the documentation:
+# For the full list of built-in configuration values, see the documentation:
 # https://www.sphinx-doc.org/en/master/usage/configuration.html
 
-# -- Path setup --------------------------------------------------------------
-
-# If extensions (or modules to document with autodoc) are in another directory,
-# add these directories to sys.path here. If the directory is relative to the
-# documentation root, use os.path.abspath to make it absolute, like shown here.
-#
-# import os
-# import sys
-# sys.path.insert(0, os.path.abspath('.'))
-
-
 # -- Project information -----------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
-project = 'cegpy'
-copyright = '2021, Gareth Walley, Aditi Shenvi'
-author = 'Gareth Walley, Aditi Shenvi'
-
-# The full version, including alpha/beta/rc tags
-release = 'x.x.x'
-
+project = "Cegpy"
+copyright = "2022, Gareth Walley, Aditi Shenvi"
+author = "Gareth Walley, Aditi Shenvi"
+release = "1.0"
 
 # -- General configuration ---------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
-# Add any Sphinx extension module names here, as strings. They can be
-# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
-# ones.
 extensions = [
-    'sphinx.ext.autodoc'
+    "sphinx.ext.duration",
+    "sphinx.ext.autodoc",
 ]
 
-# Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
-
-# List of patterns, relative to source directory, that match files and
-# directories to ignore when looking for source files.
-# This pattern also affects html_static_path and html_extra_path.
+templates_path = ["_templates"]
 exclude_patterns = []
 
 
 # -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 
-# The theme to use for HTML and HTML Help pages.  See the documentation for
-# a list of builtin themes.
-#
-html_theme = 'alabaster'
+html_theme = "sphinx_rtd_theme"
+html_static_path = ["_static"]
 
-# Add any paths that contain custom static files (such as style sheets) here,
-# relative to this directory. They are copied after the builtin static files,
-# so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
+# -- Options for autodoc -----------------------------------------------------
+autodoc_member_order = "bysource"
diff --git a/docs/source/event_tree.rst b/docs/source/event_tree.rst
new file mode 100644
index 0000000..c7bb5e3
--- /dev/null
+++ b/docs/source/event_tree.rst
@@ -0,0 +1,14 @@
+EventTree
+=========
+.. autoclass:: cegpy.EventTree
+    :members:
+        create_figure,
+        dot_event_graph,
+        root,
+        variables,
+        sampling_zeros,
+        situations,
+        leaves,
+        edge_counts,
+        categories_per_variable
+
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 99a595a..ac61d87 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -1,17 +1,26 @@
 .. cegpy documentation master file, created by
-   sphinx-quickstart on Sat Oct  2 17:20:26 2021.
+   sphinx-quickstart on Tue Aug 16 07:45:48 2022.
    You can adapt this file completely to your liking, but it should at least
    contain the root `toctree` directive.
 
-Welcome to cegpy's documentation!
+Welcome to Cegpy's documentation!
 =================================
+**Cegpy** (/segpaɪ/) is a Python package for working with Chain Event Graphs. It supports learning the graphical structure of a Chain Event Graph from data, encoding of parametric and structural priors, estimating its parameters, and performing inference.
 
-.. toctree::
-   :maxdepth: 2
-   :caption: Contents:
+It is built on top of the Python network modelling package NetworkX.
+
+.. note::
+
+   Check out the :doc:`usage` section for further information, including how to :ref:`install <installation>` the project.
 
+.. note::
 
+   After a year of development, we are releasing v1.0!
 
+.. toctree::
+
+   usage
+   api
 
 Indices and tables
 ==================
@@ -19,24 +28,3 @@ Indices and tables
 * :ref:`genindex`
 * :ref:`modindex`
 * :ref:`search`
-
-=================
-Event Tree Class
-=================
-
-.. automodule:: cegpy.trees.event
-   :members:
-
-=================
-Staged Tree Class
-=================
-
-.. automodule:: cegpy.trees.staged
-   :members:
-
-=======================
-Chain Event Graph Class
-=======================
-
-.. automodule:: cegpy.graphs.ceg
-   :members:
diff --git a/docs/source/staged_tree.rst b/docs/source/staged_tree.rst
new file mode 100644
index 0000000..b9dbb4f
--- /dev/null
+++ b/docs/source/staged_tree.rst
@@ -0,0 +1,16 @@
+StagedTree
+==========
+.. autoclass:: cegpy.StagedTree
+    :members:
+        create_figure,
+        dot_event_graph,
+        calculate_AHC_transitions,
+        create_event_tree_figure,
+        prior,
+        prior_list,
+        posterior,
+        posterior_list,
+        alpha,
+        hyperstage,
+        edge_countset,
+        ahc_output
diff --git a/docs/source/usage.rst b/docs/source/usage.rst
new file mode 100644
index 0000000..3bde503
--- /dev/null
+++ b/docs/source/usage.rst
@@ -0,0 +1,21 @@
+Usage
+=====
+
+.. _installation:
+
+Installation
+------------
+
+To use cegpy, first install it using pip:
+
+.. code-block:: console
+
+    $ pip install cegpy
+
+
+Worked Examples
+---------------
+
+Examples of typical use cases of the package can be found `in a binder on Github`_.
+
+.. _in a binder on Github: https://github.com/peterrhysstrong/cegpy-binder
diff --git a/example.py b/example.py
index 0ae3ddb..4be7366 100644
--- a/example.py
+++ b/example.py
@@ -5,8 +5,7 @@
 from pathlib import Path
 import pandas as pd
 from datetime import datetime
-import pickle
-# from time import sleep
+
 
 def create_path(filename, add_time=False, filetype='png'):
     dt_string = ''
@@ -29,22 +28,15 @@ def create_path(filename, add_time=False, filetype='png'):
 med_et = EventTree(dataframe=df)
 med_st = StagedTree(dataframe=df)
 
-et_fig_path = create_path('out/medical_dm_event_tree', True, 'pdf')
-st_fig_path = create_path('out/medical_dm_staged_tree', True, 'pdf')
+et_fig_path = create_path('out/medical_dm_event_tree', False, 'pdf')
+st_fig_path = create_path('out/medical_dm_staged_tree', False, 'pdf')
+et_from_st_fig_path = create_path('out/medical_dm_event_tree_from_staged', False, 'pdf')
 med_et.create_figure(et_fig_path)
 med_st.calculate_AHC_transitions()
 med_st.create_figure(st_fig_path)
+med_st.create_figure(et_from_st_fig_path, staged=False)
 
-# pickle_path = create_path('out/med_st', True, 'pickle')
-# with open(pickle_path, 'wb') as f:
-#     pickle.dump(med_st, f, pickle.HIGHEST_PROTOCOL)
-
-# with open(pickle_path, 'rb') as f:
-#     med_st2 = pickle.load(f)
-#     med_st2.create_figure(st_fig_path)
-
-
-# # FALLS
+# FALLS
 df_path = Path(__file__).resolve().parent.joinpath(
     'data/Falls_Data.xlsx')
 logger.info(str(df_path))
@@ -52,22 +44,24 @@ def create_path(filename, add_time=False, filetype='png'):
 falls_et = EventTree(dataframe=df)
 falls_st = StagedTree(dataframe=df)
 colours = [
-    '#8dd3c7', '#ffffb3','#bebada', '#fb8072', 
+    '#8dd3c7', '#ffffb3', '#bebada', '#fb8072',
     '#80b1d3', '#fdb462', '#b3de69', '#fccde5'
 ]
-et_fig_path = create_path('out/falls_event_tree', True, 'pdf')
-st_fig_path = create_path('out/falls_staged_tree', True, 'pdf')
+et_fig_path = create_path('out/falls_event_tree', False, 'pdf')
+st_fig_path = create_path('out/falls_staged_tree', False, 'pdf')
+et_from_st_fig_path = create_path('out/falls_event_tree_from_staged', False, 'pdf')
 falls_et.create_figure(et_fig_path)
 falls_st.calculate_AHC_transitions(colour_list=colours)
 falls_st.create_figure(st_fig_path)
+falls_st.create_figure(et_from_st_fig_path, staged=False)
 
-# get the dot graph 
-falls_st_dot = falls_st.dot_graph
+# get the dot staged tree graph
+falls_st_dot = falls_st.dot_staged_graph
 # modify nodes and edges
 falls_st_dot.get_edge("s1", "s5")[0].set_style("dotted")
 falls_st_dot.get_edge("s1", "s6")[0].set_style("dotted")
 falls_st_dot.get_node("s1")[0].set_shape("square")
 # save as a pdf
-falls_st_dot.write("out/falls_modified_staged_tree.pdf", format="pdf") 
+falls_st_dot.write("out/falls_modified_staged_tree.pdf", format="pdf")
 # or save as a dot
-falls_st_dot.write("out/falls_modified_staged_tree.dot", format="dot") 
\ No newline at end of file
+falls_st_dot.write("out/falls_modified_staged_tree.dot", format="dot")
\ No newline at end of file
diff --git a/scripts/build_cegpy.sh b/scripts/build_cegpy.sh
index 5942bfb..190f205 100644
--- a/scripts/build_cegpy.sh
+++ b/scripts/build_cegpy.sh
@@ -1,8 +1,6 @@
 # Move old distribution to the old folder.
-mv dist/* .dist_old/
-
+mkdir --parents .dist_old/; mv dist/* $_
 # build the package
-python3 -m build
-
+python -m build
 # Upload new package
 twine upload dist/*
\ No newline at end of file
diff --git a/setup.cfg b/setup.cfg
index 9f23d99..1e80833 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = cegpy
-version = 0.2
+version = attr: cegpy.__version__
 author = Gareth Walley
 author_email = gareth@walley.me
 description = Package for creating, manipulating, and analysing Chain Event Graphs. Built on networkx.
@@ -31,5 +31,8 @@ install_requires =
 
 license_files=
     LICENSE
+
 [options.packages.find]
 where = src
+exclude =
+    tests*
\ No newline at end of file
diff --git a/src/cegpy/__init__.py b/src/cegpy/__init__.py
index ede1eaf..7438762 100644
--- a/src/cegpy/__init__.py
+++ b/src/cegpy/__init__.py
@@ -8,13 +8,14 @@
 See documentation of more information.
 
 """
-from .utilities.util import *
-from .trees.event import *
-from .trees.staged import *
-from .graphs.ceg import *
 import logging
-# from graphs import ctceg
 
+from cegpy.graphs._ceg import ChainEventGraph
+from cegpy.graphs._ceg_reducer import ChainEventGraphReducer
+from cegpy.trees._event import EventTree
+from cegpy.trees._staged import StagedTree
 
 logging.basicConfig(level=logging.ERROR)
-logger = logging.getLogger('cegpy')
+logger = logging.getLogger("cegpy")
+
+__version__ = "1.0.0"
diff --git a/src/cegpy/graphs/__init__.py b/src/cegpy/graphs/__init__.py
index f373bcb..e69de29 100644
--- a/src/cegpy/graphs/__init__.py
+++ b/src/cegpy/graphs/__init__.py
@@ -1,1 +0,0 @@
-from .ceg import ChainEventGraph
\ No newline at end of file
diff --git a/src/cegpy/graphs/_ceg.py b/src/cegpy/graphs/_ceg.py
new file mode 100644
index 0000000..a2375cc
--- /dev/null
+++ b/src/cegpy/graphs/_ceg.py
@@ -0,0 +1,441 @@
+"""Chain Event Graph"""
+
+from collections import defaultdict
+from copy import deepcopy
+import itertools as it
+import logging
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+)
+import pydotplus as pdp
+import networkx as nx
+from IPython.display import Image
+from IPython import get_ipython
+
+from cegpy.utilities._util import generate_filename_and_mkdir
+from cegpy.trees._staged import StagedTree
+
+logger = logging.getLogger("cegpy.chain_event_graph")
+
+
+class CegAlreadyGenerated(Exception):
+    """Raised when a CEG is generated twice."""
+
+
+class ChainEventGraph(nx.MultiDiGraph):
+    """
+    Representation of a Chain Event Graph.
+
+    A Chain Event Graph reduces a staged tree.
+
+    The class is an extension of NetworkX MultiDiGraph.
+
+    :param staged_tree: A staged tree object where the stages have been calculated.
+    :type staged_tree: StagedTree
+
+    :param node_prefix: The prefix that is used for the nodes in the Chain Event Graph.
+        Default = "w"
+    :type node_prefix: str
+
+    :param generate: Automatically generate the Chain Event Graph upon creation of the
+        object. Default = True.
+    :type generate: bool
+    """
+
+    _edge_attributes: List = ["count", "prior", "posterior", "probability"]
+
+    sink_suffix: str = "&infin;"
+    node_prefix: str
+    generated: bool = False
+
+    def __init__(
+        self,
+        staged_tree: Optional[StagedTree] = None,
+        node_prefix: str = "w",
+        generate: bool = True,
+    ):
+        self.ahc_output = deepcopy(getattr(staged_tree, "ahc_output", {}))
+        super().__init__(staged_tree, attr={})
+        self.node_prefix = node_prefix
+        self._stages = {}
+        self.staged_root = staged_tree.root if staged_tree is not None else None
+
+        if generate and staged_tree is not None:
+            self.generate()
+
+    @property
+    def sink(self) -> str:
+        """
+        :return: Sink node name
+        :rtype: str
+        """
+        return f"{self.node_prefix}_infinity"
+
+    @property
+    def root(self) -> str:
+        """
+        :return: Root node name
+        :rtype: str
+        """
+        return f"{self.node_prefix}0"
+
+    @property
+    def stages(self) -> Mapping[str, Set[str]]:
+        """
+        :return: Mapping of stages to constituent nodes.
+        :rtype: Mapping[str, Set[str]]
+        """
+        node_stages = dict(self.nodes(data="stage", default=None))
+        stages = defaultdict(set)
+        for node, stage in node_stages.items():
+            stages[stage].add(node)
+
+        return stages
+
+    @property
+    def path_list(self) -> List[List[Tuple[str]]]:
+        """
+        :return: All the paths through the CEG, as a list of lists of edge tuples.
+        :rtype: List[List[Tuple[str]]]"""
+        path_list: List[Tuple[str]] = list(
+            nx.all_simple_edge_paths(self, self.root, self.sink)
+        )
+        return path_list
+
+    def generate(self) -> None:
+        """
+        Identifies the positions i.e. the nodes of the CEG and the edges of the CEG
+        along with their edge labels and edge counts. Here we use the
+        algorithm from our paper with the optimal stopping time.
+        """
+        if self.generated:
+            raise CegAlreadyGenerated("CEG has already been generated.")
+
+        if self.ahc_output is None or self.ahc_output == {}:
+            raise ValueError(
+                "There is no AHC output in your StagedTree. "
+                "Run StagedTree.calculate_AHC_transitions() first."
+            )
+
+        # rename root node:
+        nx.relabel_nodes(self, {self.staged_root: self.root}, copy=False)
+        self._trim_leaves_from_graph()
+        self._update_distances_to_sink()
+        self._backwards_construction(self._gen_nodes_with_increasing_distance(start=1))
+        self._relabel_nodes()
+        self.generated = True
+
+    def _backwards_construction(self, node_generator: Iterable[str]) -> None:
+        """Working backwards from the sink, the algorithm constructs the CEG."""
+        next_set_of_nodes: List = next(node_generator)
+
+        while next_set_of_nodes != [self.root]:
+            nodes_to_merge = set()
+            while len(next_set_of_nodes) > 1:
+                node_1 = next_set_of_nodes.pop(0)
+                for node_2 in next_set_of_nodes:
+                    mergeable = self._check_nodes_can_be_merged(node_1, node_2)
+                    if mergeable:
+                        nodes_to_merge.add((node_1, node_2))
+
+            if nodes_to_merge:
+                self._merge_nodes(nodes_to_merge)
+
+            try:
+                next_set_of_nodes: List = next(node_generator)
+            except StopIteration:
+                break
+
+    def _merge_nodes(self, nodes_to_merge: Set):
+        """nodes to merge should be a set of 2 element tuples"""
+        temp_1 = "temp_1"
+        temp_2 = "temp_2"
+        while nodes_to_merge:
+            nodes = nodes_to_merge.pop()
+            new_node = nodes[0]
+            # Copy nodes to temp nodes
+            node_map = {nodes[0]: temp_1, nodes[1]: temp_2}
+            nx.relabel_nodes(self, node_map, copy=False)
+            self.add_node(new_node)
+
+            edges_to_remove = self._merge_and_add_edges(
+                new_node,
+                temp_1,
+                temp_2,
+            )
+            self.remove_edges_from(edges_to_remove)
+            nx.relabel_nodes(
+                G=self, mapping={temp_1: new_node, temp_2: new_node}, copy=False
+            )
+
+            # Some nodes have been removed, we need to update the
+            # mergeable list to point to new nodes if required
+            temp_list = list(nodes_to_merge)
+            for pair in temp_list:
+                if nodes[1] in pair:
+                    new_pair = (
+                        # the other node of the pair
+                        pair[pair.index(nodes[1]) - 1],
+                        # the new node it will be merged to
+                        new_node,
+                    )
+                    nodes_to_merge.remove(pair)
+                    if new_pair[0] != new_pair[1]:
+                        nodes_to_merge.add(new_pair)
+
+    def dot_graph(self, edge_info: str = "probability") -> pdp.Dot:
+        """Returns Dot graph representation of the CEG.
+        :param edge_info: Optional - Chooses which summary measure to be displayed
+        on edges. Defaults to "count".
+        Options: ["count", "prior", "posterior", "probability"]
+
+        :type edge_info: str
+        :return: A graphviz Dot representation of the graph.
+        :rtype: pydotplus.Dot"""
+        return self._generate_dot_graph(edge_info=edge_info)
+
+    def _generate_dot_graph(self, edge_info="probability"):
+        graph = pdp.Dot(graph_type="digraph", rankdir="LR")
+        if edge_info in self._edge_attributes:
+            edge_info_dict = nx.get_edge_attributes(self, edge_info)
+        else:
+            logger.warning(
+                "edge_info '%s' does not exist for the %s class. "
+                "Using the default of 'probability' values "
+                "on edges instead. For more information, see the "
+                "documentation.",
+                edge_info,
+                self.__class__.__name__,
+            )
+            edge_info_dict = nx.get_edge_attributes(self, "probability")
+
+        for (src, dst, label), attribute in edge_info_dict.items():
+            if edge_info == "count":
+                edge_details = str(label) + "\n" + str(attribute)
+            else:
+                edge_details = f"{label}\n{float(attribute):.2f}"
+
+            if dst[1:] == "_infinity":
+                dst = f"{self.node_prefix}&infin;"
+
+            graph.add_edge(
+                pdp.Edge(
+                    src,
+                    dst,
+                    label=edge_details,
+                    labelfontcolor="#009933",
+                    fontsize="10.0",
+                    color="black",
+                )
+            )
+        nodes = list(nx.topological_sort(self))
+        for node in nodes:
+            try:
+                fill_colour = self.nodes[node]["colour"]
+            except KeyError:
+                fill_colour = "white"
+            if node[1:] == "_infinity":
+                node = f"{self.node_prefix}&infin;"
+            label = "<" + node[0] + "<SUB>" + node[1:] + "</SUB>" + ">"
+            graph.add_node(
+                pdp.Node(name=node, label=label, style="filled", fillcolor=fill_colour)
+            )
+        return graph
+
+    def create_figure(
+        self,
+        filename=None,
+        edge_info: str = "probability",
+    ) -> Union[Image, None]:
+        """Draws the coloured chain event graph for the staged_tree.
+
+        :param filename: Optional - When provided, file is saved to the filename,
+            local to the current working directory.
+            e.g. if filename = "output/ceg.svg", the file will be saved to:
+            cwd/output/ceg.svg
+            Otherwise, if function is called inside an interactive notebook, image
+            will be displayed in the notebook, even if filename is omitted.
+
+            Supports any filetype that graphviz supports. e.g: "ceg.png" or
+            "ceg.svg" etc.
+
+        :type filename: str
+
+        :param edge_info: Optional - Chooses which summary measure to be displayed on
+            edges. Value can take: "count", "prior", "posterior", "probability"
+        :type edge_info: str
+
+        :return: The event tree Image object.
+        :rtype: IPython.display.Image or None
+        """
+        graph = self.dot_graph(edge_info=edge_info)
+        if filename is None:
+            logger.warning("No filename. Figure not saved.")
+        else:
+            filename, filetype = generate_filename_and_mkdir(filename)
+            logger.info("--- generating graph ---")
+            logger.info("--- writing %s file ---", filetype)
+            graph.write(str(filename), format=filetype)
+
+        if get_ipython() is not None:
+            logger.info("--- Exporting graph to notebook ---")
+            graph_image = Image(graph.create_png())  # pylint: disable=no-member
+        else:
+            graph_image = None
+
+        return graph_image
+
+    def _trim_leaves_from_graph(self):
+        """Trims all the leaves from the graph, and points each incoming
+        edge to the sink node."""
+        # Create new CEG sink node
+        self.add_node(self.sink, colour="lightgrey")
+        outgoing_edges = deepcopy(self.succ).items()
+        # Check to see if any nodes have no outgoing edges.
+        mapping = {}
+        for node, out_edges in outgoing_edges:
+            if not out_edges and node != self.sink:
+                mapping[node] = self.sink
+
+        nx.relabel_nodes(self, mapping, copy=False)
+
+    def _update_distances_to_sink(self) -> None:
+        """
+        Iterates through the graph until it finds the root node.
+        For each node, it determines the maximum number of edges
+        from that node to the sink node.
+        """
+        max_dist = "max_dist_to_sink"
+        self.nodes[self.sink][max_dist] = 0
+        node_queue = [self.sink]
+
+        while node_queue != [self.root]:
+            node = node_queue.pop(0)
+            for pred in self.predecessors(node):
+                max_dist_to_sink = set()
+                for succ in self.successors(pred):
+                    try:
+                        max_dist_to_sink.add(self.nodes[succ][max_dist])
+                        self.nodes[pred][max_dist] = max(max_dist_to_sink) + 1
+                    except KeyError:
+                        break
+
+                if pred not in node_queue:
+                    node_queue.append(pred)
+
+    def _gen_nodes_with_increasing_distance(self, start=0) -> list:
+        """Generates nodes that are either the same or further
+        from the sink node than the last node generated."""
+        max_dists = nx.get_node_attributes(self, "max_dist_to_sink")
+        distance_dict: Mapping[int, Iterable[str]] = {}
+        for node, distance in max_dists.items():
+            dist_list: List = distance_dict.setdefault(distance, [])
+            dist_list.append(node)
+
+        for dist in range(0, max(distance_dict) + 1):
+            nodes = distance_dict.get(dist)
+            if dist >= start and nodes is not None:
+                yield nodes
+
+    def _relabel_nodes(self):
+        """Relabels nodes whilst maintaining ordering."""
+        num_iterator = it.count(1, 1)
+        nodes_to_rename = list(self.succ[self.root].keys())
+        # first, relabel the successors of this node
+        node_mapping = {}
+        while nodes_to_rename:
+            for node in nodes_to_rename.copy():
+                node_mapping[node] = f"{self.node_prefix}{next(num_iterator)}"
+                for succ in self.succ[node].keys():
+                    if succ != self.sink and succ not in nodes_to_rename:
+                        nodes_to_rename.append(succ)
+                nodes_to_rename.remove(node)
+
+        nx.relabel_nodes(self, node_mapping, copy=False)
+
+    def _merge_and_add_edges(
+        self,
+        new_node: str,
+        node_1: str,
+        node_2: str,
+    ) -> List[Tuple]:
+        """Merges outgoing edges of two nodes so that the two nodes can be
+        merged."""
+        old_edges_to_remove = []
+        for succ, t1_edge_dict in self.succ[node_1].items():
+            edge_labels = list(t1_edge_dict.keys())
+            while edge_labels:
+                label = edge_labels.pop(0)
+                n1_edge_data = t1_edge_dict[label]
+                n2_edge_data = self.succ[node_2][succ][label]
+
+                new_edge_data = _merge_edge_data(
+                    edge_1=n1_edge_data,
+                    edge_2=n2_edge_data,
+                )
+                self.add_edge(
+                    u_for_edge=new_node,
+                    v_for_edge=succ,
+                    key=label,
+                    **new_edge_data,
+                )
+                old_edges_to_remove.extend(
+                    [(node_1, succ, label), (node_2, succ, label)]
+                )
+
+        return old_edges_to_remove
+
+    def _check_nodes_can_be_merged(self, node_1, node_2) -> bool:
+        """Determine if the two nodes are able to be merged."""
+        have_same_successor_nodes = set(self.adj[node_1].keys()) == set(
+            self.adj[node_2].keys()
+        )
+
+        if have_same_successor_nodes:
+            have_same_outgoing_edges = True
+            n1_adj = self.succ[node_1]
+            for succ_node in list(n1_adj.keys()):
+                n1_edges = self.succ[node_1][succ_node]
+                n2_edges = self.succ[node_2][succ_node]
+
+                n2_edge_labels = list(n2_edges.keys())
+
+                for label in n1_edges.keys():
+                    if label not in n2_edge_labels:
+                        have_same_outgoing_edges &= False
+                        break
+                    have_same_outgoing_edges &= True
+        else:
+            have_same_outgoing_edges = False
+
+        try:
+            in_same_stage = self.nodes[node_1]["stage"] == self.nodes[node_2]["stage"]
+        except KeyError:
+            in_same_stage = False
+
+        return in_same_stage and (
+            have_same_successor_nodes and have_same_outgoing_edges
+        )
+
+
+def _merge_edge_data(
+    edge_1: Dict[str, Any],
+    edge_2: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Merges the counts, priors, and posteriors of two edges."""
+    new_edge_data = {}
+    edge = edge_1 if len(edge_1) > len(edge_2) else edge_2
+    for key in edge:
+        if key == "probability":
+            new_edge_data[key] = edge_1.get(key, 1)
+        else:
+            new_edge_data[key] = edge_1.get(key, 0) + edge_2.get(key, 0)
+    return new_edge_data
diff --git a/src/cegpy/graphs/_ceg_reducer.py b/src/cegpy/graphs/_ceg_reducer.py
new file mode 100644
index 0000000..4725ae8
--- /dev/null
+++ b/src/cegpy/graphs/_ceg_reducer.py
@@ -0,0 +1,497 @@
+"""Reduced Chain Event Graph"""
+
+from copy import deepcopy
+from typing import Iterable, List, Set, Tuple
+import networkx as nx
+
+from cegpy.graphs._ceg import ChainEventGraph
+
+
+class ChainEventGraphReducer:
+    """
+    Reduces Chain Event Graphs given certain and/or uncertain evidence.
+
+    :param ceg: Chain event graph object to reduce.
+    :type ceg: ChainEventGraph
+    """
+
+    # pylint: disable=too-many-instance-attributes,too-many-public-methods
+    _path_list: List[List[Tuple[str]]]
+
+    def __init__(self, ceg: ChainEventGraph):
+        self._ceg = ceg
+        self._path_list = ceg.path_list
+        self._certain_edges = []
+        self._uncertain_edges = []
+        self._certain_nodes = set()
+        self._uncertain_nodes = []
+        self._edges = set()
+        self._vertices = set()
+
+    def __repr__(self) -> str:
+        return (
+            f"Evidence(certain_edges={str(self.certain_edges)}, "
+            f"certain_nodes={str(self.certain_nodes)}, "
+            f"uncertain_edges={str(self.uncertain_edges)}, "
+            f"uncertain_nodes={str(self.uncertain_nodes)})"
+        )
+
+    def __str__(self) -> str:
+        """Returns human readable version of the evidence you've provided."""
+
+        def evidence_str(base_str, edges, nodes):
+            if edges == []:
+                base_str += "   Edges = []\n"
+            else:
+                base_str += "   Edges = [\n"
+                for edge in edges:
+                    base_str += f"     {str(edge)},\n"
+                base_str += "   ]\n"
+
+            if nodes == set():
+                base_str += "   Nodes = {}\n"
+            else:
+                base_str += "   Nodes = {\n"
+                for node in nodes:
+                    base_str += f"     {str(node)},\n"
+                base_str += "   }\n\n"
+            return base_str
+
+        base_str = "The evidence you have given is as follows:\n"
+        base_str += " Evidence you are certain of:\n"
+        base_str = evidence_str(base_str, self.certain_edges, self.certain_nodes)
+        base_str += " Evidence you are uncertain of:\n"
+        base_str = evidence_str(base_str, self.uncertain_edges, self.uncertain_nodes)
+        return base_str
+
+    @property
+    def certain_edges(self) -> List[Tuple[str]]:
+        """A list of all edges of the ChainEventGraph that have been observed.
+
+        `certain_edges` is a list of edge tuples of the form:
+        `[edge_1, edge_2, ... edge_n]`
+
+        Each edge tuple takes the form:
+        `("source_node_name", "destination_node_name", "edge_label")`
+        """
+        return self._certain_edges
+
+    @property
+    def uncertain_edges(self) -> List[Tuple[str]]:
+        """A list of sets of edges of the ChainEventGraph which might have occured.
+        `uncertain_edges` is a list of sets of edge tuples of the form:
+
+        `[{(a, b, label), (a, c, label)}, {(x, y, label), (x, z, label)}]`
+
+        Each edge tuple takes the form:
+
+        `("source_node_name", "destination_node_name", "edge_label")`
+        """
+        return self._uncertain_edges
+
+    @property
+    def certain_nodes(self) -> Set[str]:
+        """A set of all nodes of the ChainEventGraph that have been observed.
+
+        `certain_nodes` is a set of nodes of the form:
+        `{"node_1", "node_2", "node_3", ... "node_n"}`"""
+        return self._certain_nodes
+
+    @property
+    def uncertain_nodes(self) -> List[Set[str]]:
+        """A list of sets of nodes of the ChainEventGraph where there is uncertainty
+        which of the nodes in each set happened.
+        `uncertain_nodes` is a list of sets of nodes of the form:
+        `[{"node_1", "node_2"}, {"node_3", "node_4"}, ...]`"""
+        return self._uncertain_nodes
+
+    @property
+    def paths(self) -> List[List[Tuple[str]]]:
+        """
+        :return: A list of all paths through the reduced ChainEventGraph."""
+        return self._path_list
+
+    @property
+    def graph(self) -> ChainEventGraph:
+        """
+        :return: The reduced graph once all evidence has been taken into account.
+        :rtype: ChainEventGraph
+        """
+        self._update_path_list()
+        self._update_edges_and_vertices()
+        subgraph_view = nx.subgraph_view(
+            G=self._ceg, filter_node=self._filter_node, filter_edge=self._filter_edge
+        )
+        reduced_graph = ChainEventGraph(subgraph_view, generate=False).copy()
+
+        self._propagate_reduced_graph_probabilities(reduced_graph)
+
+        return reduced_graph
+
+    def clear_all_evidence(self) -> None:
+        """Resets the evidence provided."""
+        self._certain_edges = []
+        self._uncertain_edges = []
+        self._certain_nodes = set()
+        self._uncertain_nodes = []
+        self._edges = set()
+        self._vertices = set()
+
+    def add_certain_edge(self, src: str, dst: str, label: str) -> None:
+        """
+        Specify an edge that has been observed.
+        :param src: Edge source node label
+        :type src: str
+
+        :param dst: Edge destination node label
+        :type dst: str
+
+        :param label: Label of certain edge
+        :type label: str
+        """
+
+        edge = (src, dst, label)
+        if edge not in self._ceg.edges:
+            raise ValueError(
+                f"This edge {edge}, does not exist" f" in the Chain Event Graph."
+            )
+        self._certain_edges.append(edge)
+
+    def remove_certain_edge(self, src: str, dst: str, label: str) -> None:
+        """
+        Specify an edge to remove from the certain edges.
+        :param src: Edge source node label
+        :type src: str
+
+        :param dst: Edge destination node label
+        :type dst: str
+
+        :param label: Label of certain edge
+        :type label: str
+        """
+        try:
+            edge = (src, dst, label)
+            self._certain_edges.remove(edge)
+        except ValueError as err:
+            raise ValueError(
+                f"Edge {(src, dst, label)} not found in the certain " f"edge list."
+            ) from err
+
+    def add_certain_edge_list(self, edges: List[Tuple[str]]) -> None:
+        """
+        Specify a list of edges that have all been observed.
+
+        :param edges: List of edge tuples of the form ("src", "dst", "label")
+        :type edges: List[Tuple[str]]
+        """
+        for edge in edges:
+            self.add_certain_edge(*edge)
+
+    def remove_certain_edge_list(self, edges: List[Tuple[str]]) -> None:
+        """
+        Specify a list of edges that in the certain edge list
+        to remove.
+
+        :param edges: List of edge tuples of the form ("src", "dst", "label")
+        :type edges: List[Tuple[str]]
+        """
+        for edge in edges:
+            self.remove_certain_edge(*edge)
+
+    def add_uncertain_edge_set(self, edge_set: Set[Tuple[str]]) -> None:
+        """
+        Specify a set of edges where one of the edges has
+        occured, but you are uncertain of which one it is.
+
+        :param edge_set: Set of edge tuples of the form ("src", "dst", "label")
+        :type edge_set: Set[Tuple[str]]
+        """
+        for edge in edge_set:
+            if edge not in self._ceg.edges:
+                raise ValueError(
+                    f"The edge {edge}, does not exist" f" in the Chain Event Graph."
+                )
+
+        self._uncertain_edges.append(edge_set)
+
+    def remove_uncertain_edge_set(self, edge_set: Set[Tuple[str]]) -> None:
+        """
+        Specify a set of edges to remove from the uncertain edges.
+
+        :param edge_set: Set of edge tuples of the form ("src", "dst", "label")
+        :type edge_set: Set[Tuple[str]]
+        """
+        try:
+            self._uncertain_edges.remove(edge_set)
+        except ValueError as err:
+            raise ValueError(
+                f"{edge_set} not found in the uncertain edge list."
+            ) from err
+
+    def add_uncertain_edge_set_list(self, edge_sets: List[Set[Tuple[str]]]) -> None:
+        """
+        Specify a list of sets of edges where one of the edges has
+        occured, but you are uncertain of which one it is.
+
+        :param edge_set: List of sets of edge tuples of the form ("src", "dst", "label")
+        :type edge_set: List[Set[Tuple[str]]]
+        """
+        for edge_set in edge_sets:
+            self.add_uncertain_edge_set(edge_set)
+
+    def remove_uncertain_edge_set_list(self, edge_sets: List[Set[Tuple[str]]]) -> None:
+        """
+        Specify a list of sets of edges to remove from the evidence list.
+
+        :param edge_set: List of sets of edge tuples of the form ("src", "dst", "label")
+        :type edge_set: List[Set[Tuple[str]]]
+        """
+        for edge_set in edge_sets:
+            self.remove_uncertain_edge_set(edge_set)
+
+    def add_certain_node(self, node: str) -> None:
+        """Specify a node in the graph that has been observed.
+
+        :param node: A node label e.g. "w4"
+        :type node: str
+        """
+        if node not in self._ceg.nodes:
+            raise ValueError(
+                f"The node {node}, does not exist in the Chain Event Graph."
+            )
+        self._certain_nodes.add(node)
+
+    def remove_certain_node(self, node: str) -> None:
+        """
+        Specify a node to be removed from the certain nodes list.
+
+        :param node: A node label e.g. "w4"
+        :type node: str
+        """
+        try:
+            self._certain_nodes.remove(node)
+        except KeyError as err:
+            raise ValueError(
+                f"Node {node} not found in the set of certain nodes."
+            ) from err
+
+    def add_certain_node_set(self, nodes: Set[str]) -> None:
+        """
+        Specify a set of nodes that have been observed.
+
+        :param nodes: A set of node labels e.g. {"w4", "w8"}
+        :type nodes: Set[str]
+        """
+        for node in nodes:
+            self.add_certain_node(node)
+
+    def remove_certain_node_set(self, nodes: Set[str]) -> None:
+        """
+        Specify a list of nodes to remove from the list of nodes that have
+        been observed.
+
+        :param nodes: A set of node labels e.g. {"w4", "w8"}
+        :type nodes: Set[str]
+        """
+        for node in nodes:
+            self.remove_certain_node(node)
+
+    def add_uncertain_node_set(self, node_set: Set[str]) -> None:
+        """
+        Specify a set of nodes where one of the nodes has
+        occured, but you are uncertain of which one it is.
+
+        :param nodes: A set of node labels e.g. {"w4", "w8"}
+        :type nodes: Set[str]
+        """
+        for node in node_set:
+            if node not in self._ceg.nodes:
+                raise ValueError(
+                    f"The node {node}, does not exist" f" in the Chain Event Graph."
+                )
+
+        self._uncertain_nodes.append(node_set)
+
+    def remove_uncertain_node_set(self, node_set: Set[str]) -> None:
+        """
+        Specify a set of nodes to be removed from the uncertain
+        nodes set list.
+
+        :param nodes: A set of node labels e.g. {"w4", "w8"}
+        :type nodes: Set[str]
+        """
+        try:
+            self._uncertain_nodes.remove(node_set)
+        except ValueError as err:
+            raise ValueError(
+                f"{node_set} not found in the uncertain node list."
+            ) from err
+
+    def add_uncertain_node_set_list(self, node_sets: List[Set[str]]) -> None:
+        """
+        Specify a list of sets of nodes where in each set, one of
+        the nodes has occured, but you are uncertain of which one it is.
+
+        :param nodes: A collection of sets of uncertain nodes.
+        :type nodes: List[Set[str]]"""
+        for node_set in node_sets:
+            self.add_uncertain_node_set(node_set)
+
+    def remove_uncertain_node_set_list(self, node_sets: List[Set[str]]) -> None:
+        """
+        Specify a list of sets nodes to remove from the list of uncertain
+        sets of nodes.
+
+        :param nodes: A collection of sets of uncertain nodes.
+        :type nodes: List[Set[str]]
+        """
+        for node_set in node_sets:
+            self.remove_uncertain_node_set(node_set)
+
+    @staticmethod
+    def _remove_paths(paths: List, to_remove: Iterable):
+        for path in to_remove:
+            paths.remove(path)
+
+    def _apply_certain_edges(self, paths):
+        to_remove = []
+        for edge in self.certain_edges:
+            for path in paths:
+                if (edge not in path) and (path not in to_remove):
+                    to_remove.append(path)
+
+        self._remove_paths(paths, to_remove)
+
+    def _apply_certain_nodes(self, paths):
+        to_remove = []
+        for node in self.certain_nodes:
+            for path in paths:
+                for (src, dst, _) in path:
+                    if node in (src, dst):
+                        break
+                else:
+                    if path not in to_remove:
+                        to_remove.append(path)
+
+        self._remove_paths(paths, to_remove)
+
+    def _apply_uncertain_edges(self, paths):
+        to_remove = []
+
+        def find_edge_and_add_to_remove():
+            edge_found = False
+            for edge in path:
+                if edge in edge_set:
+                    if edge_found:
+                        if path not in to_remove:
+                            to_remove.append(path)
+                        break
+                    edge_found = True
+            else:
+                if not edge_found and path not in to_remove:
+                    to_remove.append(path)
+
+        for edge_set in self.uncertain_edges:
+            for path in paths:
+                find_edge_and_add_to_remove()
+
+        self._remove_paths(paths, to_remove)
+
+    def _apply_uncertain_nodes(self, paths):
+        to_remove = []
+
+        def find_node_and_add_to_remove():
+            node_found = False
+            for (src, _, _) in path:
+                if src in node_set:
+                    if node_found:
+                        if path not in to_remove:
+                            to_remove.append(path)
+                        break
+                    node_found = True
+            else:
+                if not node_found and path not in to_remove:
+                    to_remove.append(path)
+
+        for node_set in self.uncertain_nodes:
+            for path in paths:
+                find_node_and_add_to_remove()
+
+        self._remove_paths(paths, to_remove)
+
+    def _update_path_list(self):
+        path_list = deepcopy(self._ceg.path_list)
+
+        if self.certain_edges:
+            self._apply_certain_edges(path_list)
+
+        if self.certain_nodes:
+            self._apply_certain_nodes(path_list)
+
+        if self.uncertain_edges:
+            self._apply_uncertain_edges(path_list)
+
+        if self.uncertain_nodes:
+            self._apply_uncertain_nodes(path_list)
+
+        self._path_list = path_list
+
+    def _update_edges_and_vertices(self):
+        edges = set()
+        vertices = set()
+
+        for path in self.paths:
+            for (src, dst, label) in path:
+                edges.add((src, dst, label))
+                vertices.add(src)
+                vertices.add(dst)
+
+        self._edges = edges
+        self._vertices = vertices
+
+    def _filter_edge(self, src, dst, label) -> bool:
+        return bool((src, dst, label) in self._edges)
+
+    def _filter_node(self, node) -> bool:
+        return bool(node in self._vertices)
+
+    @staticmethod
+    def _propagate_reduced_graph_probabilities(graph: ChainEventGraph) -> None:
+        # pylint: disable=too-many-locals
+        sink = graph.sink
+        root = graph.root
+        graph.nodes[sink]["emphasis"] = 1
+        node_set = set([sink])
+
+        while node_set != {root}:
+            try:
+                dst = node_set.pop()
+            except KeyError as err:
+                raise KeyError(
+                    "Graph has more than one root... Propagation cannot continue."
+                ) from err
+
+            for src, edges in graph.pred[dst].items():
+                for edge_label, data in edges.items():
+                    edge = (src, dst, edge_label)
+                    emph = graph.nodes[dst]["emphasis"]
+                    prob = data["probability"]
+                    graph.edges[edge]["potential"] = prob * emph
+                successors = graph.succ[src]
+
+                try:
+                    emphasis = 0
+                    for _, succ_edges in successors.items():
+                        for edge_label, data in succ_edges.items():
+                            emphasis += data["potential"]
+
+                    graph.nodes[src]["emphasis"] = emphasis
+                    node_set.add(src)
+                except KeyError:
+                    pass
+
+        for (src, dst, label) in list(graph.edges):
+            edge_potential = graph.edges[(src, dst, label)]["potential"]
+            pred_node_emphasis = graph.nodes[src]["emphasis"]
+            probability = edge_potential / pred_node_emphasis
+            graph.edges[(src, dst, label)]["probability"] = probability
diff --git a/src/cegpy/graphs/ceg.py b/src/cegpy/graphs/ceg.py
deleted file mode 100644
index e3ef112..0000000
--- a/src/cegpy/graphs/ceg.py
+++ /dev/null
@@ -1,728 +0,0 @@
-import pydotplus as pdp
-import networkx as nx
-from copy import deepcopy
-import itertools as it
-
-from ..utilities.util import Util
-from IPython.display import Image
-from IPython import get_ipython
-import logging
-
-logger = logging.getLogger('cegpy.chain_event_graph')
-
-
-class ChainEventGraph(nx.MultiDiGraph):
-    """
-    Class: Chain Event Graph
-
-    Input: Staged tree object (StagedTree)
-    Output: Chain event graphs
-    """
-    def __init__(self, staged_tree=None, **attr):
-        super().__init__(staged_tree, **attr)
-        self.sink_suffix = '&infin;'
-        self.node_prefix = 'w'
-
-        if staged_tree is not None:
-            try:
-                self.ahc_output = deepcopy(staged_tree.ahc_output)
-            except AttributeError:
-                self.ahc_output = {}
-        else:
-            logger.info("Class called with no incoming graph.")
-        self.evidence = Evidence(self)
-
-    @property
-    def node_prefix(self):
-        return self._node_prefix
-
-    @node_prefix.setter
-    def node_prefix(self, value):
-        self._node_prefix = str(value)
-
-    @property
-    def sink_suffix(self):
-        return self._sink_suffix
-
-    @sink_suffix.setter
-    def sink_suffix(self, value):
-        self._sink_suffix = str(value)
-
-    @property
-    def sink_node(self):
-        return "%s%s" % (self.node_prefix, self.sink_suffix)
-
-    @property
-    def root_node(self):
-        return ("%s0" % self.node_prefix)
-
-    @property
-    def certain_evidence(self):
-        return self._certain_evidence
-
-    @certain_evidence.setter
-    def certain_evidence(self, value):
-        self._certain_evidence = value
-
-    @property
-    def uncertain_evidence(self):
-        return self._uncertain_evidence
-
-    @uncertain_evidence.setter
-    def uncertain_evidence(self, value):
-        self._uncertain_evidence = value
-
-    @property
-    def path_list(self):
-        return self.__path_list
-
-    @path_list.setter
-    def path_list(self, value):
-        self.__path_list = value
-
-    @property
-    def reduced(self):
-        return self.evidence.reduced_graph
-
-    @property
-    def stages(self):
-        self.__stages = {}
-        node_stages = dict(self.nodes(data='stage', default=None))
-        for k, v in node_stages.items():
-            try:
-                self.__stages[v].append(k)
-            except KeyError:
-                self.__stages[v] = [k]
-
-        return self.__stages
-
-    def clear_evidence(self):
-        self.evidence = Evidence(self)
-
-    def generate(self):
-        '''
-        This function takes the output of the AHC algorithm and identifies
-        the positions i.e. the vertices of the CEG and the edges of the CEG
-        along with their edge labels and edge counts. Here we use the
-        algorithm in our paper with the optimal stopping time.
-        '''
-        def check_vertices_can_be_merged(v1, v2) -> bool:
-            has_same_successor_nodes = \
-                set(self.adj[v1].keys()) == set(self.adj[v2].keys())
-
-            if has_same_successor_nodes:
-                has_same_outgoing_edges = True
-                v1_adj = self.succ[v1]
-                for succ_node in list(v1_adj.keys()):
-                    v1_edges = self.succ[v1][succ_node]
-                    v2_edges = self.succ[v2][succ_node]
-
-                    if v1_edges is None or v2_edges is None:
-                        has_same_outgoing_edges &= False
-                        break
-
-                    v2_edge_labels = \
-                        [label for label in v2_edges.keys()]
-
-                    for label in v1_edges.keys():
-                        if label not in v2_edge_labels:
-                            has_same_outgoing_edges &= False
-                            break
-                        else:
-                            has_same_outgoing_edges &= True
-            else:
-                has_same_outgoing_edges = False
-
-            try:
-                in_same_stage = \
-                    self.nodes[v1]['stage'] == self.nodes[v2]['stage']
-            except KeyError:
-                in_same_stage = False
-
-            return in_same_stage and \
-                has_same_successor_nodes and has_same_outgoing_edges
-
-        def merge_nodes(nodes_to_merge):
-            """nodes to merge should be a set of 2 element tuples"""
-            temp_1 = 'temp_1'
-            temp_2 = 'temp_2'
-            while nodes_to_merge != set():
-                nodes = nodes_to_merge.pop()
-                new_node = nodes[0]
-                # Copy nodes to temp nodes
-                node_map = {
-                    nodes[0]: temp_1,
-                    nodes[1]: temp_2
-                }
-                nx.relabel_nodes(self, node_map, copy=False)
-                ebunch_to_remove = []  # List of edges to remove
-                self.add_node(new_node)
-                for succ, t1_edge_dict in self.succ[temp_1].items():
-                    edge_labels = list(t1_edge_dict.keys())
-                    while edge_labels != []:
-                        new_edge_data = {}
-                        label = edge_labels.pop(0)
-                        t1_edge = t1_edge_dict[label]
-                        t2_edge = self.succ[temp_2][succ][label]
-                        new_edge_data['count'] = \
-                            t1_edge['count'] + t2_edge['count']
-                        new_edge_data['prior'] = \
-                            t1_edge['prior'] + t2_edge['prior']
-                        new_edge_data['posterior'] = \
-                            t1_edge['posterior'] + t2_edge['posterior']
-                        try:
-
-                            new_edge_data['probability'] = \
-                                t1_edge['probability']
-                            self.add_edge(
-                                u_for_edge=new_node,
-                                v_for_edge=succ,
-                                key=label,
-                                count=new_edge_data['count'],
-                                prior=new_edge_data['prior'],
-                                posterior=new_edge_data['posterior'],
-                                probability=new_edge_data['probability']
-                            )
-                        except KeyError:
-                            self.add_edge(
-                                u_for_edge=new_node,
-                                v_for_edge=succ,
-                                key=label,
-                                count=new_edge_data['count'],
-                                prior=new_edge_data['prior'],
-                                posterior=new_edge_data['posterior']
-                            )
-                        ebunch_to_remove.append((temp_1, succ, label))
-                        ebunch_to_remove.append((temp_2, succ, label))
-
-                self.remove_edges_from(ebunch_to_remove)
-                nx.relabel_nodes(
-                    G=self,
-                    mapping={temp_1: new_node, temp_2: new_node},
-                    copy=False
-                )
-
-                # Some nodes have been removed, we need to update the
-                # mergeable list to point to new nodes if required
-                temp_list = list(nodes_to_merge)
-                for pair in temp_list:
-                    if nodes[1] in pair:
-                        new_pair = (
-                            # the other node of the pair
-                            pair[pair.index(nodes[1]) - 1],
-                            # the new node it will be merged to
-                            new_node
-                        )
-                        nodes_to_merge.remove(pair)
-                        if new_pair[0] != new_pair[1]:
-                            nodes_to_merge.add(new_pair)
-                pass
-
-        def relabel_nodes(base_nodes, renamed_nodes=[]):
-            next_level = []
-            # first, relabel the successors of this node
-            for node in base_nodes:
-                node_mapping = {}
-                for succ in self.succ[node].keys():
-                    if succ != self.sink_node and succ not in renamed_nodes:
-                        node_mapping[succ] = self.__get_next_node_name()
-                        next_level.append(node_mapping[succ])
-                        renamed_nodes.append(node_mapping[succ])
-
-                if node_mapping != {}:
-                    nx.relabel_nodes(
-                        self,
-                        node_mapping,
-                        copy=False
-                    )
-            if next_level != []:
-                relabel_nodes(next_level, renamed_nodes)
-
-        if self.ahc_output == {}:
-            raise ValueError("Run staged tree AHC transitions first.")
-        # rename root node:
-        nx.relabel_nodes(self, {'s0': self.root_node}, copy=False)
-        self.__update_probabilities()
-        self.__trim_leaves_from_graph()
-        self.__update_distances_of_nodes_to_sink_node()
-        src_node_gen = self.__gen_nodes_with_increasing_distance(
-            start=1
-        )
-        next_set_of_nodes = next(src_node_gen)
-
-        while next_set_of_nodes != [self.root_node]:
-            nodes_to_merge = set()
-            while len(next_set_of_nodes) > 1:
-                node_1 = next_set_of_nodes.pop(0)
-                for node_2 in next_set_of_nodes:
-                    mergeable = check_vertices_can_be_merged(node_1, node_2)
-                    if mergeable:
-                        nodes_to_merge.add((node_1, node_2))
-
-            merge_nodes(nodes_to_merge)
-
-            try:
-                next_set_of_nodes = next(src_node_gen)
-            except StopIteration:
-                next_set_of_nodes = []
-
-        relabel_nodes([self.root_node])
-        self.__update_path_list()
-
-    @property
-    def dot_graph(self):
-        return self._generate_dot_graph()
-
-    def _generate_dot_graph(self):
-        graph = pdp.Dot(graph_type='digraph', rankdir='LR')
-        edge_probabilities = list(
-            self.edges(data='probability', default=1, keys=True)
-        )
-
-        for (u, v, k, p) in edge_probabilities:
-            full_label = "{}\n{:.2f}".format(k, p)
-            graph.add_edge(
-                pdp.Edge(
-                    src=u,
-                    dst=v,
-                    label=full_label,
-                    labelfontcolor='#009933',
-                    fontsize='10.0',
-                    color='black'
-                )
-            )
-        nodes = list(nx.topological_sort(self))
-        for node in nodes:
-            try:
-                fill_colour = self.nodes[node]['colour']
-            except KeyError:
-                fill_colour = 'white'
-            label = "<" + node[0] + "<SUB>" + node[1:] + "</SUB>" + ">"
-            graph.add_node(
-                pdp.Node(
-                    name=node,
-                    label=label,
-                    style='filled',
-                    fillcolor=fill_colour
-                )
-            )
-        return graph
-
-    def create_figure(self, filename):
-        """
-        Draws the chain event graph representation of the stage tree,
-        and saves it to "<filename>.filetype". Supports any filetype that
-        graphviz supports. e.g: "event_tree.png" or "event_tree.svg" etc.
-        """
-        filename, filetype = Util.generate_filename_and_mkdir(filename)
-        graph = self.dot_graph
-        graph.write(str(filename), format=filetype)
-
-        if get_ipython() is None:
-            return None
-        else:
-            return Image(graph.create_png())
-
-    def __update_probabilities(self):
-        count_total_lbl = 'count_total'
-        edge_counts = list(self.edges(data='count', keys=True, default=0))
-
-        for stage, stage_nodes in self.stages.items():
-            count_total = 0
-            stage_edges = {}
-            if stage is not None:
-
-                for (u, _, k, c) in edge_counts:
-                    if u in stage_nodes:
-                        count_total += c
-                        try:
-                            stage_edges[k] += c
-                        except KeyError:
-                            stage_edges[k] = c
-
-                for node in stage_nodes:
-                    self.nodes[node][count_total_lbl] = count_total
-
-                for (u, v, k, _) in edge_counts:
-                    if u in stage_nodes:
-                        self.edges[u, v, k]['probability'] =\
-                            stage_edges[k] / count_total
-            else:
-                for node in stage_nodes:
-                    count_total = 0
-                    stage_edges = {}
-                    for (u, _, k, c) in edge_counts:
-                        if u == node:
-                            count_total += c
-                            try:
-                                stage_edges[k] += c
-                            except KeyError:
-                                stage_edges[k] = c
-
-                    self.nodes[node][count_total_lbl] = count_total
-                    for (u, v, k, _) in edge_counts:
-                        if u == node:
-                            self.edges[u, v, k]['probability'] =\
-                                stage_edges[k] / count_total
-
-    def __update_path_list(self) -> None:
-        path_generator = nx.all_simple_edge_paths(
-            self,
-            self.root_node,
-            self.sink_node
-        )
-        path_list = []
-        while True:
-            try:
-                path_list.append(next(path_generator))
-            except StopIteration:
-                self.path_list = path_list
-                break
-
-    def __update_distances_of_nodes_to_sink_node(self) -> None:
-        """
-        Iterates through the graph until it finds the root node.
-        For each node, it determines the maximum number of edges
-        from that node to the sink node.
-        """
-        max_dist = 'max_dist_to_sink'
-        self.nodes[self.sink_node][max_dist] = 0
-        node_queue = [self.sink_node]
-
-        while node_queue != [self.root_node]:
-            node = node_queue.pop(0)
-            for pred in self.predecessors(node):
-                max_dist_to_sink = set()
-                for succ in self.successors(pred):
-                    try:
-                        max_dist_to_sink.add(
-                            self.nodes[succ][max_dist]
-                        )
-                        self.nodes[pred][max_dist] = max(max_dist_to_sink) + 1
-                    except KeyError:
-                        break
-
-                if pred not in node_queue:
-                    node_queue.append(pred)
-
-    def __gen_nodes_with_increasing_distance(self, start=0) -> list:
-        max_dists = nx.get_node_attributes(self, 'max_dist_to_sink')
-        distance_dict = {}
-        for key, value in max_dists.items():
-            distance_dict.setdefault(value, []).append(key)
-
-        for dist in range(len(distance_dict)):
-            if dist >= start:
-                yield distance_dict[dist]
-
-    def __get_next_node_name(self):
-        try:
-            num = str(next(self._num_iter))
-        except AttributeError:
-            self._num_iter = it.count(1, 1)
-            num = str(next(self._num_iter))
-
-        return str(self.node_prefix) + num
-
-    def __trim_leaves_from_graph(self):
-        # Create new CEG sink node
-        self.add_node(self.sink_node, colour='lightgrey')
-        outgoing_edges = deepcopy(self.succ).items()
-        # Check to see if any nodes have no outgoing edges.
-        for node, outgoing_edges in outgoing_edges:
-            if outgoing_edges == {} and node != self.sink_node:
-                incoming_edges = deepcopy(self.pred[node]).items()
-                # When node is identified as a leaf check the
-                # predessesor nodes that have edges that enter this node.
-                for pred_node, edges in incoming_edges:
-                    for edge_label, edge in edges.items():
-                        # Create new edge that points to the sink node,
-                        # with all the same data as the edge we will delete.
-                        try:
-                            prob = edge['probability']
-                        except KeyError:
-                            prob = 1
-                        self.add_edge(
-                            pred_node,
-                            self.sink_node,
-                            key=edge_label,
-                            count=edge['count'],
-                            prior=edge['prior'],
-                            posterior=edge['posterior'],
-                            probability=prob
-                        )
-                self.remove_node(node)
-
-
-class Evidence:
-    CERTAIN = True
-    UNCERTAIN = False
-
-    def __init__(self, graph):
-        self.__graph = graph
-
-        self.certain_edges = []
-        self.uncertain_edges = []
-        self.certain_vertices = set()
-        self.uncertain_vertices = set()
-
-    @property
-    def reduced_graph(self):
-        return self.__create_reduced_graph()
-
-    @property
-    def path_list(self):
-        return self._path_list
-
-    @path_list.setter
-    def path_list(self, value):
-        self._path_list = value
-
-    @property
-    def edges(self):
-        return list(self._edges)
-
-    @edges.setter
-    def edges(self, value):
-        self._edges = value
-
-    @property
-    def vertices(self):
-        return list(self._vertices)
-
-    @vertices.setter
-    def vertices(self, value):
-        self._vertices = value
-
-    def add_edge(self, u, v, label, certain):
-        edge = (u, v, label)
-        if certain:
-            self.certain_edges.append(edge)
-        else:
-            self.uncertain_edges.append(edge)
-
-    def add_edges_from(self, edges, certain):
-        for (u, v, k) in edges:
-            self.add_edge(u, v, k, certain)
-
-    def remove_edge(self, u, v, label, certain):
-        if certain:
-            self.certain_edges.remove((u, v, label))
-        else:
-            self.uncertain_edges.remove((u, v, label))
-
-    def remove_edges_from(self, edges, certain):
-        for (u, v, k) in edges:
-            self.remove_edge(u, v, k, certain)
-
-    def add_node(self, node, certain):
-        if certain:
-            self.certain_vertices.add(node)
-        else:
-            self.uncertain_vertices.add(node)
-
-    def add_nodes_from(self, nodes, certain):
-        for node in nodes:
-            self.add_node(node, certain)
-
-    def remove_node(self, node, certain):
-        if certain:
-            self.certain_vertices.remove(node)
-        else:
-            self.uncertain_vertices.remove(node)
-
-    def remove_nodes_from(self, nodes, certain):
-        for node in nodes:
-            self.remove_node(node, certain)
-
-    def __repr__(self) -> str:
-        repr = "Evidence(CertainEdges={}, CertainVertices={}," +\
-            " UncertainEdges={}, UncertainVertices={})"
-        return repr.format(
-            str(self.certain_edges),
-            str(self.certain_vertices),
-            str(self.uncertain_edges),
-            str(self.uncertain_vertices)
-        )
-
-    def __str__(self) -> str:
-        """Returns human readable version of the evidence you've provided."""
-        def evidence_str(base_str, edges, vertices):
-            if edges == []:
-                base_str += '   Edges = []\n'
-            else:
-                base_str += '   Edges = [\n'
-                for edge in edges:
-                    base_str += '     %s,\n' % (str(edge))
-                base_str += '   ]\n'
-
-            if vertices == set():
-                base_str += '   Vertices = {}\n'
-            else:
-                base_str += '   Vertices = {\n'
-                for vertex in vertices:
-                    base_str += "     '%s',\n" % (str(vertex))
-                base_str += '   }\n\n'
-            return base_str
-
-        base_str = 'The evidence you have given is as follows:\n'
-        base_str += ' Evidence you are certain of:\n'
-        base_str = evidence_str(
-            base_str,
-            self.certain_edges,
-            self.certain_vertices
-        )
-        base_str += ' Evidence you are uncertain of:\n'
-        base_str = evidence_str(
-            base_str,
-            self.uncertain_edges,
-            self.uncertain_vertices
-        )
-        return base_str
-
-    def __update_path_list(self):
-        def remove_paths(paths_list, paths_to_remove):
-            for path in paths_to_remove:
-                paths_list.remove(path)
-
-        # Certain Evidence
-        certain_path_list = self.__graph.path_list.copy()
-        # Uncertain Evidence
-        uncertain_path_list = self.__graph.path_list.copy()
-
-        if self.certain_edges or self.certain_vertices:
-            # Apply certain Edges
-            paths_to_remove = []
-            for edge in self.certain_edges:
-                for path in certain_path_list:
-                    if edge not in path:
-                        paths_to_remove.append(path)
-
-            remove_paths(certain_path_list, paths_to_remove)
-
-            # Apply certain vertices
-            paths_to_remove = []
-            for vertex in self.certain_vertices:
-                for path in certain_path_list:
-                    vertex_in_path = False
-                    for (u, v, _) in path:
-                        if vertex == u or vertex == v:
-                            vertex_in_path = True
-                            break
-
-                    if not vertex_in_path:
-                        paths_to_remove.append(path)
-
-            remove_paths(certain_path_list, paths_to_remove)
-
-        if self.uncertain_edges or self.uncertain_vertices:
-            # Apply uncertain Edges
-            paths_to_include = []
-            for edge in self.uncertain_edges:
-                for path in uncertain_path_list:
-                    if edge in path:
-                        paths_to_include.append(path)
-
-            # Apply uncertain Vertices
-            for vertex in self.uncertain_vertices:
-                for path in uncertain_path_list:
-                    vertex_in_path = False
-                    for (u, v, _) in path:
-                        if vertex == u or vertex == v:
-                            vertex_in_path = True
-                            break
-
-                    if vertex_in_path:
-                        paths_to_include.append(path)
-
-            uncertain_path_list = paths_to_include
-
-        # Take paths that are found in both certain and uncertain lists
-        if not self.uncertain_edges and not self.uncertain_vertices:
-            self.path_list = certain_path_list
-        elif not self.certain_edges and not self.certain_vertices:
-            self.path_list = uncertain_path_list
-        else:
-            new_path_list = []
-            for path in certain_path_list:
-                if path in uncertain_path_list:
-                    new_path_list.append(path)
-            self.path_list = new_path_list
-
-    def __update_edges_and_vertices(self):
-        edges = set()
-        vertices = set()
-
-        for path in self.path_list:
-            for (u, v, k) in path:
-                edges.add((u, v, k))
-                vertices.add(u)
-                vertices.add(v)
-
-        self.edges = edges
-        self.vertices = vertices
-
-    def __filter_edge(self, u, v, k) -> bool:
-        if (u, v, k) in self.edges:
-            return True
-        else:
-            return False
-
-    def __filter_node(self, n) -> bool:
-        if n in self.vertices:
-            return True
-        else:
-            return False
-
-    def __propagate_reduced_graph_probabilities(self, graph) -> None:
-        sink = graph.sink_node
-        root = graph.root_node
-        graph.nodes[sink]['emphasis'] = 1
-        node_set = set([sink])
-
-        while node_set != {root}:
-            try:
-                v_node = node_set.pop()
-            except KeyError:
-                break
-
-            for u_node, edges in graph.pred[v_node].items():
-                for edge_label, data in edges.items():
-                    edge = (u_node, v_node, edge_label)
-                    emph = graph.nodes[v_node]['emphasis']
-                    prob = data['probability']
-                    graph.edges[edge]['potential'] = \
-                        prob * emph
-                successors = graph.succ[u_node]
-
-                try:
-                    emphasis = 0
-                    for _, edges in successors.items():
-                        for edge_label, data in edges.items():
-                            emphasis += data['potential']
-
-                    graph.nodes[u_node]['emphasis'] = emphasis
-                    node_set.add(u_node)
-                except KeyError:
-                    pass
-
-        for (u, v, k) in list(graph.edges):
-            edge_potential = graph.edges[(u, v, k)]['potential']
-            pred_node_emphasis = graph.nodes[u]['emphasis']
-            probability = edge_potential / pred_node_emphasis
-            graph.edges[(u, v, k)]['probability'] = probability
-
-    def __create_reduced_graph(self) -> ChainEventGraph:
-        self.__update_path_list()
-        self.__update_edges_and_vertices()
-        subgraph = ChainEventGraph(
-            nx.subgraph_view(
-                G=self.__graph,
-                filter_node=self.__filter_node,
-                filter_edge=self.__filter_edge
-            )
-        ).copy()
-
-        self.__propagate_reduced_graph_probabilities(subgraph)
-
-        return subgraph
diff --git a/src/cegpy/graphs/ctceg.py b/src/cegpy/graphs/ctceg.py
deleted file mode 100644
index ad74ace..0000000
--- a/src/cegpy/graphs/ctceg.py
+++ /dev/null
@@ -1,6 +0,0 @@
-from ..graphs.ceg import ChainEventGraph
-
-
-class CTChainEventGraph(ChainEventGraph):
-    def __init__(self, params) -> None:
-        pass
diff --git a/src/cegpy/trees/__init__.py b/src/cegpy/trees/__init__.py
index 7df9aba..e69de29 100644
--- a/src/cegpy/trees/__init__.py
+++ b/src/cegpy/trees/__init__.py
@@ -1,2 +0,0 @@
-from .event import EventTree
-from .staged import StagedTree
\ No newline at end of file
diff --git a/src/cegpy/trees/_event.py b/src/cegpy/trees/_event.py
new file mode 100644
index 0000000..7a5f8c8
--- /dev/null
+++ b/src/cegpy/trees/_event.py
@@ -0,0 +1,485 @@
+"""cegpy event tree."""
+
+from collections import defaultdict
+from typing import Dict, List, Mapping, Optional, Tuple, Union
+import logging
+import textwrap
+import numpy as np
+import pydotplus as pdp
+from IPython.display import Image
+from IPython import get_ipython
+import pandas as pd
+import networkx as nx
+from cegpy.utilities._util import (
+    generate_filename_and_mkdir,
+    create_sampling_zeros,
+)
+
+# create logger object for this module
+logger = logging.getLogger("cegpy.event_tree")
+
+
+class EventTree(nx.MultiDiGraph):
+    """
+    This class extends the NetworkX MultiDiGraph class to allow the creation
+    of event tree representations of data.
+
+    :param dataframe: Required - DataFrame containing variables as column headers,
+        with event name strings in each cell. These event names will be used to
+        create the edges of the event tree. Counts of each event will
+        be extracted and attached to each edge.
+    :type dataframe: pandas.DataFrame
+
+    :param sampling_zero_paths: Optional - Paths to sampling
+        zeros.
+
+        Format is as follows: [('edge_1',), ('edge_1', 'edge_2'), ...]
+
+        If no paths are specified, default setting is that no sampling zero paths
+        are created.
+    :type sampling_zero_paths: List[Tuple[str]] or None
+
+    :param var_order: Optional - Specifies the ordering of variables to be adopted
+        in the event tree.
+        Default var_order is obtained from the order of columns in dataframe.
+        String labels in the list should match the column names in dataframe.
+    :type var_order: List[str] or None
+
+    :param struct_missing_label: Optional - Label in the dataframe for observations
+        which are structurally missing; e.g: Post operative health status is
+        irrelevant for a dead patient.
+        Label example: "struct".
+    :type struct_missing_label: str or None
+
+    :param missing_label: Optional - Label in the dataframe for observations which are
+        missing values that are not structurally missing.
+        e.g: Missing height for some individuals in the sample.
+        Label example: "miss"
+        Whatever label is provided will be renamed in the event tree to "missing".
+    :type missing_label: str or None
+
+    :param complete_case: Optional - If True, all entries (rows) with non-structural
+        missing values are removed. Default setting: False.
+    :type complete_case: bool
+    """
+
+    _sampling_zero_paths: Optional[List[Tuple]] = None
+    _sorted_paths: Mapping[Tuple[str], int]
+    _edge_attributes: List = ["count"]
+
+    def __init__(
+        self,
+        dataframe: pd.DataFrame,
+        sampling_zero_paths=None,
+        var_order=None,
+        struct_missing_label=None,
+        missing_label=None,
+        complete_case=False,
+    ) -> None:
+        # Checking argument inputs are sensible
+        if not isinstance(dataframe, pd.DataFrame):
+            raise ValueError("The dataframe parameter must be a pandas.DataFrame")
+        # Initialise Networkx DiGraph class
+        incoming_graph_data = None
+        attr = dict()
+        super().__init__(incoming_graph_data, **attr)
+
+        if struct_missing_label is not None and not isinstance(
+            struct_missing_label, str
+        ):
+            raise ValueError("struct_missing_label should be a string")
+
+        if missing_label is not None and not isinstance(missing_label, str):
+            raise ValueError("missing_label should be a string")
+
+        if not isinstance(complete_case, bool):
+            raise ValueError("complete_case should be a boolean")
+
+        self.sampling_zeros = sampling_zero_paths
+
+        # Dealing with structural and non-structural...
+        # ... missing value labels
+        if struct_missing_label is not None:
+            dataframe.replace(
+                struct_missing_label,
+                "",
+                inplace=True,
+            )
+
+        if missing_label is not None:
+            if complete_case is True:
+                dataframe.replace(
+                    missing_label,
+                    "missing",
+                    inplace=True,
+                )
+                rows_to_delete = np.where(dataframe == "missing")[0].tolist()
+                dataframe.drop(
+                    rows_to_delete,
+                    axis=0,
+                    inplace=True,
+                )
+                dataframe.reset_index(drop=True, inplace=True)
+            else:
+                dataframe.replace(
+                    missing_label,
+                    "missing",
+                    inplace=True,
+                )
+
+        # Paths sorted alphabetically in order of length
+        self._sorted_paths = defaultdict(int)
+
+        # pandas dataframe passed via parameters
+        self.dataframe = (
+            dataframe[var_order].astype(str)
+            if var_order is not None
+            else dataframe.astype(str)
+        )
+
+        self._construct_event_tree()
+        logger.info("Initialisation complete!")
+
+    @property
+    def root(self) -> str:
+        """
+        :return: The name of the root node of the event tree, currently hard coded to 's0'.
+        :rtype: str
+        """
+        return "s0"
+
+    @property
+    def variables(self) -> List:
+        """
+        :return: The column headers of the dataset.
+        :rtype: List[str]
+        """
+        variables = list(self.dataframe.columns)
+        logger.info("Variables extracted from dataframe were:")
+        logger.info(variables)
+        return variables
+
+    @property
+    def sampling_zeros(self) -> Union[List[Tuple[str]], None]:
+        """
+        Setting this property will apply sampling zero paths to the tree.
+        If different to previous value, the event tree will be regenerated.
+
+        :return: Sampling zero paths provided by the user.
+
+        :rtype: List[Tuple[str]] or None
+        """
+        if self._sampling_zero_paths is None:
+            logger.info(
+                "EventTree.sampling_zero_paths \
+                    has not been set."
+            )
+        return self._sampling_zero_paths
+
+    @sampling_zeros.setter
+    def sampling_zeros(self, sz_paths):
+        if sz_paths is None:
+            self._sampling_zero_paths = None
+        else:
+            # checkes if the user has inputted sz paths correctly
+            sz_paths = self._check_sampling_zero_paths_param(sz_paths)
+
+            if sz_paths:
+                self._sampling_zero_paths = sz_paths
+            else:
+                raise ValueError(
+                    "Parameter 'sampling_zero_paths' not in expected format. "
+                    "Should be a list of tuples like so:\n"
+                    "[('edge_1',), ('edge_1', 'edge_2'), ...]"
+                )
+
+    @property
+    def situations(self) -> List[str]:
+        """
+        :return: The situations of the tree (non-leaf nodes).
+        :rtype: List[str]
+        """
+        return [node for node, out_degree in self.out_degree if out_degree != 0]
+
+    @property
+    def leaves(self) -> List[str]:
+        """
+        :return: The leaves of the tree.
+        :rtype: List[str]
+        """
+        return [node for node, out_degree in self.out_degree if out_degree == 0]
+
+    @property
+    def edge_counts(self) -> Dict:
+        """
+        The counts along edges all edges in the tree, where edges are a
+        Tuple like so: ("source_node", "destination_node", "edge_label").
+
+        :return: A mapping of edges to their counts.
+
+        :rtype: Dict[Tuple[str], Int]
+        """
+        return nx.get_edge_attributes(self, "count")
+
+    @property
+    def categories_per_variable(self) -> Dict:
+        """
+        The number of unique categories/levels for each variable
+        (column headings in dataframe).
+
+        :return: A mapping of variables to the number of unique categories/levels.
+        :rtype: Dict[str, Int]
+        """
+
+        def display_nan_warning():
+            logger.warning(
+                textwrap.dedent(
+                    """   --- NaNs found in the dataframe!! ---
+                    cegpy assumes that NaNs are either structural zeros or
+                    structural missing values.
+                    Any non-structural missing values must be dealt with
+                    prior to providing the dataset to any of the cegpy
+                    functions. Any non-structural zeros should be explicitly
+                    added into the cegpy objects.
+                    --- See documentation for more information. ---"""
+                )
+            )
+
+        categories_to_ignore = {"N/A", "NA", "n/a", "na", "NAN", "nan"}
+        catagories_per_variable = {}
+        nans_filtered = False
+
+        for var in self.variables:
+            categories = set(self.dataframe[var].unique().tolist())
+            # remove nan with pandas
+            pd_filtered_categories = {x for x in categories if pd.notna(x)}
+
+            # remove any string nans that might have made it in.
+            filtered_cats = pd_filtered_categories - categories_to_ignore
+            if pd_filtered_categories != filtered_cats:
+                nans_filtered = True
+
+            catagories_per_variable[var] = len(filtered_cats)
+
+        if nans_filtered:
+            display_nan_warning()
+
+        return catagories_per_variable
+
+    def dot_event_graph(self, edge_info: str = "count") -> pdp.Dot:
+        """Returns Dot graph representation of the event tree.
+        :param edge_info: Optional - Chooses which summary measure to be displayed on edges.
+        In event trees, only "count" can be displayed, so this can be omitted.
+
+        :type edge_info: str
+        :return: A graphviz Dot representation of the graph.
+        :rtype: pydotplus.Dot
+        """
+        return self._generate_dot_graph(fill_colour="lightgrey", edge_info=edge_info)
+
+    def _generate_dot_graph(self, fill_colour=None, edge_info="count"):
+        node_list = list(self)
+        graph = pdp.Dot(graph_type="digraph", rankdir="LR")
+        if edge_info in self._edge_attributes:
+            edge_info_dict = nx.get_edge_attributes(self, edge_info)
+        else:
+            logger.warning(
+                "edge_info '%s' does not exist for the "
+                "%s class. Using the default of 'count' values "
+                "on edges instead. For more information, see the "
+                "documentation.",
+                edge_info,
+                self.__class__.__name__,
+            )
+            edge_info_dict = nx.get_edge_attributes(self, "count")
+
+        for edge, attribute in edge_info_dict.items():
+            if edge_info == "count":
+                edge_details = str(edge[2]) + "\n" + str(attribute)
+            else:
+                edge_details = f"{edge[2]}\n{float(attribute):.2f}"
+
+            graph.add_edge(
+                pdp.Edge(
+                    edge[0],
+                    edge[1],
+                    label=edge_details,
+                    labelfontcolor="#009933",
+                    fontsize="10.0",
+                    color="black",
+                )
+            )
+        for node in node_list:
+            if fill_colour is None:
+                try:
+                    fill_node_colour = self.nodes[node]["colour"]
+                except KeyError:
+                    fill_node_colour = "lightgrey"
+            else:
+                fill_node_colour = fill_colour
+            label = "<" + node[0] + "<SUB>" + node[1:] + "</SUB>" + ">"
+            graph.add_node(
+                pdp.Node(
+                    name=node, label=label, style="filled", fillcolor=fill_node_colour
+                )
+            )
+        return graph
+
+    def _create_figure(self, graph: pdp.Dot, filename: str):
+        """Draws the event tree for the process described by the dataset,
+        and saves it to "<filename>.filetype". Supports any filetype that
+        graphviz supports. e.g: "event_tree.png" or "event_tree.svg" etc.
+        """
+        if filename is None:
+            logger.warning("No filename. Figure not saved.")
+        else:
+            filename, filetype = generate_filename_and_mkdir(filename)
+            logger.info("--- generating graph ---")
+            logger.info("--- writing %s file ---", filetype)
+            graph.write(str(filename), format=filetype)
+            graph_image = None
+
+        if get_ipython() is not None:
+            logger.info("--- Exporting graph to notebook ---")
+            graph_image = Image(graph.create_png())
+        else:
+            graph_image = None
+
+        return graph_image
+
+    def create_figure(
+        self, filename=None, edge_info: str = "count"
+    ) -> Union[Image, None]:
+        """Creates event tree from the dataframe.
+
+        :param filename: Optional - When provided, file is saved to the filename,
+            local to the current working directory.
+            e.g. if filename = "output/event_tree.svg", the file will be saved to:
+            cwd/output/event_tree.svg
+            Otherwise, if function is called inside an interactive notebook, image
+            will be displayed in the notebook, even if filename is omitted.
+            Supports any filetype that graphviz supports. e.g: "event_tree.png" or
+            "event_tree.svg" etc.
+
+        :type filename: str
+
+        :param edge_info: Optional - Chooses which summary measure to be displayed on edges.
+            In event trees, only "count" can be displayed, so this can be omitted.
+        :type edge_info: str
+
+        :return: The event tree Image object.
+        :rtype: IPython.display.Image or None
+        """
+        return self._create_figure(self.dot_event_graph(edge_info=edge_info), filename)
+
+    def _create_unsorted_paths_dict(self) -> defaultdict:
+        """Creates and populates a dictionary of all paths provided in the dataframe,
+        in the order in which they are given."""
+        unsorted_paths = defaultdict(int)
+
+        for variable_number in range(0, len(self.variables)):
+            dataframe_upto_variable = self.dataframe.loc[
+                :, self.variables[0 : variable_number + 1]
+            ]
+
+            for row in dataframe_upto_variable.itertuples():
+                row = row[1:]
+                new_row = [
+                    edge_label
+                    for edge_label in row
+                    if edge_label != np.nan
+                    and str(edge_label) != "NaN"
+                    and str(edge_label) != "nan"
+                    and edge_label != ""
+                ]
+                new_row = tuple(new_row)
+
+                # checking if the last edge label in row was nan. That would
+                # result in double counting nan must be identified as string
+                if (
+                    row[-1] != np.nan
+                    and str(row[-1]) != "NaN"
+                    and str(row[-1]) != "nan"
+                    and row[-1] != ""
+                ):
+                    unsorted_paths[new_row] += 1
+
+        return unsorted_paths
+
+    def _create_path_dict_entries(self):
+        """Create path dict entries for each path, including the
+        sampling zero paths if any.
+        Each path is an ordered sequence of edge labels starting
+        from the root.
+        The keys in the dict are ordered alphabetically.
+        Also calls the method self._sampling_zeros to ensure
+        manually added path format is correct.
+        Added functionality to remove NaN/null edge labels
+        assuming they are structural zeroes"""
+        unsorted_paths = self._create_unsorted_paths_dict()
+
+        if self.sampling_zeros is not None:
+            unsorted_paths = create_sampling_zeros(self.sampling_zeros, unsorted_paths)
+
+        depth = len(max(list(unsorted_paths.keys()), key=len))
+        keys_of_list = list(unsorted_paths.keys())
+        sorted_keys = []
+        for deep in range(0, depth + 1):
+            unsorted_mini_list = [key for key in keys_of_list if len(key) == deep]
+            sorted_keys = sorted_keys + sorted(unsorted_mini_list)
+
+        for key in sorted_keys:
+            self._sorted_paths[key] = unsorted_paths[key]
+
+        node_list = self._create_node_list_from_paths(self._sorted_paths)
+        self.add_nodes_from(node_list)
+
+    def _check_sampling_zero_paths_param(self, sampling_zero_paths) -> List:
+        """Check param 'sampling_zero_paths' is in the correct format"""
+        coerced_sampling_zero_paths = []
+        for tup in sampling_zero_paths:
+            if not isinstance(tup, tuple):
+                return None
+            coerced_sampling_zero_paths.append(tuple([str(elem) for elem in tup]))
+
+        return coerced_sampling_zero_paths
+
+    def _create_node_list_from_paths(self, paths) -> list:
+        """Creates list of all nodes: includes root, situations, leaves"""
+        node_list = [self.root]
+
+        for node_number, _ in enumerate(list(paths.keys()), start=1):
+            node_list.append(f"s{node_number}")
+
+        return node_list
+
+    def _construct_event_tree(self):
+        """Constructs event_tree DiGraph.
+        Takes the paths, and adds all the nodes and edges to the Graph"""
+
+        logger.info("Starting construction of event tree")
+        self._create_path_dict_entries()
+        # Taking a list of a networkx graph object (self) provides a list
+        # of all the nodes
+        node_list = list(self)
+
+        # Work through the sorted paths list to build the event tree.
+        edge_labels_list = ["root"]
+        for path, count in list(self._sorted_paths.items()):
+            path = list(path)
+            edge_labels_list.append(path)
+            if path[:-1] in edge_labels_list:
+                path_edge_comes_from = edge_labels_list.index(path[:-1])
+                self.add_edge(
+                    u_for_edge=node_list[path_edge_comes_from],
+                    v_for_edge=node_list[edge_labels_list.index(path)],
+                    key=path[-1],
+                    count=count,
+                )
+            else:
+                self.add_edge(
+                    u_for_edge=node_list[0],
+                    v_for_edge=node_list[edge_labels_list.index(path)],
+                    key=path[-1],
+                    count=count,
+                )
diff --git a/src/cegpy/trees/_staged.py b/src/cegpy/trees/_staged.py
new file mode 100644
index 0000000..b11b81a
--- /dev/null
+++ b/src/cegpy/trees/_staged.py
@@ -0,0 +1,803 @@
+"""cegpy Staged Tree"""
+from copy import deepcopy
+from fractions import Fraction
+from itertools import combinations, chain
+import logging
+from operator import add, sub, itemgetter
+from typing import Dict, List, Optional, Tuple, Union
+import networkx as nx
+import scipy.special
+from IPython.display import Image
+
+from cegpy.trees._event import EventTree
+from cegpy.utilities._util import generate_colours
+
+logger = logging.getLogger("cegpy.staged_tree")
+
+
+# pylint: disable=too-many-instance-attributes
+class StagedTree(EventTree):
+    """Representation of a Staged Tree.
+
+    A staged tree is a tree where each node is a situation, and each edge is a
+    transition from one situation to another. Each situation is given a 'stage'
+    which groups it with other situations which have the same outgoing edges,
+    with equivalent probabilities of occurring.
+
+    The class is an extension of EventTree.
+
+    :param dataframe: Required - DataFrame containing variables as column headers,
+        with event name strings in each cell. These event names will be used to
+        create the edges of the event tree. Counts of each event will
+        be extracted and attached to each edge.
+    :type dataframe: pandas.DataFrame
+
+    :param sampling_zero_paths: Optional - Paths to sampling
+        zeros.
+
+        Format is as follows: [('edge_1',), ('edge_1', 'edge_2'), ...]
+
+        If no paths are specified, default setting is that no sampling zero paths
+        are created.
+    :type sampling_zero_paths: List[Tuple[str]] or None
+
+    :param var_order: Optional - Specifies the ordering of variables to be adopted
+        in the event tree.
+        Default var_order is obtained from the order of columns in dataframe.
+        String labels in the list should match the column names in dataframe.
+    :type var_order: List[str] or None
+
+    :param struct_missing_label: Optional - Label in the dataframe for observations
+        which are structurally missing; e.g: Post operative health status is
+        irrelevant for a dead patient.
+        Label example: "struct".
+    :type struct_missing_label: str or None
+
+    :param missing_label: Optional - Label in the dataframe for observations which are
+        missing values that are not structurally missing.
+        e.g: Missing height for some individuals in the sample.
+        Label example: "miss"
+        Whatever label is provided will be renamed in the event tree to "missing".
+    :type missing_label: str or None
+
+    :param complete_case: Optional - If True, all entries (rows) with non-structural
+        missing values are removed. Default setting: False.
+    :type complete_case: bool
+    """
+
+    _edge_attributes: List = ["count", "prior", "posterior", "probability"]
+
+    # pylint: disable=too-many-arguments
+    def __init__(
+        self,
+        dataframe,
+        sampling_zero_paths=None,
+        var_order=None,
+        struct_missing_label=None,
+        missing_label=None,
+        complete_case=False,
+    ) -> None:
+        # Call event tree init to generate event tree
+        super().__init__(
+            dataframe=dataframe,
+            sampling_zero_paths=sampling_zero_paths,
+            var_order=var_order,
+            struct_missing_label=struct_missing_label,
+            missing_label=missing_label,
+            complete_case=complete_case,
+        )
+
+        self._mean_posterior_probs = []
+        self._merged_situations = []
+        self._stage_colours = []
+        self._sort_count = 0
+        self._colours_for_situations = []
+        self._ahc_output = {}
+        logger.debug("Starting Staged Tree")
+
+    @property
+    def prior(self) -> Dict[Tuple[str], List[Fraction]]:
+        """A mapping of priors keyed by edge. Keys are
+        3-tuples of the form: (src, dst, edge_label).
+
+        :return: A mapping edge -> priors.
+        :rtype: Dict[Tuple[str], List[Fraction]]
+        """
+        return nx.get_edge_attributes(self, "prior")
+
+    @prior.setter
+    def prior(self, prior: List[List[Fraction]]) -> None:
+        offset = 0
+        for node_idx, node_priors in enumerate(prior):
+            node_name = f"s{node_idx + offset}"
+            while not self.succ[node_name]:
+                offset += 1
+                node_name = f"s{node_idx + offset}"
+            for edge_prior_idx, succ_key in enumerate(self.succ[node_name].keys()):
+                label = list(self.succ[node_name][succ_key])[0]
+                self.edges[(node_name, succ_key, label)]["prior"] = Fraction(
+                    node_priors[edge_prior_idx]
+                )
+
+    @property
+    def prior_list(self) -> List[List[Fraction]]:
+        """
+        :return: Priors in the form of a list of lists.
+        :rtype: List[List[Fraction]]
+        """
+
+        prior_list = []
+        prev_node = list(self.prior)[0][0]
+        succ_list = []
+
+        for edge, prior in self.prior.items():
+            node = edge[0]
+            if node != prev_node:
+                prior_list.append(succ_list)
+                succ_list = []
+            succ_list.append(prior)
+            prev_node = node
+
+        if succ_list:
+            prior_list.append(succ_list)
+        return prior_list
+
+    @property
+    def posterior(self) -> Dict[Tuple[str], List[Fraction]]:
+        """Posteriors along each edge, calculated by adding edge count to the prior for
+        each edge.
+        Keys are 3-tuples of the form: (src, dst, edge_label).
+
+        :return: Mapping of edge -> edge_count + prior
+        :rtype: Dict[Tuple[str], List[Fraction]]
+        """
+        try:
+            posterior = nx.get_edge_attributes(self, "posterior")
+            if posterior == {}:
+                raise AttributeError("Posterior not yet set.")
+            return posterior
+        except AttributeError:
+            for edge in self.edges:
+                edge_dict = self.edges[edge]
+                posterior = edge_dict["prior"] + edge_dict["count"]
+                edge_dict["posterior"] = posterior
+            return nx.get_edge_attributes(self, "posterior")
+
+    @property
+    def posterior_list(self) -> List[List[Fraction]]:
+        """
+        :return: Posteriors in the form of a list of lists.
+        :rtype: List[List[Fraction]]
+        """
+        posterior_list = []
+        prev_node = list(self.posterior)[0][0]
+        succ_list = []
+
+        for edge, posterior in self.posterior.items():
+            node = edge[0]
+            if node != prev_node:
+                posterior_list.append(succ_list)
+                succ_list = []
+            succ_list.append(posterior)
+            prev_node = node
+
+        if succ_list:
+            posterior_list.append(succ_list)
+        return posterior_list
+
+    @property
+    def alpha(self) -> float:
+        """
+        The equivalent sample size set for the root node which is then uniformly
+        propagated through the tree.
+
+        :return: The value of Alpha.
+        :rtype: float
+        """
+        return self._alpha
+
+    @alpha.setter
+    def alpha(self, value):
+        self._alpha = value
+
+    @property
+    def hyperstage(self) -> List[List[str]]:
+        """
+        Indication of which nodes are allowed to be in the same stage. Each
+        list is a list of node names e.g. "s0".
+
+        :return: The List of all hyperstages.
+        :rtype: List[List[str]]
+        """
+        return self._hyperstage
+
+    @hyperstage.setter
+    def hyperstage(self, value):
+        self._hyperstage = value
+
+    @property
+    def edge_countset(self) -> List[List]:
+        """
+        Indexed the same as situations.
+        :return: Edge counts emination from each node of the tree.
+        :rtype: List[List]
+        """
+        return self._create_edge_countset()
+
+    @property
+    def ahc_output(self) -> Dict:
+        """
+        Contains a List of Lists containing all the situations that were merged,
+        and the log likelihood.
+
+        :return: The output from the AHC algorithm.
+        :rtype: Dict
+        """
+        return self._ahc_output
+
+    def _check_hyperstages(self, hyperstage) -> None:
+        hyper_situations = chain(*hyperstage)
+        hyper_situations_set = set(hyper_situations)
+        situations_set = set(self.situations)
+        # Check if all situations are present
+        missing = situations_set.difference(hyper_situations_set)
+        if missing:
+            raise ValueError(
+                f"Situation(s) {missing} are missing from the list of " "hyperstages."
+            )
+        # Check if all situations provided exist
+        extra = hyper_situations_set.difference(situations_set)
+        if extra:
+            raise ValueError(f"Situation(s) {extra} are not present in the tree.")
+        # Check if all situations in a stage have the same number of edges
+        for stage in hyperstage:
+            n_edges = self.out_degree[stage[0]]
+            for node in stage:
+                if self.out_degree[node] != n_edges:
+                    raise ValueError(
+                        "Situations in the same hyperstage "
+                        "must have the same number of outgoing edges."
+                    )
+
+    def _check_prior(self, node_priors_list) -> None:
+        if len(node_priors_list) != len(self.edge_countset):
+            raise ValueError(
+                "Number of sub-lists in the list of priors "
+                "must agree with the number of situations."
+            )
+
+        for node_idx, node_priors in enumerate(node_priors_list):
+            if len(node_priors) != len(self.edge_countset[node_idx]):
+                raise ValueError(
+                    "The length of each sub-list in the list of priors "
+                    "must agree with the number of edges emanating from "
+                    "its corresponding situation."
+                )
+            for node_prior in node_priors:
+                if node_prior < 0:
+                    raise ValueError("All priors must be non-negative.")
+
+    def _store_params(self, prior, alpha, hyperstage) -> None:
+        """User has passed in AHC params, this function processes them,
+        and generates any default AHC params if required."""
+        if prior:
+            if alpha:
+                self.alpha = None
+                logging.warning(
+                    "Params Warning!! When prior is given, alpha is not required!"
+                )
+            self._check_prior(prior)
+            self.prior = prior
+        else:
+            if alpha is None:
+                self.alpha = self._calculate_default_alpha()
+                logging.warning(
+                    "Params Warning!! Neither prior nor alpha "
+                    "were provided. Using default alpha "
+                    "value of %d.",
+                    self.alpha,
+                )
+            else:
+                self.alpha = alpha
+
+            # No matter what alpha is, generate default prior
+            self.prior = self._create_default_prior(self.alpha)
+
+        if hyperstage is None:
+            self.hyperstage = self._create_default_hyperstage()
+        else:
+            self._check_hyperstages(hyperstage)
+            self.hyperstage = hyperstage
+
+    def _calculate_default_alpha(self) -> int:
+        """If no alpha is given, a default value is calculated.
+        The value is calculated by determining the maximum number
+        of categories that any one variable has"""
+        logger.info("Calculating default prior")
+        max_count = max(list(self.categories_per_variable.values()))
+        return max_count
+
+    def _create_default_prior(self, alpha) -> list:
+        """default prior set for the AHC method using the mass conservation property.
+        That is, the alpha param is the phantom sample starting at
+        the root, and it is spread equally across all edges along the tree.
+        (see chapter 5 of Collazo, Gorgen & Smith 'Chain Event Graphs', 2018)
+        The prior is a list of lists. Each list gives the prior along the
+        edges of a specific situation.
+        Indexed same as self.situations & self.egde_countset"""
+
+        logger.info("Generating default prior")
+        default_prior = [0] * len(self.situations)
+        sample_size_at_node = {}
+
+        # Root node is assigned phantom sample (alpha)
+        if isinstance(alpha, float):
+            alpha = Fraction.from_float(alpha)
+        elif isinstance(alpha, (int, str)):
+            alpha = Fraction.from_float(float(alpha))
+        else:
+            raise TypeError(
+                "Prior generator param alpha is in a strange format.\
+                            .."
+            )
+
+        sample_size_at_node[self.root] = alpha
+
+        for node_idx, node in enumerate(self.situations):
+            # How many nodes emanate from the current node?
+            number_of_emanating_nodes = self.out_degree[node]
+
+            # Divide the sample size from the current node equally among
+            # emanating nodes
+            equal_distribution_of_sample = (
+                sample_size_at_node[node] / number_of_emanating_nodes
+            )
+            default_prior[node_idx] = [
+                equal_distribution_of_sample
+            ] * number_of_emanating_nodes
+
+            relevant_terminating_nodes = [
+                edge[1] for edge in list(self.edges) if edge[0] == node
+            ]
+
+            for terminating_node in relevant_terminating_nodes:
+                sample_size_at_node[terminating_node] = equal_distribution_of_sample
+
+        return default_prior
+
+    def _create_default_hyperstage(self) -> list:
+        """Generates default hyperstage for the AHC method.
+        A hyperstage is a list of lists such that two situaions can be in the
+        same stage only if there are elements of the same list for some list
+        in the hyperstage.
+        The default is to allow all situations with the same number of
+        outgoing edges and the same edge labels to be in a common list."""
+        logger.info("Creating default hyperstage")
+        hyperstage = []
+        info_of_edges = []
+
+        for node in self.situations:
+            labels = [edge[2] for edge in self.edges if edge[0] == node]
+            labels.sort()
+
+            info_of_edges.append([self.out_degree[node], labels])
+
+        sorted_info = []
+        for info in info_of_edges:
+            if info not in sorted_info:
+                sorted_info.append(info)
+
+        for info in sorted_info:
+            situations_with_value_edges = []
+            for idx, situation in enumerate(self.situations):
+                if info_of_edges[idx] == info:
+                    situations_with_value_edges.append(situation)
+            hyperstage = hyperstage + [situations_with_value_edges]
+
+        return hyperstage
+
+    def _create_edge_countset(self) -> list:
+        """Each element of list contains a list with counts along edges emanating from
+        a specific situation. Indexed same as self.situations"""
+        logger.info("Creating edge countset")
+        edge_countset = []
+
+        for node in self.situations:
+            edge_countset.append(
+                [count for edge, count in self.edge_counts.items() if edge[0] == node]
+            )
+        return edge_countset
+
+    @staticmethod
+    def _calculate_lg_of_sum(array) -> float:
+        """function to calculate log gamma of the sum of an array"""
+        array = [float(x) for x in array]
+        return scipy.special.gammaln(sum(array))
+
+    @staticmethod
+    def _calculate_sum_of_lg(array) -> float:
+        """function to calculate log gamma of each element of an array"""
+        return sum([scipy.special.gammaln(float(x)) for x in array])
+
+    def _calculate_initial_loglikelihood(self, prior, posterior) -> float:
+        """calculating log likelihood given a prior and posterior"""
+        # Calculate prior contribution
+        logger.info("Calculating initial loglikelihood")
+
+        pri_lg_of_sum = [self._calculate_lg_of_sum(elem) for elem in prior]
+        pri_sum_of_lg = [self._calculate_sum_of_lg(elem) for elem in prior]
+        pri_contribution = list(map(sub, pri_lg_of_sum, pri_sum_of_lg))
+
+        # Calculate posterior contribution
+        post_lg_of_sum = [self._calculate_lg_of_sum(elem) for elem in posterior]
+        post_sum_of_lg = [self._calculate_sum_of_lg(elem) for elem in posterior]
+        post_contribution = list(map(sub, post_sum_of_lg, post_lg_of_sum))
+
+        sum_pri_contribution = sum(pri_contribution)
+        sum_post_contribution = sum(post_contribution)
+        return sum_pri_contribution + sum_post_contribution
+
+    def _calculate_bayes_factor(self, prior1, posterior1, prior2, posterior2) -> float:
+        """calculates the bayes factor comparing two models which differ in
+        only one stage"""
+        new_prior = list(map(add, prior1, prior2))
+        new_posterior = list(map(add, posterior1, posterior2))
+        return (
+            self._calculate_lg_of_sum(new_prior)
+            - self._calculate_lg_of_sum(prior1)
+            - self._calculate_lg_of_sum(prior2)
+            - self._calculate_lg_of_sum(new_posterior)
+            + self._calculate_lg_of_sum(posterior1)
+            + self._calculate_lg_of_sum(posterior2)
+            + self._calculate_sum_of_lg(new_posterior)
+            - self._calculate_sum_of_lg(posterior1)
+            - self._calculate_sum_of_lg(posterior2)
+            - self._calculate_sum_of_lg(new_prior)
+            + self._calculate_sum_of_lg(prior1)
+            + self._calculate_sum_of_lg(prior2)
+        )
+
+    def _sort_list(self, list_of_tuples) -> list:
+        """function to sort a list of lists to remove repetitions"""
+
+        for idx_1, tup_1 in enumerate(list_of_tuples):
+            for idx_2, tup_2 in enumerate(list_of_tuples):
+                tups_intersect = set(tup_1) & set(tup_2)
+                if tups_intersect:
+                    union = tuple(set(tup_1) | set(tup_2))
+                    list_of_tuples[idx_1] = []
+                    list_of_tuples[idx_2] = union
+
+        new_list_of_tuples = [elem for elem in list_of_tuples if elem != []]
+
+        if new_list_of_tuples == list_of_tuples:
+            return new_list_of_tuples
+
+        return self._sort_list(new_list_of_tuples)
+
+    def _calculate_mean_posterior_probs(
+        self,
+        merged_situations: List,
+        posteriors: List,
+    ) -> List:
+        """Iterates through array of lists, calculates mean
+        posterior probabilities"""
+        mean_posterior_probs = []
+
+        for sit in self.situations:
+            if sit not in list(chain(*merged_situations)):
+                merged_situations.append((sit,))
+        for stage in merged_situations:
+            for sit in stage:
+                sit_idx = self.situations.index(sit)
+                if all(posteriors[sit_idx]) != 0:
+                    stage_probs = posteriors[sit_idx]
+                    break
+                stage_probs = []
+
+            total = sum(stage_probs)
+            mean_posterior_probs.append(
+                [round(elem / total, 3) for elem in stage_probs]
+            )
+
+        return mean_posterior_probs
+
+    @staticmethod
+    def _independent_hyperstage_generator(hyperstages: List[List]) -> List[List[List]]:
+        """Spit out the next hyperstage that can be dealt with
+        independently."""
+        new_hyperstages = [[hyperstages[0]]]
+
+        for sublist in hyperstages[1:]:
+            hs_to_add = [sublist]
+
+            for hyperstage in new_hyperstages.copy():
+                for other_sublist in hyperstage:
+                    if not set(other_sublist).isdisjoint(set(sublist)):
+                        hs_to_add.extend(hyperstage)
+                        new_hyperstages.remove(hyperstage)
+
+            new_hyperstages.append(hs_to_add)
+
+        return new_hyperstages
+
+    def _execute_ahc(self, hyperstage=None) -> Tuple[List, float, List]:
+        """finds all subsets and scores them"""
+        if hyperstage is None:
+            hyperstage = deepcopy(self.hyperstage)
+
+        priors = deepcopy(self.prior_list)
+        posteriors = deepcopy(self.posterior_list)
+
+        loglikelihood = self._calculate_initial_loglikelihood(priors, posteriors)
+
+        merged_situation_list = []
+
+        # Which list in hyperstage have only 1 edge coming out of them
+        # For that list, add the list to the merged situation list, and
+        # add together their priors/posteriors, and remove it from the
+        # hyperstage
+        for sub_hyper in deepcopy(hyperstage):
+            if self.out_degree[sub_hyper[0]] == 1:
+                merged_situation_list.append(tuple(sub_hyper))
+                indexes = [self.situations.index(situ) for situ in sub_hyper]
+                if len(indexes) == 1:
+                    sub_hyper_priors = itemgetter(*indexes)(priors)
+                    sub_hyper_posteriors = itemgetter(*indexes)(posteriors)
+                else:
+                    sub_hyper_priors = list(chain(*itemgetter(*indexes)(priors)))
+                    sub_hyper_posteriors = list(
+                        chain(*itemgetter(*indexes)(posteriors))
+                    )
+                for loop_idx, p_index in enumerate(indexes):
+                    if loop_idx == 0:
+                        priors[p_index] = [sum(sub_hyper_priors)]
+                        posteriors[p_index] = [sum(sub_hyper_posteriors)]
+                    else:
+                        priors[p_index] = [0]
+                        posteriors[p_index] = [0]
+
+                hyperstage.remove(sub_hyper)
+
+        hyperstage_combinations = [
+            item for sub_hyper in hyperstage for item in combinations(sub_hyper, 2)
+        ]
+
+        while True:
+            newscores_list = [
+                self._calculate_bayes_factor(
+                    priors[self.situations.index(sub_hyper[0])],
+                    posteriors[self.situations.index(sub_hyper[0])],
+                    priors[self.situations.index(sub_hyper[1])],
+                    posteriors[self.situations.index(sub_hyper[1])],
+                )
+                for sub_hyper in hyperstage_combinations
+            ]
+
+            local_score = max(newscores_list)
+
+            if local_score > 0:
+                local_merged = hyperstage_combinations[
+                    newscores_list.index(local_score)
+                ]
+                merge_situ_1, merge_situ_2 = local_merged
+                merge_situ_1_idx = self.situations.index(merge_situ_1)
+                merge_situ_2_idx = self.situations.index(merge_situ_2)
+                merged_situation_list.append(local_merged)
+
+                priors[merge_situ_1_idx] = list(
+                    map(add, priors[merge_situ_1_idx], priors[merge_situ_2_idx])
+                )
+                posteriors[merge_situ_1_idx] = list(
+                    map(add, posteriors[merge_situ_1_idx], posteriors[merge_situ_2_idx])
+                )
+                priors[merge_situ_2_idx] = [0] * len(priors[merge_situ_1_idx])
+                posteriors[merge_situ_2_idx] = [0] * len(posteriors[merge_situ_1_idx])
+
+                loglikelihood += local_score
+            else:
+                break
+
+        merged_situation_list = self._sort_list(merged_situation_list)
+        for sit in self.situations:
+            if sit not in list(chain(*merged_situation_list)):
+                merged_situation_list.append((sit,))
+
+        self._apply_mean_posterior_probs(
+            merged_situations=merged_situation_list,
+            mean_posterior_probs=_calculate_mean_posterior_probs(
+                self.situations, merged_situation_list, posteriors
+            ),
+        )
+
+        return loglikelihood, merged_situation_list
+
+    def _mark_nodes_with_stage_number(self, merged_situations):
+        """AHC algorithm creates a list of indexes to the situations list.
+        This function takes those indexes and creates a new list which is
+        in a string representation of nodes."""
+        self._sort_count = 0
+        stage_count = 0
+        for stage in merged_situations:
+            if len(stage) > 1:
+                for node in stage:
+                    self.nodes[node]["stage"] = stage_count
+                stage_count += 1
+
+    def _generate_colours_for_situations(self, merged_situations, colour_list):
+        """Colours each stage of the tree with an individual colour"""
+        num_colours = len([m for m in merged_situations if len(m) > 1])
+        if colour_list is None:
+            stage_colours = generate_colours(num_colours)
+        else:
+            stage_colours = colour_list
+            if len(colour_list) < num_colours:
+                raise IndexError(
+                    f"The number of colours provided ({len(colour_list)}) is "
+                    "less than the number of distinct colours required "
+                    f"({num_colours})."
+                )
+        self._stage_colours = stage_colours
+        for node in self.nodes:
+            try:
+                stage = self.nodes[node]["stage"]
+                self.nodes[node]["colour"] = stage_colours[stage]
+            except KeyError:
+                self.nodes[node]["colour"] = "lightgrey"
+
+    def calculate_AHC_transitions(
+        self, prior=None, alpha=None, hyperstage=None, colour_list=None
+    ) -> Dict:
+        """Bayesian Agglommerative Hierarchical Clustering algorithm
+        implementation. It returns a list of lists of the situations which
+        have been merged together, the likelihood of the final model.
+
+        :param prior: Optional - A mapping of priors keyed by edge. Keys are
+            3-tuples of the form: (src, dst, edge_label).
+        :type prior: Dict[Tuple[str], List[Fraction]]
+
+        :param alpha: Optional - The equivalent sample size set for the root node
+            which is then uniformly propagated through the tree.
+        :type alpha: float
+
+        :param hyperstage: Optional - Indication of which nodes are allowed to be
+            in the same stage. Each list is a list of node names e.g. "s0".
+        :type hyperstage: List[List[str]]
+
+        :param colour_list: Optional - a list of hex colours to be used for stages.
+            Otherwise, colours evenly spaced around the colour spectrum are used.
+        :type colour_list: List[str]
+
+        :return: The output from the AHC algorithm, specified above.
+        :rtype: Dict
+        """
+        logger.info("\n\n --- Starting AHC Algorithm ---")
+
+        self._store_params(prior, alpha, hyperstage)
+
+        loglikelihood, merged_situations = self._execute_ahc()
+
+        self._mark_nodes_with_stage_number(merged_situations)
+
+        self._generate_colours_for_situations(merged_situations, colour_list)
+
+        self._ahc_output = {
+            "Merged Situations": merged_situations,
+            "Log Likelihood": loglikelihood,
+        }
+        return self.ahc_output
+
+    def dot_staged_graph(self, edge_info: str = "count"):
+        """Returns Dot graph representation of the staged tree.
+        :param edge_info: Optional - Chooses which summary measure to be displayed
+        on edges. Defaults to "count".
+        Options: ["count", "prior", "posterior", "probability"]
+
+        :type edge_info: str
+        :return: A graphviz Dot representation of the graph.
+        :rtype: pydotplus.Dot"""
+        return self._generate_dot_graph(edge_info=edge_info)
+
+    def create_event_tree_figure(
+        self,
+        filename: Optional[str] = None,
+        edge_info: str = "count",
+    ) -> Union[Image, None]:
+        """Creates event tree from the dataframe.
+
+        :param filename: Optional - When provided, file is saved to the filename,
+            local to the current working directory.
+            e.g. if filename = "output/event_tree.svg", the file will be saved to:
+            cwd/output/event_tree.svg
+            Otherwise, if function is called inside an interactive notebook, image
+            will be displayed in the notebook, even if filename is omitted.
+            Supports any filetype that graphviz supports. e.g: "event_tree.png" or
+            "event_tree.svg" etc.
+
+        :type filename: str
+
+        :param edge_info: Optional - Chooses which summary measure to be displayed on
+            edges. Value can take: "count", "prior", "posterior", "probability"
+        :type edge_info: str
+
+        :return: The event tree Image object.
+        :rtype: IPython.display.Image or None
+        """
+        return super().create_figure(filename, edge_info=edge_info)
+
+    def create_figure(
+        self,
+        filename: Optional[str] = None,
+        edge_info: str = "count",
+    ) -> Union[Image, None]:
+        """Draws the coloured staged tree for the process described by
+        the dataset.
+
+        :param filename: Optional - When provided, file is saved to the filename,
+            local to the current working directory.
+            e.g. if filename = "output/event_tree.svg", the file will be saved to:
+            cwd/output/staged_tree.svg
+            Otherwise, if function is called inside an interactive notebook, image
+            will be displayed in the notebook, even if filename is omitted.
+
+            Supports any filetype that graphviz supports. e.g: "staged_tree.png" or
+            "staged_tree.svg" etc.
+
+        :type filename: str
+
+        :param edge_info: Optional - Chooses which summary measure to be displayed on edges.
+            In event trees, only "count" can be displayed, so this can be omitted.
+        :type edge_info: str
+
+        :return: The event tree Image object.
+        :rtype: IPython.display.Image or None
+        """
+        try:
+            if not self._ahc_output:
+                raise AttributeError
+
+            graph = self.dot_staged_graph(edge_info)
+            graph_image = super()._create_figure(graph, filename)
+
+        except AttributeError:
+            logger.error(
+                "----- PLEASE RUN AHC ALGORITHM before trying to"
+                " export a staged tree graph -----"
+            )
+            graph_image = None
+
+        return graph_image
+
+    def _apply_mean_posterior_probs(
+        self, merged_situations: List, mean_posterior_probs: List
+    ) -> None:
+        """Apply the mean posterior probabilities to each edge."""
+        for stage_idx, stage in enumerate(merged_situations):
+            for sit in stage:
+                dst_nodes = list(chain(self.succ[sit]))
+                edge_labels = list(chain(*self.succ[sit].values()))
+                for edge_idx, label in enumerate(edge_labels):
+                    self.edges[(sit, dst_nodes[edge_idx], label)][
+                        "probability"
+                    ] = mean_posterior_probs[stage_idx][edge_idx]
+
+
+def _calculate_mean_posterior_probs(
+    all_situations: List, merged_situations: List, posteriors: List
+) -> List:
+    """Given a staged tree, calculate the mean posterior probs."""
+    # Add all situations that are not in a stage.
+    mean_posterior_probs = []
+    for stage in merged_situations:
+        for sit in stage:
+            sit_idx = all_situations.index(sit)
+            if all(posteriors[sit_idx]) != 0:
+                stage_posteriors = posteriors[sit_idx]
+                break
+
+        total = sum(stage_posteriors)
+        mean_posterior_probs.append(
+            [round(posterior / total, 3) for posterior in stage_posteriors]
+        )
+
+    return mean_posterior_probs
diff --git a/src/cegpy/trees/event.py b/src/cegpy/trees/event.py
deleted file mode 100644
index b799977..0000000
--- a/src/cegpy/trees/event.py
+++ /dev/null
@@ -1,416 +0,0 @@
-from collections import defaultdict
-import numpy as np
-import pydotplus as pdp
-import logging
-from ..utilities.util import Util
-from IPython.display import Image
-from IPython import get_ipython
-import pandas as pd
-import textwrap
-import networkx as nx
-# create logger object for this module
-logger = logging.getLogger('cegpy.event_tree')
-
-
-class EventTree(nx.MultiDiGraph):
-    """
-    Class for event trees.
-
-    This class extends the networkx DiGraph class to allow the creation
-    of event trees from data provided in a pandas dataframe.
-
-    A DiGraph stores nodes and edges with optional data, or attributes.
-
-    DiGraphs hold directed edges.  Self loops are allowed but multiple
-    (parallel) edges are not.
-
-    Nodes can be arbitrary (hashable) Python objects with optional
-    key/value attributes. By convention `None` is not used as a node.
-
-    Edges are represented as links between nodes with optional
-    key/value attributes.
-
-    Parameters
-    ----------
-    dataframe : Pandas dataframe (required)
-        Dataframe containing variables as column headers, with event
-        name strings in each cell. These event names will be used to
-        create the edges of the event tree. Counts of each event will
-        be extracted and attached to each edge.
-
-    sampling_zero_paths: list of tuples containing paths to sampling
-        zeros.
-        Format is as follows: \
-            [('edge_1',), ('edge_1', 'edge_2'), ...]
-
-    incoming_graph_data : input graph (optional, default: None)
-        Data to initialize graph.  If None (default) an empty
-        graph is created.  The data can be an edge list, or any
-        NetworkX graph object.  If the corresponding optional Python
-        packages are installed the data can also be a NumPy matrix
-        or 2d ndarray, a SciPy sparse matrix, or a PyGraphviz graph.
-
-    attr : keyword arguments, optional (default= no attributes)
-        Attributes to add to graph as key=value pairs.
-
-    See Also
-    --------
-    StagedTrees \n
-    ChainEventGraph
-
-    Examples
-    --------
-    >>> G = nx.Graph()  # or DiGraph, MultiGraph, MultiDiGraph, etc
-    >>> G = nx.Graph(name="my graph")
-    >>> e = [(1, 2), (2, 3), (3, 4)]  # list of edges
-    >>> G = nx.Graph(e)
-
-    Arbitrary graph attribute pairs (key=value) may be assigned
-
-    >>> G = nx.Graph(e, day="Friday")
-    >>> G.graph
-    {'day': 'Friday'}
-
-    """
-    def __init__(self, dataframe, sampling_zero_paths=None,
-                 incoming_graph_data=None, var_order=None, **attr) -> None:
-        """Initialize an event tree graph with edges, name, or graph attributes.
-        This class extends the networkx DiGraph class to allow the creation
-        of event trees from data provided in a pandas dataframe.
-
-        Parameters
-        ----------
-        dataframe : Pandas dataframe (required)
-            Dataframe containing variables as column headers, with event
-            name strings in each cell. These event names will be used to
-            create the edges of the event tree. Counts of each event will
-            be extracted and attached to each edge.
-
-        sampling_zero_paths: list of tuples containing paths to sampling
-            zeros.
-            Format is as follows: \
-                [('edge_1',), ('edge_1', 'edge_2'), ...]
-
-        incoming_graph_data : input graph (optional, default: None)
-            Data to initialize graph.  If None (default) an empty
-            graph is created.  The data can be an edge list, or any
-            NetworkX graph object.  If the corresponding optional Python
-            packages are installed the data can also be a NumPy matrix
-            or 2d ndarray, a SciPy sparse matrix, or a PyGraphviz graph.
-        
-        var_order : ordered list of variable names. (optional, default order 
-            of variables in the event tree adopted from the order of columns in 
-            the dataframe). 
-
-        attr : keyword arguments, optional (default= no attributes)
-            Attributes to add to graph as key=value pairs.
-
-        See Also
-        --------
-        convert
-
-        Examples
-        --------
-        >>> G = nx.Graph()  # or DiGraph, MultiGraph, MultiDiGraph, etc
-        >>> G = nx.Graph(name="my graph")
-        >>> e = [(1, 2), (2, 3), (3, 4)]  # list of edges
-        >>> G = nx.Graph(e)
-
-        Arbitrary graph attribute pairs (key=value) may be assigned
-
-        >>> G = nx.Graph(e, day="Friday")
-        >>> G.graph
-        {'day': 'Friday'}
-
-        """
-        logger.info('Initialising')
-        # Initialise Networkx DiGraph class
-        super().__init__(incoming_graph_data, **attr)
-        self._sampling_zero_paths = None
-        self.sampling_zeros = sampling_zero_paths
-
-        # Paths sorted alphabetically in order of length
-        self._sorted_paths = defaultdict(int)
-
-        # pandas dataframe passed via parameters
-        if var_order is not None:
-            self.dataframe = dataframe[var_order]
-        else:
-            self.dataframe = dataframe
-        self.__construct_event_tree()
-        logger.info('Initialisation complete!')
-
-    @property
-    def root(self) -> str:
-        """Root node of the event tree.
-        Currently hard coded to 's0'"""
-        return 's0'
-
-    @property
-    def variables(self) -> list:
-        """The column headers of the dataset"""
-        vars = list(self._dataframe.columns)
-        logger.info('Variables extracted from dataframe were:')
-        logger.info(vars)
-        return vars
-
-    @property
-    def dataframe(self):
-        return self._dataframe
-
-    @dataframe.setter
-    def dataframe(self, value: pd.DataFrame):
-        if isinstance(value, pd.DataFrame):
-            self._dataframe = value
-        else:
-            raise ValueError(
-                "Package currently only supports Pandas DataFrame"
-                " objects provided as the dataframe")
-
-    @property
-    def sampling_zeros(self):
-        if self._sampling_zero_paths is None:
-            logger.info("EventTree.sampling_zero_paths \
-                    has not been set.")
-        return self._sampling_zero_paths
-
-    @sampling_zeros.setter
-    def sampling_zeros(self, sz_paths):
-        """Use this function to set the sampling zero paths.
-        If different to previous value, will re-generate the event tree."""
-        if sz_paths is None:
-            self._sampling_zero_paths = None
-        else:
-            # checkes if the user has inputted sz paths correctly
-            sz_paths = self.__check_sampling_zero_paths_param(sz_paths)
-
-            if sz_paths:
-                self._sampling_zero_paths = sz_paths
-            else:
-                error_str = "Parameter 'sampling_zero_paths' not in expected format. \
-                             Should be a list of tuples like so:\n \
-                             [('edge_1',), ('edge_1', 'edge_2'), ...]"
-                raise ValueError(error_str)
-
-    @property
-    def situations(self) -> list:
-        """List of situations of the tree.
-        (non-leaf nodes)"""
-        return [
-            node for node, out_degree in self.out_degree
-            if out_degree != 0
-        ]
-
-    @property
-    def leaves(self) -> list:
-        """List of leaves of the tree."""
-        # if not already generated, create self.leaves
-        return [
-            node for node, out_degree in self.out_degree
-            if out_degree == 0
-        ]
-
-    @property
-    def edge_counts(self) -> dict:
-        '''list of counts along edges. Indexed same as edges and edge_labels'''
-        return nx.get_edge_attributes(self, 'count')
-
-    @property
-    def categories_per_variable(self) -> dict:
-        '''list of number of unique categories/levels for each variable
-        (a column in the df)'''
-        def display_nan_warning():
-            logger.warning(
-                textwrap.dedent(
-                    """   --- NaNs found in the dataframe!! ---
-                    cegpy assumes that NaNs are either structural zeros or
-                    structural missing values.
-                    Any non-structural missing values must be dealt with
-                    prior to providing the dataset to any of the cegpy
-                    functions. Any non-structural zeros should be explicitly
-                    added into the cegpy objects.
-                    --- See documentation for more information. ---"""
-                )
-            )
-
-        categories_to_ignore = {"N/A", "NA", "n/a", "na", "NAN", "nan"}
-        self._catagories_per_variable = {}
-        nans_filtered = False
-
-        for var in self.variables:
-            categories = set(self._dataframe[var].unique().tolist())
-            # remove nan with pandas
-            pd_filtered_categories = {x for x in categories if pd.notna(x)}
-            if pd_filtered_categories != categories:
-                nans_filtered = True
-
-            # remove any string nans that might have made it in.
-            filtered_cats = pd_filtered_categories - categories_to_ignore
-            if pd_filtered_categories != filtered_cats:
-                nans_filtered = True
-
-            self._catagories_per_variable[var] = len(filtered_cats)
-
-        if nans_filtered:
-            display_nan_warning()
-
-        return self._catagories_per_variable
-
-    @property
-    def dot_graph(self):
-        return self._generate_dot_graph()
-
-    def _generate_dot_graph(self):
-        node_list = list(self)
-        graph = pdp.Dot(graph_type='digraph', rankdir='LR')
-        for edge, count in self.edge_counts.items():
-            edge_details = str(edge[2]) + '\n' + str(count)
-
-            graph.add_edge(
-                pdp.Edge(
-                    edge[0],
-                    edge[1],
-                    label=edge_details,
-                    labelfontcolor="#009933",
-                    fontsize="10.0",
-                    color="black"
-                )
-            )
-
-        for node in node_list:
-            try:
-                fill_colour = self.nodes[node]['colour']
-            except KeyError:
-                fill_colour = 'lightgrey'
-            label = "<" + node[0] + "<SUB>" + node[1:] + "</SUB>" + ">"
-            graph.add_node(
-                pdp.Node(
-                    name=node,
-                    label=label,
-                    style="filled",
-                    fillcolor=fill_colour))
-        return graph
-
-    def create_figure(self, filename):
-        """Draws the event tree for the process described by the dataset,
-        and saves it to "<filename>.filetype". Supports any filetype that
-        graphviz supports. e.g: "event_tree.png" or "event_tree.svg" etc.
-        """
-        filename, filetype = Util.generate_filename_and_mkdir(filename)
-        logger.info("--- generating graph ---")
-        graph = self.dot_graph
-        logger.info("--- writing " + filetype + " file ---")
-        graph.write(str(filename), format=filetype)
-
-        if get_ipython() is None:
-            return None
-        else:
-            logger.info("--- Exporting graph to notebook ---")
-            return Image(graph.create_png())
-
-    def __create_unsorted_paths_dict(self) -> defaultdict:
-        """Creates and populates a dictionary of all paths provided in the dataframe,
-        in the order in which they are given."""
-        unsorted_paths = defaultdict(int)
-
-        for variable_number in range(0, len(self.variables)):
-            dataframe_upto_variable = self.dataframe.loc[
-                :, self.variables[0:variable_number+1]]
-
-            for row in dataframe_upto_variable.itertuples():
-                row = row[1:]
-                new_row = [edge_label for edge_label in row if
-                           edge_label != np.nan and
-                           str(edge_label) != 'NaN' and
-                           str(edge_label) != 'nan' and
-                           edge_label != '']
-                new_row = tuple(new_row)
-
-                # checking if the last edge label in row was nan. That would
-                # result in double counting nan must be identified as string
-                if (row[-1] != np.nan and str(row[-1]) != 'NaN' and
-                   str(row[-1]) != 'nan' and row[-1] != ''):
-                    unsorted_paths[new_row] += 1
-
-        return unsorted_paths
-
-    def __create_path_dict_entries(self):
-        '''Create path dict entries for each path, including the
-        sampling zero paths if any.
-        Each path is an ordered sequence of edge labels starting
-        from the root.
-        The keys in the dict are ordered alphabetically.
-        Also calls the method self._sampling_zeros to ensure
-        manually added path format is correct.
-        Added functionality to remove NaN/null edge labels
-        assuming they are structural zeroes'''
-        unsorted_paths = self.__create_unsorted_paths_dict()
-
-        if self.sampling_zeros is not None:
-            unsorted_paths = Util.create_sampling_zeros(
-                self.sampling_zeros, unsorted_paths)
-
-        depth = len(max(list(unsorted_paths.keys()), key=len))
-        keys_of_list = list(unsorted_paths.keys())
-        sorted_keys = []
-        for deep in range(0, depth + 1):
-            unsorted_mini_list = [key for key in keys_of_list if
-                                  len(key) == deep]
-            sorted_keys = sorted_keys + sorted(unsorted_mini_list)
-
-        for key in sorted_keys:
-            self._sorted_paths[key] = unsorted_paths[key]
-
-        node_list = self.__create_node_list_from_paths(self._sorted_paths)
-        self.add_nodes_from(node_list)
-
-    def __check_sampling_zero_paths_param(self, sampling_zero_paths) -> list:
-        """Check param 'sampling_zero_paths' is in the correct format"""
-        for tup in sampling_zero_paths:
-            if not isinstance(tup, tuple):
-                return None
-            else:
-                if not Util.check_tuple_contains_strings(tup):
-                    return None
-
-        return sampling_zero_paths
-
-    def __create_node_list_from_paths(self, paths) -> list:
-        """Creates list of all nodes: includes root, situations, leaves"""
-        node_list = [self.root]
-
-        for vertex_number, _ in enumerate(list(paths.keys()), start=1):
-            node_list.append('s%d' % vertex_number)
-
-        return node_list
-
-    def __construct_event_tree(self):
-        """Constructs event_tree DiGraph.
-        Takes the paths, and adds all the nodes and edges to the Graph"""
-
-        logger.info('Starting construction of event tree')
-        self.__create_path_dict_entries()
-        # Taking a list of a networkx graph object (self) provides a list
-        # of all the nodes
-        node_list = list(self)
-
-        # Work through the sorted paths list to build the event tree.
-        edge_labels_list = ['root']
-        for path, count in list(self._sorted_paths.items()):
-            path = list(path)
-            edge_labels_list.append(path)
-            if path[:-1] in edge_labels_list:
-                path_edge_comes_from = edge_labels_list.index(path[:-1])
-                self.add_edge(
-                    u_for_edge=node_list[path_edge_comes_from],
-                    v_for_edge=node_list[edge_labels_list.index(path)],
-                    key=path[-1],
-                    count=count
-                )
-            else:
-                self.add_edge(
-                    u_for_edge=node_list[0],
-                    v_for_edge=node_list[edge_labels_list.index(path)],
-                    key=path[-1],
-                    count=count
-                )
diff --git a/src/cegpy/trees/staged.py b/src/cegpy/trees/staged.py
deleted file mode 100644
index 44a2e05..0000000
--- a/src/cegpy/trees/staged.py
+++ /dev/null
@@ -1,580 +0,0 @@
-from copy import deepcopy
-from fractions import Fraction
-from operator import add, sub, itemgetter
-from IPython.display import Image
-from IPython import get_ipython
-from itertools import combinations, chain
-from typing import List, Tuple
-import networkx as nx
-import scipy.special
-import logging
-from ..utilities.util import Util
-from ..trees.event import EventTree
-
-logger = logging.getLogger('cegpy.staged_tree')
-
-
-class StagedTree(EventTree):
-    def __init__(
-            self,
-            dataframe,
-            prior=None,
-            alpha=None,
-            hyperstage=None,
-            sampling_zero_paths=None,
-            incoming_graph_data=None,
-            **attr) -> None:
-
-        # Call event tree init to generate event tree
-        super().__init__(
-            dataframe,
-            sampling_zero_paths,
-            incoming_graph_data,
-            **attr
-        )
-
-        self.__store_params(prior, alpha, hyperstage)
-        self._mean_posterior_probs = []
-        self._merged_situations = []
-        self._stage_colours = []
-        self._sort_count = 0
-        self._colours_for_situations = []
-        logger.debug("Starting Staged Tree")
-
-    @property
-    def prior(self):
-        return nx.get_edge_attributes(self, 'prior')
-
-    @prior.setter
-    def prior(self, prior):
-        offset = 0
-        for node_idx, node_priors in enumerate(prior):
-            node_name = ('s%d' % (node_idx + offset))
-            while self.succ[node_name] == {}:
-                offset += 1
-                node_name = ('s%d' % (node_idx + offset))
-            pass
-            for edge_prior_idx, succ_key in enumerate(
-                    self.succ[node_name].keys()):
-                label = list(self.succ[node_name][succ_key])[0]
-                self.edges[(node_name, succ_key, label)]['prior'] = \
-                    node_priors[edge_prior_idx]
-
-    @property
-    def prior_list(self):
-        """Priors provided as a list of lists"""
-        prior_list = []
-        prev_node = list(self.prior)[0][0]
-        succ_list = []
-
-        for edge, prior in self.prior.items():
-            node = edge[0]
-            if node != prev_node:
-                prior_list.append(succ_list)
-                succ_list = []
-            succ_list.append(prior)
-            prev_node = node
-
-        if succ_list != []:
-            prior_list.append(succ_list)
-        return prior_list
-
-    @property
-    def posterior(self):
-        '''Posterior is calculated such that the edge count is added
-        to the prior for each edge.'''
-        try:
-            posterior = nx.get_edge_attributes(self, 'posterior')
-            if posterior == {}:
-                raise AttributeError('Posterior not yet set.')
-            else:
-                return posterior
-        except AttributeError:
-            for edge in self.edges:
-                edge_dict = self.edges[edge]
-                posterior = edge_dict['prior'] + edge_dict['count']
-                edge_dict['posterior'] = posterior
-            return nx.get_edge_attributes(self, 'posterior')
-
-    @property
-    def posterior_list(self):
-        posterior_list = []
-        prev_node = list(self.posterior)[0][0]
-        succ_list = []
-
-        for edge, posterior in self.posterior.items():
-            node = edge[0]
-            if node != prev_node:
-                posterior_list.append(succ_list)
-                succ_list = []
-            succ_list.append(posterior)
-            prev_node = node
-
-        if succ_list != []:
-            posterior_list.append(succ_list)
-        return posterior_list
-
-    @property
-    def alpha(self):
-        return self._alpha
-
-    @alpha.setter
-    def alpha(self, value):
-        self._alpha = value
-
-    @property
-    def hyperstage(self):
-        return self._hyperstage
-
-    @hyperstage.setter
-    def hyperstage(self, value):
-        self._hyperstage = value
-
-    @property
-    def edge_countset(self):
-        return self.__create_edge_countset()
-
-    @property
-    def ahc_output(self):
-        return self._ahc_output
-
-    @ahc_output.setter
-    def ahc_output(self, value):
-        self._ahc_output = value
-
-    def __store_params(self, prior, alpha, hyperstage) -> None:
-        """User has passed in AHC params, this function processes them,
-        and generates any default AHC params if required."""
-        if prior:
-            if alpha:
-                self.alpha = None
-                logging.warning("Params Warning!! When prior is given, " +
-                                "alpha is not required!")
-        else:
-            if alpha is None:
-                self.alpha = self.__calculate_default_alpha()
-                logging.warning("Params Warning!! Neither prior nor alpha " +
-                                "were provided. Using default alpha " +
-                                "value of %d.", self.alpha)
-            else:
-                self.alpha = alpha
-
-            # No matter what alpha is, generate default prior
-            self.prior = self.__create_default_prior(self.alpha)
-
-        if hyperstage is None:
-            self.hyperstage = self.__create_default_hyperstage()
-        else:
-            self.hyperstage = hyperstage
-
-    def __calculate_default_alpha(self) -> int:
-        """If no alpha is given, a default value is calculated.
-        The value is calculated by determining the maximum number
-        of categories that any one variable has"""
-        logger.info("Calculating default prior")
-        max_count = max(list(self.categories_per_variable.values()))
-        return max_count
-
-    def __create_default_prior(self, alpha) -> list:
-        """default prior set for the AHC method using the mass conservation property.
-        That is, the alpha param is the phantom sample starting at
-        the root, and it is spread equally across all edges along the tree.
-        (see chapter 5 of Collazo, Gorgen & Smith 'Chain Event Graphs', 2018)
-        The prior is a list of lists. Each list gives the prior along the
-        edges of a specific situation.
-        Indexed same as self.situations & self.egde_countset"""
-
-        logger.info("Generating default prior")
-        default_prior = [0] * len(self.situations)
-        sample_size_at_node = dict()
-
-        # Root node is assigned phantom sample (alpha)
-        if isinstance(alpha, float):
-            alpha = Fraction.from_float(alpha)
-        elif isinstance(alpha, int) or isinstance(alpha, str):
-            alpha = Fraction.from_float(float(alpha))
-        else:
-            logger.warning("Prior generator param alpha is in a strange format.\
-                            ..")
-
-        sample_size_at_node[self.root] = alpha
-
-        for node_idx, node in enumerate(self.situations):
-            # How many nodes emanate from the current node?
-            number_of_emanating_nodes = self.out_degree[node]
-
-            # Divide the sample size from the current node equally among
-            # emanating nodes
-            equal_distribution_of_sample = \
-                sample_size_at_node[node] / number_of_emanating_nodes
-            default_prior[node_idx] = \
-                [equal_distribution_of_sample] * number_of_emanating_nodes
-
-            relevant_terminating_nodes = [
-                edge[1] for edge in list(self.edges) if edge[0] == node
-            ]
-
-            for terminating_node in relevant_terminating_nodes:
-                sample_size_at_node[terminating_node] = \
-                 equal_distribution_of_sample
-
-        return default_prior
-
-    def __create_default_hyperstage(self) -> list:
-        '''Generates default hyperstage for the AHC method.
-        A hyperstage is a list of lists such that two situaions can be in the
-        same stage only if there are elements of the same list for some list
-        in the hyperstage.
-        The default is to allow all situations with the same number of
-        outgoing edges and the same edge labels to be in a common list. '''
-        logger.info("Creating default hyperstage")
-        hyperstage = []
-        info_of_edges = []
-
-        for node in self.situations:
-            labels = [
-                edge[2] for edge in self.edges
-                if edge[0] == node
-            ]
-            labels.sort()
-
-            info_of_edges.append([self.out_degree[node], labels])
-
-        sorted_info = []
-        for info in info_of_edges:
-            if info not in sorted_info:
-                sorted_info.append(info)
-
-        for info in sorted_info:
-            situations_with_value_edges = []
-            for idx, situation in enumerate(self.situations):
-                if info_of_edges[idx] == info:
-                    situations_with_value_edges.append(situation)
-            hyperstage = hyperstage + [situations_with_value_edges]
-
-        return hyperstage
-
-    def __create_edge_countset(self) -> list:
-        '''Each element of list contains a list with counts along edges emanating from
-        a specific situation. Indexed same as self.situations'''
-        logger.info("Creating edge countset")
-        edge_countset = []
-
-        for node in self.situations:
-            edge_countset.append([
-                count for edge, count in self.edge_counts.items()
-                if edge[0] == node
-            ])
-        return edge_countset
-
-    def _calculate_lg_of_sum(self, array) -> float:
-        '''function to calculate log gamma of the sum of an array'''
-        array = [float(x) for x in array]
-        return scipy.special.gammaln(sum(array))
-
-    def _calculate_sum_of_lg(self, array) -> float:
-        '''function to calculate log gamma of each element of an array'''
-        return sum([scipy.special.gammaln(float(x)) for x in array])
-
-    def _calculate_initial_loglikelihood(self, prior, posterior) -> float:
-        '''calculating log likelihood given a prior and posterior'''
-        # Calculate prior contribution
-        logger.info("Calculating initial loglikelihood")
-
-        pri_lg_of_sum = [
-            self._calculate_lg_of_sum(elem) for elem in prior
-        ]
-        pri_sum_of_lg = [
-            self._calculate_sum_of_lg(elem) for elem in prior
-        ]
-        pri_contribution = list(map(sub, pri_lg_of_sum, pri_sum_of_lg))
-
-        # Calculate posterior contribution
-        post_lg_of_sum = [
-            self._calculate_lg_of_sum(elem) for elem in posterior
-        ]
-        post_sum_of_lg = [
-            self._calculate_sum_of_lg(elem) for elem in posterior
-        ]
-        post_contribution = list(map(sub, post_sum_of_lg, post_lg_of_sum))
-
-        sum_pri_contribution = sum(pri_contribution)
-        sum_post_contribution = sum(post_contribution)
-        return (sum_pri_contribution + sum_post_contribution)
-
-    def _check_issubset(self, item, hyperstage) -> bool:
-        '''function to check if two situations belong to the same set in the
-        hyperstage'''
-        return any(set(item).issubset(element) for element in hyperstage)
-
-    def _calculate_bayes_factor(self, prior1, posterior1,
-                                prior2, posterior2) -> float:
-        '''calculates the bayes factor comparing two models which differ in
-        only one stage'''
-        new_prior = list(map(add, prior1, prior2))
-        new_posterior = list(map(add, posterior1, posterior2))
-        return (
-            self._calculate_lg_of_sum(new_prior)
-            - self._calculate_lg_of_sum(prior1)
-            - self._calculate_lg_of_sum(prior2)
-            - self._calculate_lg_of_sum(new_posterior)
-            + self._calculate_lg_of_sum(posterior1)
-            + self._calculate_lg_of_sum(posterior2)
-            + self._calculate_sum_of_lg(new_posterior)
-            - self._calculate_sum_of_lg(posterior1)
-            - self._calculate_sum_of_lg(posterior2)
-            - self._calculate_sum_of_lg(new_prior)
-            + self._calculate_sum_of_lg(prior1)
-            + self._calculate_sum_of_lg(prior2)
-        )
-
-    def _sort_list(self, list_of_tuples) -> list:
-        '''function to sort a list of lists to remove repetitions'''
-
-        for l1_idx in range(0, len(list_of_tuples)):
-            for l2_idx in range(l1_idx+1, len(list_of_tuples)):
-                tup_1 = list_of_tuples[l1_idx]
-                tup_2 = list_of_tuples[l2_idx]
-                tups_intersect = set(tup_1) & set(tup_2)
-
-                if tups_intersect:
-                    union = tuple(set(tup_1) | set(tup_2))
-                    list_of_tuples[l1_idx] = []
-                    list_of_tuples[l2_idx] = union
-
-        new_list_of_tuples = [
-            elem for elem in list_of_tuples
-            if elem != []
-        ]
-
-        if new_list_of_tuples == list_of_tuples:
-            return new_list_of_tuples
-        else:
-            return self._sort_list(new_list_of_tuples)
-
-    def _calculate_mean_posterior_probs(
-        self,
-        merged_situations: List,
-        posteriors: List,
-    ) -> List:
-        '''Iterates through array of lists, calculates mean
-        posterior probabilities'''
-        mean_posterior_probs = []
-
-        for sit in self.situations:
-            if sit not in list(chain(*merged_situations)):
-                merged_situations.append((sit,))
-        for stage in merged_situations:
-            for sit in stage:
-                sit_idx = self.situations.index(sit)
-                if all(posteriors[sit_idx]) != 0:
-                    stage_probs = posteriors[sit_idx]
-                    break
-                else:
-                    stage_probs = []
-
-            total = sum(stage_probs)
-            mean_posterior_probs.append(
-                [round(elem/total, 3) for elem in stage_probs]
-            )
-
-        return mean_posterior_probs
-
-    def _independent_hyperstage_generator(
-            self, hyperstage: List[List]) -> List[List[List]]:
-        """Spit out the next hyperstage that can be dealt with
-        independently."""
-        new_hyperstages = [[hyperstage[0]]]
-
-        for sublist in hyperstage[1:]:
-            hs_to_add = [sublist]
-
-            for hs in new_hyperstages.copy():
-                for other_sublist in hs:
-                    if not set(other_sublist).isdisjoint(set(sublist)):
-                        hs_to_add.extend(hs)
-                        new_hyperstages.remove(hs)
-
-            new_hyperstages.append(hs_to_add)
-
-        return new_hyperstages
-
-    def _execute_AHC(self, hyperstage=None) -> Tuple[List, float, List]:
-        """finds all subsets and scores them"""
-        if hyperstage is None:
-            hyperstage = deepcopy(self.hyperstage)
-
-        priors = deepcopy(self.prior_list)
-        posteriors = deepcopy(self.posterior_list)
-
-        loglikelihood = self._calculate_initial_loglikelihood(
-            priors, posteriors
-        )
-
-        merged_situation_list = []
-
-        # Which list in hyperstage have only 1 edge coming out of them
-        # For that list, add the list to the merged situation list, and
-        # add together their priors/posteriors, and remove it from the
-        # hyperstage
-        for sub_hyper in deepcopy(hyperstage):
-            if self.out_degree[sub_hyper[0]] == 1:
-                merged_situation_list.append(tuple(sub_hyper))
-                indexes = [self.situations.index(situ) for situ in sub_hyper]
-                sub_hyper_priors = list(chain(
-                    *itemgetter(*indexes)(priors)
-                ))
-                sub_hyper_posteriors = list(chain(
-                    *itemgetter(*indexes)(posteriors)
-                ))
-                for loop_idx, p_index in enumerate(indexes):
-                    if loop_idx == 0:
-                        priors[p_index] = [sum(sub_hyper_priors)]
-                        posteriors[p_index] = [sum(sub_hyper_posteriors)]
-                    else:
-                        priors[p_index] = [0]
-                        posteriors[p_index] = [0]
-
-                hyperstage.remove(sub_hyper)
-                print("blash")
-
-        hyperstage_combinations = [
-            item for sub_hyper in hyperstage
-            for item in combinations(sub_hyper, 2)
-        ]
-
-        while True:
-            newscores_list = [self._calculate_bayes_factor(
-                priors[self.situations.index(sub_hyper[0])],
-                posteriors[self.situations.index(sub_hyper[0])],
-                priors[self.situations.index(sub_hyper[1])],
-                posteriors[self.situations.index(sub_hyper[1])],
-            ) for sub_hyper in hyperstage_combinations]
-
-            local_score = max(newscores_list)
-
-            if local_score > 0:
-                local_merged = hyperstage_combinations[
-                    newscores_list.index(local_score)
-                ]
-                merge_situ_1, merge_situ_2 = local_merged
-                merge_situ_1_idx = self.situations.index(merge_situ_1)
-                merge_situ_2_idx = self.situations.index(merge_situ_2)
-                merged_situation_list.append(local_merged)
-
-                priors[merge_situ_1_idx] = list(
-                    map(
-                        add,
-                        priors[merge_situ_1_idx],
-                        priors[merge_situ_2_idx]
-                    )
-                )
-                posteriors[merge_situ_1_idx] = list(
-                    map(
-                        add,
-                        posteriors[merge_situ_1_idx],
-                        posteriors[merge_situ_2_idx]
-                    )
-                )
-                priors[merge_situ_2_idx] = (
-                    [0] * len(priors[merge_situ_1_idx]))
-                posteriors[merge_situ_2_idx] = (
-                    [0] * len(posteriors[merge_situ_1_idx]))
-
-                loglikelihood += local_score
-            else:
-                break
-
-        merged_situation_list = self._sort_list(merged_situation_list)
-        mean_posterior_probs = (
-            self._calculate_mean_posterior_probs(
-                merged_situation_list,
-                posteriors
-            )
-        )
-        return mean_posterior_probs, loglikelihood, merged_situation_list
-
-    def _mark_nodes_with_stage_number(self, merged_situations):
-        """AHC algorithm creates a list of indexes to the situations list.
-        This function takes those indexes and creates a new list which is
-        in a string representation of nodes."""
-        self._sort_count = 0
-        for index, stage in enumerate(merged_situations):
-            if len(stage) > 1:
-                for node in stage:
-                    self.nodes[node]['stage'] = index
-
-    def _generate_colours_for_situations(self, merged_situations, colour_list):
-        """Colours each stage of the tree with an individual colour"""
-        number_of_stages = len(merged_situations)
-        if colour_list is None:
-            stage_colours = Util.generate_colours(number_of_stages)
-        else:
-            stage_colours = colour_list
-            if len(colour_list) < number_of_stages:
-                raise IndexError(
-                    f"The number of colours provided ({len(colour_list)}) is "
-                    f"less than the number of stages ({number_of_stages}). "
-                )
-        self._stage_colours = stage_colours
-        for node in self.nodes:
-            try:
-                stage = self.nodes[node]['stage']
-                self.nodes[node]['colour'] = stage_colours[stage]
-            except KeyError:
-                self.nodes[node]['colour'] = 'lightgrey'
-
-    def calculate_AHC_transitions(self, prior=None,
-                                  alpha=None, hyperstage=None,
-                                  colour_list=None):
-        '''Bayesian Agglommerative Hierarchical Clustering algorithm
-        implementation. It returns a list of lists of the situations which
-        have been merged together, the likelihood of the final model and
-        the mean posterior conditional probabilities of the stages.
-
-        User can specify a list of colours to be used for stages. Otherwise,
-        colours evenly spaced around the colour spectrum are used.'''
-        logger.info("\n\n --- Starting AHC Algorithm ---")
-
-        self.__store_params(prior, alpha, hyperstage)
-
-        mean_posterior_probs, loglikelihood, merged_situations = (
-            self._execute_AHC())
-
-        self._mark_nodes_with_stage_number(merged_situations)
-
-        self._generate_colours_for_situations(merged_situations, colour_list)
-
-        self.ahc_output = {
-            "Merged Situations": merged_situations,
-            "Loglikelihood": loglikelihood,
-            "Mean Posterior Probabilities": mean_posterior_probs
-        }
-        return self.ahc_output
-
-    def create_figure(self, filename):
-        """Draws the coloured staged tree for the process described by
-        the dataset, and saves it to "<filename>.filetype". Supports
-        any filetype that graphviz supports. e.g: "event_tree.png" or
-        "event_tree.svg" etc.
-        """
-        try:
-            self.ahc_output
-            filename, filetype = Util.generate_filename_and_mkdir(filename)
-            logger.info("--- generating graph ---")
-            graph = self.dot_graph
-            logger.info("--- writing " + filetype + " file ---")
-            graph.write(str(filename), format=filetype)
-
-            if get_ipython() is None:
-                return None
-            else:
-                logger.info("--- Exporting graph to notebook ---")
-                return Image(graph.create_png())
-
-        except AttributeError:
-            logger.error(
-                "----- PLEASE RUN AHC ALGORITHM before trying to" +
-                " export graph -----"
-            )
-            return None
diff --git a/src/cegpy/utilities/__init__.py b/src/cegpy/utilities/__init__.py
index e69de29..be67bc3 100644
--- a/src/cegpy/utilities/__init__.py
+++ b/src/cegpy/utilities/__init__.py
@@ -0,0 +1,11 @@
+"""Shared utilities for the cegpy package."""
+
+from cegpy.utilities._util import (
+    check_list_contains_strings,
+    check_tuple_contains_strings,
+    generate_colours,
+    generate_colour_run,
+    generate_filename_and_mkdir,
+    create_relative_path,
+    create_sampling_zeros,
+)
diff --git a/src/cegpy/utilities/_util.py b/src/cegpy/utilities/_util.py
new file mode 100644
index 0000000..ba56a7c
--- /dev/null
+++ b/src/cegpy/utilities/_util.py
@@ -0,0 +1,135 @@
+"""Utilities for use throughout the package."""
+
+from collections import defaultdict
+import os
+from pathlib import Path
+import math
+from typing import List, Union
+import colorutils
+from colorutils import Color as Colour
+
+ST_OUTPUT = {
+    "Merged Situations": [],
+    "Loglikelihood": float,
+    "Mean Posterior Probabilities": [],
+}
+
+
+def check_list_contains_strings(str_list) -> bool:
+    """Ensure that a list only contains strings"""
+    for tup in str_list:
+        if not isinstance(tup, str):
+            return False
+    return True
+
+
+def check_tuple_contains_strings(tup) -> bool:
+    """Check each element of the tuple to ensure it is a string"""
+    if isinstance(tup, tuple):
+        for elem in tup:
+            if not isinstance(elem, str):
+                return False
+        return True
+    else:
+        return False
+
+
+def generate_colours(number) -> list:
+    """Generates a list of hex colour strings that are evenly spaced
+    around the colour spectrum"""
+    hex_colours = []
+    low, start, high = 100, 150, 255
+    starts = [
+        Colour((start, low, low)),
+        Colour((low, start, low)),
+        Colour((low, low, start)),
+    ]
+    ends = [
+        Colour((high, low, low)),
+        Colour((low, high, low)),
+        Colour((low, low, high)),
+    ]
+    prim_colours = math.ceil(number / 2)
+    hex_colours += generate_colour_run(prim_colours, starts, ends)
+    starts = [
+        Colour((start, start, low)),
+        Colour((low, start, start)),
+        Colour((start, low, start)),
+    ]
+    ends = [
+        Colour((high, high, low)),
+        Colour((low, high, high)),
+        Colour((high, low, high)),
+    ]
+    sec_colours = math.floor(number / 2)
+    hex_colours += generate_colour_run(sec_colours, starts, ends)
+
+    return hex_colours
+
+
+def generate_colour_run(number, starts, ends) -> list:
+    """Generates multiple colour runs with uniform distance between colours."""
+    split = [0, 0, 0]
+    for i in range(number):
+        split[i % len(split)] += 1
+
+    jumps = []
+    for val in split:
+        jumps.append(max(val - 1, 0))
+
+    colours: List[Colour] = []
+
+    for idx, end in enumerate(ends):
+        if jumps[idx] > 0:
+            colours += colorutils.color_run(starts[idx], end, jumps[idx], to_color=True)
+        else:
+            colours.append(end)
+    hex_colours = []
+    for colour in colours:
+        hex_colours.append(colour.hex)
+
+    return hex_colours
+
+
+def generate_filename_and_mkdir(filename: Union[str, Path]) -> str:
+    """Creates a filename."""
+    if filename is not Path:
+        filename = Path(filename)
+
+    filetype = filename.suffix.strip(".")
+    if filename.suffix == filetype:
+        filename.joinpath(".png")
+        filetype = "png"
+
+    os.makedirs(filename.parent, exist_ok=True)
+    return filename, filetype
+
+
+def create_relative_path(relative_path) -> Path:
+    """Creates pathlib object from a relative path such as:
+    '/output/image.png'"""
+    return Path(__file__).resolve().parent.joinpath(relative_path)
+
+
+def create_sampling_zeros(sampling_zero_paths, path_dict) -> defaultdict:
+    """The list of paths to zero must only contain tuples.
+    Each tuple is a sampling zero path that needs to be added.
+    If multiple edges along a path need to be added, they must be
+    added in order. i.e. path[:-1] should already be a key in self.paths
+    Eg suppose edge 'eat' already exists as an edge emanating from the
+    root.
+    To add paths ('eat', 'sleep'), ('eat', 'sleep', 'repeat'),
+    the sampling zero parameter should be added as:
+    et = ceg({'dataframe': df, 'sampling_zero_paths':
+    [('eat', 'sleep',),('eat', 'sleep','repeat',)]})
+    """
+    for path in sampling_zero_paths:
+        if (path[:-1] in list(path_dict.keys())) or len(path) == 1:
+            path_dict[path] = 0
+        else:
+            raise ValueError(
+                "The path up to it's last edge should be \
+                added first. Ensure the tuple ends with a comma."
+            )
+
+    return path_dict
diff --git a/src/cegpy/utilities/util.py b/src/cegpy/utilities/util.py
deleted file mode 100644
index a3ca482..0000000
--- a/src/cegpy/utilities/util.py
+++ /dev/null
@@ -1,150 +0,0 @@
-import os
-import colorutils
-import math
-from collections import defaultdict
-from pathlib import Path
-from colorutils import Color as Colour
-from datetime import datetime
-
-ST_OUTPUT = {
-    "Merged Situations": [],
-    "Loglikelihood": float,
-    "Mean Posterior Probabilities": [],
-}
-
-
-class Util:
-    @staticmethod
-    def check_list_contains_strings(str_list) -> bool:
-        """Ensure that a list only contains strings"""
-        for tup in str_list:
-            if not isinstance(tup, str):
-                return False
-        return True
-
-    @staticmethod
-    def check_tuple_contains_strings(tup) -> bool:
-        """Check each element of the tuple to ensure it is a string"""
-        if isinstance(tup, tuple):
-            for elem in tup:
-                if not isinstance(elem, str):
-                    return False
-            return True
-        else:
-            return False
-
-    @staticmethod
-    def generate_colours(number) -> list:
-        """Generates a list of hex colour strings that are evenly spaced
-        around the colour spectrum"""
-        hex_colours = []
-        low, start, high = 100, 150, 255
-        starts = [
-            Colour((start, low, low)),
-            Colour((low, start, low)),
-            Colour((low, low, start))
-        ]
-        ends = [
-            Colour((high, low, low)),
-            Colour((low, high, low)),
-            Colour((low, low, high))
-        ]
-        prim_colours = math.ceil(number / 2)
-        hex_colours += Util.generate_colour_run(prim_colours, starts, ends)
-        starts = [
-            Colour((start, start, low)),
-            Colour((low, start, start)),
-            Colour((start, low, start))
-        ]
-        ends = [
-            Colour((high, high, low)),
-            Colour((low, high, high)),
-            Colour((high, low, high))
-        ]
-        sec_colours = math.floor(number / 2)
-        hex_colours += Util.generate_colour_run(sec_colours, starts, ends)
-
-        return hex_colours
-
-    @staticmethod
-    def generate_colour_run(number, starts, ends) -> list:
-        split = [0, 0, 0]
-        for i in range(number):
-            split[i % len(split)] += 1
-
-        jumps = []
-        for val in split:
-            jumps.append(max(val - 1, 0))
-
-        colours = []
-
-        for idx in range(len(ends)):
-            if jumps[idx] > 0:
-                colours += colorutils.color_run(
-                    starts[idx], ends[idx], jumps[idx], to_color=True)
-            else:
-                colours.append(ends[idx])
-        hex_colours = []
-        for colour in colours:
-            hex_colours.append(colour.hex)
-
-        return hex_colours
-
-    @staticmethod
-    def generate_filename_and_mkdir(filename) -> str:
-        if filename is not Path:
-            filename = Path(filename)
-
-        filetype = filename.suffix.strip('.')
-        if filename.suffix == filetype:
-            filename.joinpath('.png')
-            filetype = 'png'
-
-        os.makedirs(filename.parent, exist_ok=True)
-        return filename, filetype
-
-    @staticmethod
-    def create_relative_path(relative_path) -> Path:
-        """Creates pathlib object from a relative path such as:
-        '/output/image.png'"""
-        return Path(__file__).resolve().parent.joinpath(relative_path)
-
-    @staticmethod
-    def create_sampling_zeros(sampling_zero_paths, path_dict) -> defaultdict:
-        '''The list of paths to zero must only contain tuples.
-        Each tuple is a sampling zero path that needs to be added.
-        If multiple edges along a path need to be added, they must be
-        added in order. i.e. path[:-1] should already be a key in self.paths
-        Eg suppose edge 'eat' already exists as an edge emanating from the
-        root.
-        To add paths ('eat', 'sleep'), ('eat', 'sleep', 'repeat'),
-        the sampling zero parameter should be added as:
-        et = ceg({'dataframe': df, 'sampling_zero_paths':
-        [('eat', 'sleep',),('eat', 'sleep','repeat',)]})
-        '''
-        for path in sampling_zero_paths:
-            if (path[:-1] in list(path_dict.keys())) or len(path) == 1:
-                path_dict[path] = 0
-            else:
-                raise ValueError("The path up to it's last edge should be \
-                    added first. Ensure the tuple ends with a comma.")
-
-        return path_dict
-
-    @staticmethod
-    def create_path(filename, add_time=False, filetype='png'):
-        dt_string = ''
-        if add_time:
-            now = datetime.now()
-            dt_string = now.strftime("__%d-%m-%Y_%H-%M-%S")
-        fig_path = Path(__file__).resolve().parent\
-            .parent.parent.parent.joinpath(
-                filename + dt_string + '.' + filetype)
-        return fig_path
-
-    @staticmethod
-    def flatten_list_of_lists(list_of_lists) -> list:
-        flat_list = []
-        for sublist in list_of_lists:
-            flat_list = flat_list + sublist
-        return flat_list


Test patch:
diff --git a/integration/test_integration.py b/integration/test_integration.py
deleted file mode 100644
index 2f26b92..0000000
--- a/integration/test_integration.py
+++ /dev/null
@@ -1,279 +0,0 @@
-from pathlib import Path
-import networkx as nx
-import pandas as pd
-import os.path
-from src.cegpy.graphs.ceg import ChainEventGraph
-from src.cegpy.graphs.ceg import Evidence
-from src.cegpy.trees.staged import StagedTree
-from src.cegpy.utilities.util import Util
-
-
-class TestReducedOutputs:
-    def setup(self):
-        G = nx.MultiDiGraph()
-        nodes = ['w0', 'w1', 'w2', 'w3', 'w4', 'w5',
-                 'w6', 'w7', 'w8', 'w&infin;']
-        G.add_edge('w0', 'w1', 'a', probability=0.3)
-        G.add_edge('w0', 'w1', 'b', probability=0.3)
-        G.add_edge('w0', 'w5', 'j', probability=0.2)
-        G.add_edge('w0', 'w7', 'k', probability=0.05)
-        G.add_edge('w0', 'w6', 'l', probability=0.05)
-        G.add_edge('w0', 'w3', 's', probability=0.1)
-
-        G.add_edge('w1', 'w2', 'c', probability=0.88)
-        G.add_edge('w1', 'w3', 'd', probability=0.12)
-
-        G.add_edge('w2', 'w4', 'e', probability=0.95)
-        G.add_edge('w2', 'w4', 'f', probability=0.05)
-
-        G.add_edge('w3', 'w4', 'g', probability=0.5)
-        G.add_edge('w3', 'w8', 'h', probability=0.5)
-
-        G.add_edge('w4', 'w8', 'i', probability=0.9)
-        G.add_edge('w4', 'w&infin;', 'r', probability=0.1)
-
-        G.add_edge('w5', 'w6', 'm', probability=0.6)
-        G.add_edge('w5', 'w7', 'n', probability=0.4)
-
-        G.add_edge('w6', 'w7', 'o', probability=1)
-
-        G.add_edge('w7', 'w8', 'p', probability=1)
-
-        G.add_edge('w8', 'w&infin;', 'q', probability=1)
-
-        G.add_nodes_from(nodes)
-        H = ChainEventGraph(G)
-        self.evidence = Evidence(H)
-        print(self.evidence)
-
-    def test_subgraph(self):
-        def update_path_lists():
-            self.evidence._Evidence__graph._ChainEventGraph__update_path_list()
-            self.evidence._Evidence__update_path_list()
-
-        fname = Util.create_path('out/Evidence_pre_test', False, 'pdf')
-        self.evidence._Evidence__graph.create_figure(fname)
-        update_path_lists()
-        self.evidence.add_edge('w8', 'w&infin;', 'q', Evidence.CERTAIN)
-        self.evidence.add_node('w4', Evidence.CERTAIN)
-        reduced = self.evidence.reduced_graph
-
-        expected_reduced_probabilities = {
-            ('w0', 'w1', 'a'): 0.46,
-            ('w0', 'w1', 'b'): 0.46,
-            ('w0', 'w3', 's'): 0.08,
-            ('w1', 'w2', 'c'): 0.94,
-            ('w1', 'w3', 'd'): 0.06,
-            ('w3', 'w4', 'g'): 1.0,
-            ('w2', 'w4', 'e'): 0.95,
-            ('w2', 'w4', 'f'): 0.05,
-            ('w4', 'w8', 'i'): 1.0,
-            ('w8', 'w&infin;', 'q'): 1.0
-        }
-        actual_reduced_probabilities = nx.get_edge_attributes(
-            reduced,
-            'probability'
-        )
-
-        for edge, probability in expected_reduced_probabilities.items():
-            try:
-                actual = round(actual_reduced_probabilities[edge], 2)
-                assert actual == probability
-            except KeyError:
-                assert False  # edge doesnt exist in graph!
-
-        fname = Util.create_path('out/Evidence_test_certain', False, 'pdf')
-        reduced.create_figure(fname)
-
-        expected_paths = [
-            [('w0', 'w1', 'a'), ('w1', 'w2', 'c'), ('w2', 'w4', 'e'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'a'), ('w1', 'w2', 'c'), ('w2', 'w4', 'f'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'b'), ('w1', 'w2', 'c'), ('w2', 'w4', 'f'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'b'), ('w1', 'w2', 'c'), ('w2', 'w4', 'e'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'a'), ('w1', 'w3', 'd'), ('w3', 'w4', 'g'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'b'), ('w1', 'w3', 'd'), ('w3', 'w4', 'g'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w3', 's'), ('w3', 'w4', 'g'), ('w4', 'w8', 'i'),
-                ('w8', 'w&infin;', 'q')]
-        ]
-        for path in expected_paths:
-            assert path in self.evidence.path_list
-        assert len(expected_paths) == len(self.evidence.path_list)
-
-        self.evidence.remove_edge('w8', 'w&infin;', 'q', Evidence.CERTAIN)
-        self.evidence.remove_node('w4', Evidence.CERTAIN)
-        update_path_lists()
-        self.evidence.add_node('w6', Evidence.UNCERTAIN)
-        self.evidence.add_node('w3', Evidence.UNCERTAIN)
-        reduced = self.evidence.reduced_graph
-
-        expected_reduced_probabilities = {
-            ('w0', 'w1', 'a'): 0.105,
-            ('w0', 'w1', 'b'): 0.105,
-            ('w0', 'w5', 'j'): 0.351,
-            ('w0', 'w6', 'l'): 0.146,
-            ('w0', 'w3', 's'): 0.292,
-            ('w1', 'w3', 'd'): 1.0,
-            ('w5', 'w6', 'm'): 1.0,
-            ('w7', 'w8', 'p'): 1.0,
-            ('w6', 'w7', 'o'): 1.0,
-            ('w3', 'w4', 'g'): 0.5,
-            ('w3', 'w8', 'h'): 0.5,
-            ('w4', 'w8', 'i'): 0.9,
-            ('w4', 'w&infin;', 'r'): 0.1,
-            ('w8', 'w&infin;', 'q'): 1.0
-        }
-        actual_reduced_probabilities = nx.get_edge_attributes(
-            reduced,
-            'probability'
-        )
-
-        for edge, probability in expected_reduced_probabilities.items():
-            try:
-                actual = round(actual_reduced_probabilities[edge], 3)
-                assert actual == probability
-            except KeyError:
-                assert False  # edge doesnt exist in graph!
-
-        fname = Util.create_path('out/Evidence_test_uncertain', False, 'pdf')
-        reduced.create_figure(fname)
-
-        expected_paths = [
-            [('w0', 'w1', 'a'), ('w1', 'w3', 'd'), ('w3', 'w4', 'g'),
-                ('w4', 'w&infin;', 'r')],
-            [('w0', 'w1', 'b'), ('w1', 'w3', 'd'), ('w3', 'w4', 'g'),
-                ('w4', 'w&infin;', 'r')],
-            [('w0', 'w1', 'a'), ('w1', 'w3', 'd'), ('w3', 'w4', 'g'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'b'), ('w1', 'w3', 'd'), ('w3', 'w4', 'g'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'a'), ('w1', 'w3', 'd'), ('w3', 'w8', 'h'),
-                ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'b'), ('w1', 'w3', 'd'), ('w3', 'w8', 'h'),
-                ('w8', 'w&infin;', 'q')],
-            [('w0', 'w5', 'j'), ('w5', 'w6', 'm'), ('w6', 'w7', 'o'),
-                ('w7', 'w8', 'p'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w6', 'l'), ('w6', 'w7', 'o'), ('w7', 'w8', 'p'),
-                ('w8', 'w&infin;', 'q')],
-            [('w0', 'w3', 's'), ('w3', 'w4', 'g'), ('w4', 'w8', 'i'),
-                ('w8', 'w&infin;', 'q')],
-            [('w0', 'w3', 's'), ('w3', 'w4', 'g'), ('w4', 'w&infin;', 'r')],
-            [('w0', 'w3', 's'), ('w3', 'w8', 'h'), ('w8', 'w&infin;', 'q')]
-        ]
-        for path in expected_paths:
-            assert path in self.evidence.path_list
-        assert len(expected_paths) == len(self.evidence.path_list)
-
-        self.evidence.remove_node('w3', Evidence.UNCERTAIN)
-        self.evidence.remove_node('w6', Evidence.UNCERTAIN)
-        update_path_lists()
-        self.evidence.add_edge('w2', 'w4', 'e', Evidence.UNCERTAIN)
-        self.evidence.add_edge('w3', 'w4', 'g', Evidence.UNCERTAIN)
-        self.evidence.add_edge('w3', 'w8', 'h', Evidence.UNCERTAIN)
-        self.evidence.add_node('w1', Evidence.CERTAIN)
-
-        reduced = self.evidence.reduced_graph
-
-        expected_reduced_probabilities = {
-            ('w0', 'w1', 'a'): 0.5,
-            ('w0', 'w1', 'b'): 0.5,
-            ('w1', 'w2', 'c'): 0.87,
-            ('w1', 'w3', 'd'): 0.13,
-            ('w3', 'w4', 'g'): 0.5,
-            ('w3', 'w8', 'h'): 0.5,
-            ('w2', 'w4', 'e'): 1.0,
-            ('w4', 'w8', 'i'): 0.9,
-            ('w4', 'w&infin;', 'r'): 0.1,
-            ('w8', 'w&infin;', 'q'): 1.0
-        }
-        actual_reduced_probabilities = nx.get_edge_attributes(
-            reduced,
-            'probability'
-        )
-
-        for edge, probability in expected_reduced_probabilities.items():
-            try:
-                actual = round(actual_reduced_probabilities[edge], 2)
-                assert actual == probability
-            except KeyError:
-                assert False  # edge doesnt exist in graph!
-
-        fname = Util.create_path('out/Evidence_test_certain_and_uncertain', False, 'pdf')
-        reduced.create_figure(fname)
-
-        expected_paths = [
-            [('w0', 'w1', 'a'), ('w1', 'w2', 'c'), ('w2', 'w4', 'e'),
-                ('w4', 'w&infin;', 'r')],
-            [('w0', 'w1', 'a'), ('w1', 'w2', 'c'), ('w2', 'w4', 'e'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'b'), ('w1', 'w2', 'c'), ('w2', 'w4', 'e'),
-                ('w4', 'w&infin;', 'r')],
-            [('w0', 'w1', 'b'), ('w1', 'w2', 'c'), ('w2', 'w4', 'e'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'a'), ('w1', 'w3', 'd'), ('w3', 'w8', 'h'),
-                ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'a'), ('w1', 'w3', 'd'), ('w3', 'w4', 'g'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'a'), ('w1', 'w3', 'd'), ('w3', 'w4', 'g'),
-                ('w4', 'w&infin;', 'r')],
-            [('w0', 'w1', 'b'), ('w1', 'w3', 'd'), ('w3', 'w8', 'h'),
-                ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'b'), ('w1', 'w3', 'd'), ('w3', 'w4', 'g'),
-                ('w4', 'w8', 'i'), ('w8', 'w&infin;', 'q')],
-            [('w0', 'w1', 'b'), ('w1', 'w3', 'd'), ('w3', 'w4', 'g'),
-                ('w4', 'w&infin;', 'r')],
-        ]
-        for path in expected_paths:
-            assert path in self.evidence.path_list
-        assert len(expected_paths) == len(self.evidence.path_list)
-
-
-class TestCegOutput:
-    def setup(self):
-        df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/medical_dm_modified.xlsx')
-
-        st = StagedTree(
-            dataframe=pd.read_excel(df_path),
-            name="medical_staged"
-        )
-        st.calculate_AHC_transitions()
-        self.ceg = ChainEventGraph(st)
-
-    def test_creation_of_ceg(self) -> None:
-        self.ceg.generate()
-        fname = Util.create_path('out/medical_dm_CEG', True, 'pdf')
-        self.ceg.create_figure(fname)
-        assert os.path.isfile(fname)
-
-    def test_generation(self) -> None:
-        self.ceg.generate()
-        path_list = self.ceg.path_list
-        nodes = ['w0', 'w1', 'w4', 'w10', 'w9', 'w&infin;']
-
-        subgraph = self.ceg.subgraph(nodes).copy()
-        fname = Util.create_path('out/Subgraph_test', True, 'pdf')
-        subgraph.create_figure(fname)
-        print(path_list)
-
-
-class TestCegNonStrat:
-    def setup(self):
-        df_path = Path(__file__).resolve().parent.parent.joinpath(
-            'data/Falls_Data.xlsx'
-        )
-        st = StagedTree(
-            dataframe=pd.read_excel(df_path),
-            name='Falls Staged Tree'
-        )
-        st.calculate_AHC_transitions()
-        self.ceg = ChainEventGraph(st)
-
-    def test_generation(self):
-        self.ceg.generate()
diff --git a/pytest.ini b/pytest.ini
index 871f185..a2b4132 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -1,2 +1,2 @@
-[pytest]
-addopts= -n4 --cov=./ --cov-report term-missing
\ No newline at end of file
+# [pytest]
+# addopts= -n4 --cov=./ --cov-report term-missing
\ No newline at end of file
diff --git a/scripts/build_cegpy_test.sh b/scripts/build_cegpy_test.sh
index c5d876a..e32f7df 100644
--- a/scripts/build_cegpy_test.sh
+++ b/scripts/build_cegpy_test.sh
@@ -1,8 +1,6 @@
 # Move old distribution to the old folder.
-mv dist/* .dist_old/
-
+mkdir --parents .dist_old/; mv dist/* $_
 # build the package
-python3 -m build
-
+python -m build
 # Upload new package
-python3 -m twine upload --repository testpypi dist/*
\ No newline at end of file
+python -m twine upload --repository testpypi dist/*
diff --git a/integration/__init__.py b/src/tests/__init__.py
similarity index 100%
rename from integration/__init__.py
rename to src/tests/__init__.py
diff --git a/src/tests/test_ceg.py b/src/tests/test_ceg.py
new file mode 100644
index 0000000..e3b93f4
--- /dev/null
+++ b/src/tests/test_ceg.py
@@ -0,0 +1,791 @@
+"""Tests ChainEventGraph"""
+# pylint: disable=protected-access
+import re
+from pathlib import Path
+from typing import Dict, List, Mapping, Tuple
+import unittest
+from unittest.mock import Mock, patch
+import networkx as nx
+import pandas as pd
+import pytest
+import pytest_mock
+from cegpy import StagedTree, ChainEventGraph
+from cegpy.graphs._ceg import (
+    CegAlreadyGenerated,
+    _merge_edge_data,
+)
+
+
+class TestMockedCEGMethods:
+    """Tests that Mock functions in ChainEventGraph"""
+
+    node_prefix = "w"
+    sink_suffix = "&infin;"
+    staged: StagedTree
+
+    def setup(self):
+        """Test setup"""
+        df_path = (
+            Path(__file__)
+            .resolve()
+            .parent.parent.joinpath("../data/medical_dm_modified.xlsx")
+        )
+
+        self.staged = StagedTree(dataframe=pd.read_excel(df_path))
+        self.staged.calculate_AHC_transitions()
+
+    def test_generate_argument(self, mocker: pytest_mock.MockerFixture):
+        """When ChainEventGraph called with generate, the .generate()
+        method is called."""
+        mocker.patch("cegpy.graphs._ceg.ChainEventGraph.generate")
+        ceg = ChainEventGraph(self.staged, generate=True)
+        ceg.generate.assert_called_once()  # pylint: disable=no-member
+
+
+class TestUnitCEG(unittest.TestCase):
+    """More ChainEventGraph tests"""
+
+    def setUp(self):
+        self.node_prefix = "w"
+        self.sink_suffix = "&infin;"
+        df_path = (
+            Path(__file__)
+            .resolve()
+            .parent.parent.joinpath("../data/medical_dm_modified.xlsx")
+        )
+
+        self.staged = StagedTree(dataframe=pd.read_excel(df_path))
+        self.staged.calculate_AHC_transitions()
+
+    def test_stages_property(self):
+        """Stages is a mapping of stage names to lists of nodes"""
+        ceg = ChainEventGraph(self.staged, generate=False)
+        node_stage_mapping: Mapping = dict(ceg.nodes(data="stage", default=None))
+        stages = ceg.stages
+        self.assertEqual(len(ceg.nodes), sum(len(nodes) for nodes in stages.values()))
+        for node, stage in node_stage_mapping.items():
+            self.assertIn(node, stages[stage])
+
+    def test_create_figure(self):
+        """.create_figure() called with no filename"""
+        ceg = ChainEventGraph(self.staged, generate=False)
+        with self.assertLogs("cegpy", level="INFO") as log_cm:
+            assert ceg.create_figure() is None
+        self.assertEqual(
+            ["WARNING:cegpy.chain_event_graph:No filename. Figure not saved."],
+            log_cm.output,
+        )
+
+
+class TestCEGHelpersTestCases(unittest.TestCase):
+    """Tests some of the CEG helper functions."""
+
+    def setUp(self):
+        self.graph = nx.MultiDiGraph()
+        self.init_nodes = ["w0", "w1", "w2", "w3", "w4", "w_infinity"]
+        self.init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w1", "w3", "c"),
+            ("w1", "w4", "d"),
+            ("w2", "w3", "c"),
+            ("w2", "w4", "d"),
+            ("w3", "w_infinity", "e"),
+            ("w4", "w_infinity", "f"),
+        ]
+        self.graph.root = "w0"
+        self.graph.add_nodes_from(self.init_nodes)
+        self.graph.add_edges_from(self.init_edges)
+        self.ceg = ChainEventGraph(self.graph, generate=False)
+
+    def test_merge_edge_data(self):
+        """Edges are merged"""
+        edge_1 = dict(
+            zip(
+                ["count", "prior", "posterior", "probability"],
+                [250, 0.5, 250, 0.8],
+            )
+        )
+        edge_2 = dict(
+            zip(
+                ["count", "prior", "posterior", "probability"],
+                [550, 25, 0.4, 0.9],
+            )
+        )
+
+        new_edge_dict = _merge_edge_data(edge_1=edge_1, edge_2=edge_2)
+        self.assert_edges_merged(new_edge_dict, edge_1, edge_2)
+
+    def test_merge_edge_data_with_missing(self):
+        """Edges are merged even when some data is missing in one edge."""
+        edge_1 = dict(
+            zip(
+                ["count", "prior"],
+                [250, 0.5],
+            )
+        )
+        edge_2 = dict(
+            zip(
+                ["count", "prior", "posterior"],
+                [550, 25, 0.4],
+            )
+        )
+
+        new_edge_dict = _merge_edge_data(edge_1=edge_1, edge_2=edge_2)
+        self.assert_edges_merged(new_edge_dict, edge_1, edge_2)
+
+    def assert_edges_merged(self, new_edge: Dict, edge_1: Dict, edge_2: Dict):
+        """Edges were merged successfully."""
+
+        self.assertSetEqual(
+            set(new_edge.keys()),
+            set(edge_1.keys()).union(set(edge_2.keys())),
+            msg="Edges do not have the same keys.",
+        )
+
+        for key, value in new_edge.items():
+            if key == "probability":
+                self.assertEqual(
+                    value, edge_1.get(key), msg="Probability shouldn't be summed."
+                )
+            else:
+                self.assertEqual(
+                    value,
+                    (edge_1.get(key, 0) + edge_2.get(key, 0)),
+                    msg=f"{key} not merged. Merged value: {value}",
+                )
+
+    def test_relabel_nodes(self):
+        """Relabel nodes successfully renames all the nodes."""
+        self.ceg._relabel_nodes()
+        node_pattern = r"^w([0-9]+)|w(_infinity)$"
+        prog = re.compile(node_pattern)
+        for node in self.ceg.nodes:
+            result = prog.match(node)
+            assert result is not None, "Node does not match expected format."
+
+    def test_merge_outgoing_edges(self):
+        """Outgoing edges are merged."""
+        edges_to_add = [
+            {
+                "src": "s1",
+                "dst": "s3",
+                "key": "event_1",
+                "counts": 5,
+                "priors": 0.2,
+                "posteriors": 1,
+                "probability": 0.8,
+            },
+            {
+                "src": "s1",
+                "dst": "s4",
+                "key": "event_2",
+                "counts": 10,
+                "priors": 0.3,
+                "posteriors": 5,
+                "probability": 0.2,
+            },
+            {
+                "src": "s2",
+                "dst": "s3",
+                "key": "event_1",
+                "counts": 11,
+                "priors": 0.5,
+                "posteriors": 2,
+                "probability": 0.8,
+            },
+            {
+                "src": "s2",
+                "dst": "s4",
+                "key": "event_2",
+                "counts": 6,
+                "priors": 0.9,
+                "posteriors": 3,
+                "probability": 0.2,
+            },
+        ]
+        for edge in edges_to_add:
+            self.ceg.add_edge(
+                u_for_edge=edge["src"],
+                v_for_edge=edge["dst"],
+                key=edge["key"],
+                counts=edge["counts"],
+                priors=edge["priors"],
+                posteriors=edge["posteriors"],
+            )
+        self.ceg.remove_edges_from(
+            [
+                ("s1", "s3", "c"),
+                ("s1", "s4", "d"),
+                ("s2", "s3", "c"),
+                ("s2", "s4", "d"),
+            ]
+        )
+
+        expected = [(edge["src"], edge["dst"], edge["key"]) for edge in edges_to_add]
+        actual = self.ceg._merge_and_add_edges("s99", "s1", "s2")
+        for edge in expected:
+            assert (
+                edge in actual
+            ), f"Expected {edge} in return value, but it was not found."
+        assert len(expected) == len(
+            actual
+        ), "Actual number of edges does not match expected number of edges."
+
+
+class TestNodesCanBeMerged(unittest.TestCase):
+    """Tests nodes_can_be_merged() function."""
+
+    def setUp(self):
+        self.graph = nx.MultiDiGraph()
+        self.init_nodes = ["w0", "w1", "w2", "w3", "w4", "w_infinity"]
+        self.graph.root = "w0"
+        self.graph.add_nodes_from(self.init_nodes)
+
+    def test_check_nodes_can_be_merged(self):
+        """Nodes can be merged."""
+        init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w1", "w3", "c"),
+            ("w1", "w4", "d"),
+            ("w2", "w3", "c"),
+            ("w2", "w4", "d"),
+            ("w3", "w_infinity", "e"),
+            ("w4", "w_infinity", "f"),
+        ]
+        self.graph.add_edges_from(init_edges)
+        ceg = ChainEventGraph(self.graph, generate=False)
+        ceg.nodes["w1"]["stage"] = 2
+        ceg.nodes["w2"]["stage"] = 2
+        assert ceg._check_nodes_can_be_merged("w1", "w2"), "Nodes should be mergeable."
+
+    def test_nodes_not_in_same_stage_cannot_be_merged(self):
+        """Nodes not in same stage cannot be merged"""
+        init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w1", "w3", "c"),
+            ("w1", "w4", "d"),
+            ("w2", "w3", "c"),
+            ("w2", "w4", "d"),
+            ("w3", "w_infinity", "e"),
+            ("w4", "w_infinity", "f"),
+        ]
+        self.graph.add_edges_from(init_edges)
+        ceg = ChainEventGraph(self.graph, generate=False)
+        ceg.nodes["w1"]["stage"] = 1
+        ceg.nodes["w2"]["stage"] = 2
+        assert not ceg._check_nodes_can_be_merged(
+            "w1", "w2"
+        ), "Nodes should not be mergeable."
+
+    def test_nodes_with_different_successor_nodes(self):
+        """Nodes with different successor nodes won't be merged."""
+        init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w1", "w3", "c"),
+            ("w1", "w4", "d"),
+            ("w2", "w_infinity", "c"),
+            ("w2", "w4", "d"),
+            ("w3", "w_infinity", "e"),
+            ("w4", "w_infinity", "f"),
+        ]
+        self.graph.add_edges_from(init_edges)
+        ceg = ChainEventGraph(self.graph, generate=False)
+        ceg.nodes["w1"]["stage"] = 2
+        ceg.nodes["w2"]["stage"] = 2
+        assert not ceg._check_nodes_can_be_merged(
+            "w1", "w2"
+        ), "Nodes should not be mergeable."
+
+    def test_nodes_with_different_outgoing_edges(self):
+        """Nodes with different outgoing edges won't be merged."""
+        init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w1", "w3", "c"),
+            ("w1", "w4", "d"),
+            ("w2", "w3", "g"),
+            ("w2", "w4", "d"),
+            ("w3", "w_infinity", "e"),
+            ("w4", "w_infinity", "f"),
+        ]
+        self.graph.add_edges_from(init_edges)
+        ceg = ChainEventGraph(self.graph, generate=False)
+        ceg.nodes["w1"]["stage"] = 2
+        ceg.nodes["w2"]["stage"] = 2
+        assert not ceg._check_nodes_can_be_merged(
+            "w1", "w2"
+        ), "Nodes should not be mergeable."
+
+    def test_merging_of_nodes(self):
+        """The nodes are merged, and all edges are merged too."""
+        init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w1", "w3", "c"),
+            ("w1", "w4", "d"),
+            ("w2", "w3", "c"),
+            ("w2", "w4", "d"),
+            ("w3", "w_infinity", "e"),
+            ("w4", "w_infinity", "f"),
+        ]
+        self.graph.add_edges_from(init_edges)
+        ceg = ChainEventGraph(self.graph, generate=False)
+        ceg.nodes["w1"]["stage"] = 2
+        ceg.nodes["w2"]["stage"] = 2
+        ceg._merge_nodes({("w1", "w2")})
+        expected_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w1", "b"),
+            ("w1", "w3", "c"),
+            ("w1", "w4", "d"),
+            ("w3", "w_infinity", "e"),
+            ("w4", "w_infinity", "f"),
+        ]
+        for edge in expected_edges:
+            self.assertIn(edge, list(ceg.edges))
+
+    def test_merging_of_three_nodes(self):
+        """The three nodes are merged, and all edges are merged too."""
+        self.graph.add_node("w5")
+        init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w0", "w3", "c"),
+            ("w1", "w4", "d"),
+            ("w1", "w5", "e"),
+            ("w2", "w4", "d"),
+            ("w2", "w5", "e"),
+            ("w3", "w4", "d"),
+            ("w3", "w5", "e"),
+            ("w4", "w_infinity", "f"),
+            ("w5", "w_infinity", "g"),
+        ]
+        self.graph.add_edges_from(init_edges)
+        ceg = ChainEventGraph(self.graph, generate=False)
+
+        nodes_to_merge = {"w1", "w2", "w3"}
+        ceg.nodes["w1"]["stage"] = 2
+        ceg.nodes["w2"]["stage"] = 2
+        ceg.nodes["w3"]["stage"] = 2
+        ceg._merge_nodes({("w1", "w2"), ("w2", "w3"), ("w1", "w3")})
+        nodes_post_merge = set(ceg.nodes)
+        merged_node = nodes_post_merge.intersection(nodes_to_merge).pop()
+        expected_edges = [
+            ("w0", merged_node, "a"),
+            ("w0", merged_node, "b"),
+            ("w0", merged_node, "c"),
+            (merged_node, "w4", "d"),
+            (merged_node, "w5", "e"),
+            ("w4", "w_infinity", "f"),
+            ("w5", "w_infinity", "g"),
+        ]
+
+        for edge in expected_edges:
+            self.assertIn(edge, list(ceg.edges))
+
+        self.assertEqual(len(list(ceg.edges)), len(expected_edges))
+
+
+class TestTrimLeavesFromGraph(unittest.TestCase):
+    """Tests trim_leaves_from_graph"""
+
+    def setUp(self):
+        self.graph = nx.MultiDiGraph()
+        self.init_nodes = ["w0", "w1", "w2", "w3", "w4"]
+        self.init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w1", "w3", "c"),
+            ("w1", "w4", "d"),
+            ("w2", "w3", "c"),
+            ("w2", "w4", "d"),
+        ]
+        self.graph.root = "w0"
+        self.leaves = ["w3", "w4"]
+        self.graph.add_nodes_from(self.init_nodes)
+        self.graph.add_edges_from(self.init_edges)
+        self.ceg = ChainEventGraph(self.graph, generate=False)
+
+    def test_leaves_trimmed_from_graph(self) -> None:
+        """Leaves are trimmed from the graph."""
+        self.ceg._trim_leaves_from_graph()
+        for leaf in self.leaves:
+            self.assertIsNone(self.ceg.nodes.get(leaf), "Leaf was not removed.")
+
+            for edge_list_key in self.ceg.edges.keys():
+                self.assertNotEqual(
+                    edge_list_key[1],
+                    leaf,
+                    msg=f"Edge still pointing to leaf: {leaf}",
+                )
+
+        expected_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w1", self.ceg.sink, "c"),
+            ("w1", self.ceg.sink, "d"),
+            ("w2", self.ceg.sink, "c"),
+            ("w2", self.ceg.sink, "d"),
+        ]
+        for edge in expected_edges:
+            self.assertIn(
+                edge,
+                list(self.ceg.edges),
+                msg=f"Edge not found: {edge}",
+            )
+
+        self.assertEqual(
+            len(list(self.ceg.edges)),
+            len(expected_edges),
+            "Wrong number of edges.",
+        )
+
+
+class TestPathList:
+    """Tests path list generation"""
+
+    # pylint: disable=too-few-public-methods
+    graph = nx.MultiDiGraph()
+
+    def test_path_list_generation(self):
+        """Path list is generated correctly."""
+
+        init_nodes = ["w0", "w1", "w2", "w3", "w4", "w5", "w_infinity"]
+        init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w1", "w3", "e"),
+            ("w1", "w4", "f"),
+            ("w2", "w_infinity", "c"),
+            ("w3", "w_infinity", "d"),
+            ("w4", "w5", "c"),
+            ("w5", "w_infinity", "d"),
+        ]
+        self.graph.root = "w0"
+        self.graph.add_nodes_from(init_nodes)
+        self.graph.add_edges_from(init_edges)
+        ceg = ChainEventGraph(self.graph, generate=False)
+        actual_path_list = ceg.path_list
+        expected_paths = [
+            [("w0", "w1", "a"), ("w1", "w3", "e"), ("w3", "w_infinity", "d")],
+            [
+                ("w0", "w1", "a"),
+                ("w1", "w4", "f"),
+                ("w4", "w5", "c"),
+                ("w5", "w_infinity", "d"),
+            ],
+            [("w0", "w2", "b"), ("w2", "w_infinity", "c")],
+        ]
+        for path in expected_paths:
+            assert path in actual_path_list, f"Path not found: {path}"
+
+        assert len(actual_path_list) == len(
+            expected_paths
+        ), "Incorrect number of paths."
+
+
+class TestDistanceToSink(unittest.TestCase):
+    """Tests distance to sink calculation"""
+
+    def setUp(self):
+        self.graph = nx.MultiDiGraph()
+        self.init_nodes = ["w0", "w1", "w2", "w3", "w4", "w5", "w_infinity"]
+        self.init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w1", "w3", "e"),
+            ("w1", "w4", "f"),
+            ("w2", "w_infinity", "c"),
+            ("w3", "w_infinity", "d"),
+            ("w4", "w5", "c"),
+            ("w5", "w_infinity", "d"),
+        ]
+        self.graph.add_nodes_from(self.init_nodes)
+        self.graph.add_edges_from(self.init_edges)
+        self.graph.root = "w0"
+        self.ceg = ChainEventGraph(self.graph, generate=False)
+
+    def test_update_distances_to_sink(self) -> None:
+        """Distance to sink is always max length of paths to sink."""
+
+        def check_distances():
+            actual_node_dists = nx.get_node_attributes(self.ceg, "max_dist_to_sink")
+            for node, distance in actual_node_dists.items():
+                assert distance == expected_node_dists[node]
+
+        expected_node_dists = {
+            "w0": 4,
+            "w1": 3,
+            "w2": 1,
+            "w3": 1,
+            "w4": 2,
+            "w5": 1,
+            "w_infinity": 0,
+        }
+        self.ceg._update_distances_to_sink()
+        check_distances()
+
+        # Add another edge to the dictionary, to show that the path is max,
+        # and not min distance to sink
+        self.ceg.add_edge("w1", self.ceg.sink)
+        self.ceg.add_edge("w4", self.ceg.sink)
+        self.ceg._update_distances_to_sink()
+        check_distances()
+
+    def test_gen_nodes_with_increasing_distance(self) -> None:
+        """Tests generate_nodes_with_increasing_distance"""
+        expected_nodes = {
+            0: [self.ceg.sink],
+            1: ["w2", "w3", "w5"],
+            2: ["w4"],
+            3: ["w1"],
+            4: [self.ceg.root],
+        }
+        for dist, nodes in expected_nodes.items():
+            for node in nodes:
+                self.ceg.nodes[node]["max_dist_to_sink"] = dist
+
+        nodes_gen = self.ceg._gen_nodes_with_increasing_distance(start=0)
+
+        for dist, nodes in expected_nodes.items():
+            actual_node_list = next(nodes_gen)
+            self.assertEqual(actual_node_list.sort(), nodes.sort())
+
+
+class TestEdgeInfoAttributes:
+    """Test edge_info argument."""
+
+    med_s_z_paths: List[Tuple]
+    med_df: pd.DataFrame
+    med_st: StagedTree
+
+    def setup(self):
+        """Test Setup"""
+        med_df_path = (
+            Path(__file__)
+            .resolve()
+            .parent.parent.joinpath("../data/medical_dm_modified.xlsx")
+        )
+        self.med_s_z_paths = None
+        self.med_df = pd.read_excel(med_df_path)
+        self.med_st = StagedTree(
+            dataframe=self.med_df, sampling_zero_paths=self.med_s_z_paths
+        )
+
+    def test_figure_with_wrong_edge_attribute(
+        self, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        """Ensures a warning is raised when a non-existent
+        attribute is passed for the edge_info argument"""
+        msg = (
+            r"edge_info 'prob' does not exist for the "
+            r"ChainEventGraph class. Using the default of 'probability' values "
+            r"on edges instead. For more information, see the "
+            r"documentation."
+        )
+
+        # stratified medical dataset
+        ceg = ChainEventGraph(self.med_st, generate=False)
+        _ = ceg.create_figure(filename=None, edge_info="prob")
+        assert msg in caplog.text, "Expected log message not logged."
+
+
+@patch.object(ChainEventGraph, "_relabel_nodes")
+@patch.object(ChainEventGraph, "_gen_nodes_with_increasing_distance")
+@patch.object(ChainEventGraph, "_backwards_construction")
+@patch.object(ChainEventGraph, "_update_distances_to_sink")
+@patch.object(ChainEventGraph, "_trim_leaves_from_graph")
+@patch.object(nx, "relabel_nodes")
+class TestGenerate(unittest.TestCase):
+    """Tests the .generate() method"""
+
+    def setUp(self) -> None:
+        self.graph = nx.MultiDiGraph()
+        self.init_nodes = [
+            "s0",
+            "s1",
+            "s2",
+            "s3",
+            "s4",
+            "s5",
+            "s6",
+            "s7",
+            "s8",
+            "s9",
+            "s10",
+            "s11",
+            "s12",
+        ]
+        self.init_edges = [
+            ("s0", "s1", "a"),
+            ("s0", "s2", "b"),
+            ("s1", "s3", "c"),
+            ("s1", "s4", "d"),
+            ("s2", "s5", "e"),
+            ("s2", "s6", "f"),
+            ("s3", "s7", "g"),
+            ("s3", "s8", "h"),
+            ("s4", "s9", "i"),
+            ("s4", "s10", "j"),
+            ("s4", "s11", "k"),
+            ("s0", "s12", "l"),
+        ]
+        self.graph.add_nodes_from(self.init_nodes)
+        self.graph.add_edges_from(self.init_edges)
+        self.graph.root = "s0"
+        self.graph.ahc_output = {
+            "Merged Situations": [("s1", "s2")],
+            "Loglikelihood": 1234.5678,
+        }
+        self.ceg = ChainEventGraph(self.graph, generate=False)
+
+    def test_raises_exception_when_called_twice(self, *_):
+        """.generate() raises a CegAlreadyGenerated error when called twice"""
+        self.ceg.generate()
+        with self.assertRaises(CegAlreadyGenerated):
+            self.ceg.generate()
+
+    def test_raises_exception_when_there_is_no_ahc_output(self, *_):
+        """.generate() raises a ValueError when ahc_output doesn't exist"""
+        self.ceg.ahc_output = None
+        with self.assertRaises(
+            ValueError, msg="There is no AHC output in your StagedTree."
+        ):
+            self.ceg.generate()
+
+    # pylint: disable=too-many-arguments
+    def test_calls_helper_functions_in_the_correct_order(
+        self,
+        nx_relabel: Mock,
+        trim_leaves: Mock,
+        update_distances: Mock,
+        backwards_construction: Mock,
+        gen_nodes: Mock,
+        relabel_nodes: Mock,
+    ):
+        """.generate() calls the helper functions"""
+        self.ceg.generate()
+        nx_relabel.assert_called_once_with(
+            self.ceg, {self.ceg.staged_root: self.ceg.root}, copy=False
+        )
+        trim_leaves.assert_called_once_with()
+        update_distances.assert_called_once_with()
+        backwards_construction.assert_called_once_with(gen_nodes.return_value)
+        gen_nodes.assert_called_once_with(start=1)
+        relabel_nodes.assert_called_once_with()
+
+    def test_generated_flag_set(self, *_):
+        """The generated flag is set to True when .generate() is called"""
+        self.ceg.generate()
+        assert self.ceg.generated is True
+
+
+class TestBackwardsConstruction(unittest.TestCase):
+    """Tests the ._backwards_construction() method"""
+
+    @staticmethod
+    def gen_sets_of_nodes():
+        """Generates a set of nodes for backwards construction"""
+        nodes = [[f"w{i}", f"w{i+1}"] for i in range(5, 1, -1)]
+        for node in nodes:
+            yield node
+
+    def test_backwards_construction_always_ends(self):
+        """
+        ._backwards_construction() always ends, even if there are no
+        nodes to process
+        """
+        graph = nx.MultiDiGraph()
+        init_nodes = [
+            "w0",
+            "w1",
+            "w2",
+            "w3",
+            "w4",
+            "w5",
+            "w6",
+            "w7",
+            "w8",
+            "w9",
+            "w10",
+            "w11",
+            "w12",
+        ]
+        init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w1", "w3", "c"),
+            ("w1", "w4", "d"),
+            ("w2", "w5", "e"),
+            ("w2", "w6", "f"),
+            ("w3", "w7", "g"),
+            ("w3", "w8", "h"),
+            ("w4", "w9", "i"),
+            ("w4", "w10", "j"),
+            ("w4", "w11", "k"),
+            ("w0", "w12", "l"),
+        ]
+        graph.add_nodes_from(init_nodes)
+        graph.add_edges_from(init_edges)
+        graph.root = "w0"
+        graph.ahc_output = {
+            "Merged Situations": [("s1", "s2")],
+            "Loglikelihood": 1234.5678,
+        }
+        ceg = ChainEventGraph(graph, generate=False)
+
+        ceg._backwards_construction(self.gen_sets_of_nodes())
+
+    def test_backwards_construction_produces_ceg(self):
+        """The backwards construction algorithm takes the staged tree and
+        makes the ceg."""
+        nodes = ["w0", "s1", "s2", "s3", "s4", "s5", "s11", "s12", "s13", "w_infinity"]
+        edges = [
+            ("w0", "s1", "hospital"),
+            ("w0", "s2", "community"),
+            ("s1", "s3", "test"),
+            ("s1", "s4", "no test"),
+            ("s2", "s11", "test"),
+            ("s2", "s12", "no test"),
+            ("s3", "s5", "positive"),
+            ("s3", "w_infinity", "negative"),
+            ("s4", "w_infinity", "death"),
+            ("s4", "w_infinity", "recovery"),
+            ("s5", "w_infinity", "death"),
+            ("s5", "w_infinity", "recovery"),
+            ("s11", "s13", "positive"),
+            ("s11", "w_infinity", "negative"),
+            ("s12", "w_infinity", "death"),
+            ("s12", "w_infinity", "recovery"),
+            ("s13", "w_infinity", "death"),
+            ("s13", "w_infinity", "recovery"),
+        ]
+        graph = nx.MultiDiGraph()
+        graph.add_nodes_from(nodes)
+        graph.add_edges_from(edges)
+        graph.root = "w0"
+        graph.nodes["w0"]["stage"] = 0
+        graph.nodes["s1"]["stage"] = 1
+        graph.nodes["s2"]["stage"] = 2
+        graph.nodes["s3"]["stage"] = 3
+        graph.nodes["s4"]["stage"] = 4
+        graph.nodes["s5"]["stage"] = 4
+        graph.nodes["s11"]["stage"] = 3
+        graph.nodes["s12"]["stage"] = 5
+        graph.nodes["s13"]["stage"] = 4
+
+        def gen_sets_of_nodes():
+            yield ["s4", "s5", "s12", "s13"]
+            yield ["s3", "s11"]
+            yield ["s1", "s2"]
+            yield ["w0"]
+
+        ceg = ChainEventGraph(graph, generate=False)
+        ceg._backwards_construction(gen_sets_of_nodes())
+        all_nodes = set(ceg.nodes)
+        self.assertEqual(len(all_nodes.intersection({"s5", "s4", "s13"})), 1)
+        self.assertEqual(len(all_nodes.intersection({"s11", "s3"})), 1)
diff --git a/src/tests/test_ceg_reducer.py b/src/tests/test_ceg_reducer.py
new file mode 100644
index 0000000..79fd4f6
--- /dev/null
+++ b/src/tests/test_ceg_reducer.py
@@ -0,0 +1,602 @@
+import unittest
+import networkx as nx
+import pytest
+from pydotplus.graphviz import InvocationException
+from cegpy import ChainEventGraph, ChainEventGraphReducer
+
+
+class TestCEGReducer(unittest.TestCase):
+    def setUp(self):
+        G = nx.MultiDiGraph()
+        self.init_nodes = [
+            "w0",
+            "w1",
+            "w2",
+            "w3",
+            "w4",
+            "w5",
+            "w6",
+            "w7",
+            "w8",
+            "w_infinity",
+        ]
+        self.init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w1", "b"),
+            ("w1", "w2", "c"),
+            ("w1", "w3", "d"),
+            ("w2", "w4", "e"),
+            ("w2", "w4", "f"),
+            ("w3", "w4", "g"),
+            ("w3", "w8", "h"),
+            ("w4", "w8", "i"),
+            ("w0", "w5", "j"),
+            ("w0", "w7", "k"),
+            ("w0", "w6", "l"),
+            ("w5", "w6", "m"),
+            ("w5", "w7", "n"),
+            ("w6", "w7", "o"),
+            ("w7", "w8", "p"),
+            ("w8", "w_infinity", "q"),
+            ("w4", "w_infinity", "r"),
+            ("w0", "w3", "s"),
+        ]
+        G.add_nodes_from(self.init_nodes)
+        G.add_edges_from(self.init_edges)
+        G.root = "w0"
+        G.ahc_output = {
+            "Merged Situations": [("s1", "s2")],
+            "Loglikelihood": 1234.5678,
+        }
+        H = ChainEventGraph(G, generate=False)
+        self.rceg = ChainEventGraphReducer(H)
+
+        self.rceg._ceg.edges["w0", "w1", "a"]["probability"] = 0.3
+        self.rceg._ceg.edges["w0", "w1", "b"]["probability"] = 0.3
+        self.rceg._ceg.edges["w1", "w2", "c"]["probability"] = 0.88
+        self.rceg._ceg.edges["w1", "w3", "d"]["probability"] = 0.12
+        self.rceg._ceg.edges["w2", "w4", "e"]["probability"] = 0.95
+        self.rceg._ceg.edges["w2", "w4", "f"]["probability"] = 0.05
+        self.rceg._ceg.edges["w3", "w4", "g"]["probability"] = 0.5
+        self.rceg._ceg.edges["w3", "w8", "h"]["probability"] = 0.5
+        self.rceg._ceg.edges["w4", "w8", "i"]["probability"] = 0.9
+        self.rceg._ceg.edges["w0", "w5", "j"]["probability"] = 0.2
+        self.rceg._ceg.edges["w0", "w7", "k"]["probability"] = 0.05
+        self.rceg._ceg.edges["w0", "w6", "l"]["probability"] = 0.05
+        self.rceg._ceg.edges["w5", "w6", "m"]["probability"] = 0.6
+        self.rceg._ceg.edges["w5", "w7", "n"]["probability"] = 0.4
+        self.rceg._ceg.edges["w6", "w7", "o"]["probability"] = 1
+        self.rceg._ceg.edges["w7", "w8", "p"]["probability"] = 1
+        self.rceg._ceg.edges["w8", "w_infinity", "q"]["probability"] = 1
+        self.rceg._ceg.edges["w4", "w_infinity", "r"]["probability"] = 0.1
+        self.rceg._ceg.edges["w0", "w3", "s"]["probability"] = 0.1
+
+        try:
+            self.rceg._ceg.create_figure("out/test_propation_pre.pdf")
+        except InvocationException:
+            pass
+
+    def test_repr(self):
+        rep = repr(self.rceg)
+
+        assert "certain_edges=" in rep
+        assert "certain_nodes=" in rep
+        assert "uncertain_edges=" in rep
+        assert "uncertain_nodes=" in rep
+
+    def test_update_path_list(self):
+        pre_path_list = self.rceg.paths
+
+        certain_edges = [
+            ("w0", "w1", "a"),
+            ("w1", "w2", "c"),
+        ]
+        for edge in certain_edges:
+            self.rceg.add_certain_edge(*edge)
+
+        self.rceg._update_path_list()
+        after_path_list = self.rceg.paths
+        assert pre_path_list != after_path_list
+
+    def test_add_and_remove_certain_edge(self):
+        certain_edges = [
+            ("w0", "w1", "a"),
+            ("w1", "w2", "c"),
+        ]
+        for edge in certain_edges:
+            self.rceg.add_certain_edge(*edge)
+        assert certain_edges == self.rceg.certain_edges
+
+        pytest.raises(ValueError, self.rceg.add_certain_edge, *("w3", "w100", "h"))
+
+        edge_to_remove = certain_edges[1]
+        self.rceg.remove_certain_edge(*edge_to_remove)
+        certain_edges.remove(edge_to_remove)
+        assert certain_edges == self.rceg.certain_edges
+
+        pytest.raises(
+            ValueError,
+            self.rceg.remove_certain_edge,
+            *edge_to_remove,
+        )
+
+    def test_add_and_remove_certain_edge_list(self):
+        certain_edges = [
+            ("w0", "w1", "a"),
+            ("w1", "w2", "c"),
+            ("w2", "w4", "e"),
+        ]
+        self.rceg.add_certain_edge_list(certain_edges)
+        assert certain_edges == self.rceg.certain_edges
+
+        self.rceg.remove_certain_edge_list(certain_edges[1:])
+        for edge in certain_edges[1:]:
+            certain_edges.remove(edge)
+        assert certain_edges == self.rceg.certain_edges
+
+    def test_add_and_remove_uncertain_edge_set(self):
+        uncertain_edge_sets = [
+            {
+                ("w0", "w1", "a"),
+                ("w0", "w1", "b"),
+                ("w0", "w3", "s"),
+            },
+            {
+                ("w0", "w5", "j"),
+                ("w0", "w7", "k"),
+            },
+        ]
+        for edge_set in uncertain_edge_sets:
+            self.rceg.add_uncertain_edge_set(edge_set)
+        assert uncertain_edge_sets == self.rceg.uncertain_edges
+
+        pytest.raises(
+            ValueError,
+            self.rceg.add_uncertain_edge_set,
+            {("w0", "w5", "j"), ("w0", "w7", "k"), ("w3", "w100", "h")},
+        )
+
+        edge_set_to_remove = uncertain_edge_sets[0].copy()
+        self.rceg.remove_uncertain_edge_set(edge_set_to_remove)
+        uncertain_edge_sets.remove(edge_set_to_remove)
+        assert uncertain_edge_sets == self.rceg.uncertain_edges
+
+        pytest.raises(
+            ValueError,
+            self.rceg.remove_uncertain_edge_set,
+            edge_set_to_remove,
+        )
+
+    def test_add_and_remove_uncertain_edge_set_list(self):
+        uncertain_edge_sets = [
+            {
+                ("w0", "w1", "a"),
+                ("w0", "w1", "b"),
+                ("w0", "w3", "s"),
+            },
+            {
+                ("w0", "w5", "j"),
+                ("w0", "w7", "k"),
+            },
+            {
+                ("w5", "w6", "m"),
+                ("w5", "w7", "n"),
+            },
+        ]
+        self.rceg.add_uncertain_edge_set_list(uncertain_edge_sets)
+        assert uncertain_edge_sets == self.rceg.uncertain_edges
+
+        self.rceg.remove_uncertain_edge_set_list(uncertain_edge_sets[1:])
+        for edge_set in uncertain_edge_sets[1:]:
+            uncertain_edge_sets.remove(edge_set)
+        assert uncertain_edge_sets == self.rceg.uncertain_edges
+
+    def test_add_and_remove_certain_nodes(self):
+        # nodes = ['w0', 'w1', 'w2', 'w3', 'w4',
+        #          'w5', 'w6', 'w7', 'w8', 'w_infinity']
+        certain_nodes = {"w0", "w1", "w3"}
+        for node in certain_nodes:
+            self.rceg.add_certain_node(node)
+        assert certain_nodes == self.rceg.certain_nodes
+
+        pytest.raises(
+            ValueError,
+            self.rceg.add_certain_node,
+            "w150",
+        )
+
+        node = certain_nodes.pop()
+        self.rceg.remove_certain_node(node)
+        assert certain_nodes == self.rceg.certain_nodes
+
+        pytest.raises(ValueError, self.rceg.remove_certain_node, node)
+
+    def test_add_and_remove_certain_nodes_set(self):
+        certain_nodes = {"w0", "w1", "w3"}
+        self.rceg.add_certain_node_set(certain_nodes)
+        assert certain_nodes == self.rceg.certain_nodes
+
+        nodes_to_remove = {"w0", "w1"}
+        self.rceg.remove_certain_node_set(nodes_to_remove)
+        certain_nodes.difference_update(nodes_to_remove)
+        assert certain_nodes == self.rceg.certain_nodes
+
+    def test_add_and_remove_uncertain_nodes_set(self):
+        uncertain_node_sets = [
+            {"w0", "w1", "w3"},
+            {"w4", "w5"},
+            {"w7", "w8"},
+        ]
+        for node_set in uncertain_node_sets:
+            self.rceg.add_uncertain_node_set(node_set)
+        assert uncertain_node_sets == self.rceg.uncertain_nodes
+
+        pytest.raises(
+            ValueError,
+            self.rceg.add_uncertain_node_set,
+            {"w6", "w150"},
+        )
+        node_set = uncertain_node_sets.pop()
+        self.rceg.remove_uncertain_node_set(node_set)
+        assert uncertain_node_sets == self.rceg.uncertain_nodes
+
+        pytest.raises(ValueError, self.rceg.remove_uncertain_node_set, node_set)
+
+    def test_add_and_remove_uncertain_nodes_set_list(self):
+        uncertain_node_sets = [
+            {"w0", "w1", "w3"},
+            {"w4", "w5"},
+            {"w7", "w8"},
+        ]
+        self.rceg.add_uncertain_node_set_list(uncertain_node_sets)
+        assert uncertain_node_sets == self.rceg.uncertain_nodes
+
+        node_sets_to_remove = uncertain_node_sets[1:]
+        self.rceg.remove_uncertain_node_set_list(node_sets_to_remove)
+        for node_set in node_sets_to_remove:
+            uncertain_node_sets.remove(node_set)
+        assert uncertain_node_sets == self.rceg.uncertain_nodes
+
+    def test_clear_all_evidence(self):
+        self.rceg.add_certain_edge("w0", "w1", "a")
+        self.rceg.add_certain_node("w1")
+        self.rceg.add_uncertain_edge_set(
+            {
+                ("w0", "w1", "a"),
+                ("w0", "w1", "b"),
+            }
+        )
+        self.rceg.add_uncertain_node_set({"w1", "w2"})
+
+        self.rceg.clear_all_evidence()
+        assert self.rceg.certain_edges == []
+        assert self.rceg.certain_nodes == set()
+        assert self.rceg.uncertain_edges == []
+        assert self.rceg.uncertain_nodes == []
+
+    def test_propagation(self) -> None:
+        uncertain_edges = {
+            ("w3", "w4", "g"),
+            ("w3", "w8", "h"),
+        }
+        certain_nodes = {"w1"}
+        self.rceg.add_uncertain_edge_set(uncertain_edges)
+        self.rceg.add_certain_node_set(certain_nodes)
+        rceg_out = self.rceg.graph
+
+        try:
+            rceg_out.create_figure("out/test_propagation_after.pdf")
+        except InvocationException:
+            pass
+
+    def test_propagation_two(self) -> None:
+        uncertain_edges = {
+            ("w2", "w4", "e"),
+            ("w3", "w4", "g"),
+            ("w3", "w8", "h"),
+        }
+        self.rceg.add_uncertain_edge_set(uncertain_edges)
+        self.rceg.add_certain_node("w1")
+        rceg_out = self.rceg.graph
+        assert rceg_out.edges["w0", "w1", "a"]["probability"] == 0.5
+        assert rceg_out.edges["w0", "w1", "b"]["probability"] == 0.5
+        assert rceg_out.edges["w1", "w2", "c"]["probability"] == (
+            pytest.approx(0.87, abs=0.01)
+        )
+        assert rceg_out.edges["w1", "w3", "d"]["probability"] == (
+            pytest.approx(0.13, abs=0.01)
+        )
+        assert rceg_out.edges["w2", "w4", "e"]["probability"] == 1.0
+        assert rceg_out.edges["w3", "w4", "g"]["probability"] == 0.5
+        assert rceg_out.edges["w3", "w8", "h"]["probability"] == 0.5
+        assert rceg_out.edges["w4", "w8", "i"]["probability"] == 0.9
+        assert rceg_out.edges["w8", "w_infinity", "q"]["probability"] == 1.0
+        assert rceg_out.edges["w4", "w_infinity", "r"]["probability"] == 0.1
+        try:
+            rceg_out.create_figure("out/test_propagation_two.pdf")
+        except InvocationException:
+            pass
+
+    def test_propagation_three(self) -> None:
+        uncertain_nodes = {"w3", "w6"}
+        self.rceg.add_uncertain_node_set(uncertain_nodes)
+        rceg_out = self.rceg.graph
+        try:
+            rceg_out.create_figure("out/test_propagation_three.pdf")
+        except InvocationException:
+            pass
+        assert rceg_out.edges["w0", "w1", "a"]["probability"] == (
+            pytest.approx(0.11, abs=0.01)
+        )
+        assert rceg_out.edges["w0", "w1", "b"]["probability"] == (
+            pytest.approx(0.11, abs=0.01)
+        )
+        assert rceg_out.edges["w1", "w3", "d"]["probability"] == (
+            pytest.approx(1.0, abs=0.01)
+        )
+        assert rceg_out.edges["w3", "w4", "g"]["probability"] == (
+            pytest.approx(0.5, abs=0.01)
+        )
+        assert rceg_out.edges["w3", "w8", "h"]["probability"] == (
+            pytest.approx(0.5, abs=0.01)
+        )
+        assert rceg_out.edges["w4", "w8", "i"]["probability"] == (
+            pytest.approx(0.9, abs=0.01)
+        )
+        assert rceg_out.edges["w0", "w5", "j"]["probability"] == (
+            pytest.approx(0.35, abs=0.01)
+        )
+        assert rceg_out.edges["w0", "w6", "l"]["probability"] == (
+            pytest.approx(0.15, abs=0.01)
+        )
+        assert rceg_out.edges["w5", "w6", "m"]["probability"] == (
+            pytest.approx(1.0, abs=0.01)
+        )
+        assert rceg_out.edges["w6", "w7", "o"]["probability"] == (
+            pytest.approx(1.0, abs=0.01)
+        )
+        assert rceg_out.edges["w7", "w8", "p"]["probability"] == (
+            pytest.approx(1.0, abs=0.01)
+        )
+        assert rceg_out.edges["w8", "w_infinity", "q"]["probability"] == (
+            pytest.approx(1.0, abs=0.01)
+        )
+        assert rceg_out.edges["w4", "w_infinity", "r"]["probability"] == (
+            pytest.approx(0.1, abs=0.01)
+        )
+        assert rceg_out.edges["w0", "w3", "s"]["probability"] == (
+            pytest.approx(0.29, abs=0.01)
+        )
+
+    def test_propagation_four(self) -> None:
+        uncertain_edges = [
+            {
+                ("w2", "w4", "e"),
+                ("w3", "w4", "g"),
+                ("w3", "w8", "h"),
+            },
+            {
+                ("w1", "w2", "c"),
+                ("w1", "w3", "d"),
+            },
+        ]
+        self.rceg.add_uncertain_edge_set_list(uncertain_edges)
+
+        rceg_out = self.rceg.graph
+        assert rceg_out.edges["w0", "w1", "a"]["probability"] == 0.5
+        assert rceg_out.edges["w0", "w1", "b"]["probability"] == 0.5
+        assert rceg_out.edges["w1", "w2", "c"]["probability"] == (
+            pytest.approx(0.87, abs=0.01)
+        )
+        assert rceg_out.edges["w1", "w3", "d"]["probability"] == (
+            pytest.approx(0.13, abs=0.01)
+        )
+        assert rceg_out.edges["w2", "w4", "e"]["probability"] == 1.0
+        assert rceg_out.edges["w3", "w4", "g"]["probability"] == 0.5
+        assert rceg_out.edges["w3", "w8", "h"]["probability"] == 0.5
+        assert rceg_out.edges["w4", "w8", "i"]["probability"] == 0.9
+        assert rceg_out.edges["w8", "w_infinity", "q"]["probability"] == 1.0
+        assert rceg_out.edges["w4", "w_infinity", "r"]["probability"] == 0.1
+        try:
+            rceg_out.create_figure("out/test_propagation_four.pdf")
+        except InvocationException:
+            pass
+
+    def test_str_out(self):
+        uncertain_edges = {
+            ("w7", "w8", "p"),
+            ("w3", "w8", "h"),
+        }
+        self.rceg.add_uncertain_edge_set(uncertain_edges)
+
+        str_rep = self.rceg.__str__()
+        for edge in uncertain_edges:
+            assert str(edge) in str_rep
+
+        certain_edges = [
+            ("w0", "w5", "j"),
+            ("w5", "w7", "n"),
+        ]
+        self.rceg.add_certain_edge_list(certain_edges)
+        str_rep = self.rceg.__str__()
+        for edge in certain_edges:
+            assert str(edge) in str_rep
+
+        uncertain_nodes = [
+            {"w5", "w6"},
+            {"w3", "w4"},
+        ]
+        self.rceg.add_uncertain_node_set_list(uncertain_nodes)
+        str_rep = self.rceg.__str__()
+        for node_set in uncertain_nodes:
+            assert str(node_set) in str_rep
+
+        certain_nodes = {"w1", "w2"}
+        self.rceg.add_certain_node_set(certain_nodes)
+
+        str_rep = self.rceg.__str__()
+        for edge in uncertain_edges:
+            assert str(edge) in str_rep
+
+        for edge in certain_edges:
+            assert str(edge) in str_rep
+
+        for node_set in uncertain_nodes:
+            assert str(node_set) in str_rep
+
+        for node in certain_nodes:
+            assert str(node) in str_rep
+
+
+class TestReducedCEGTwo(object):
+    def setup(self):
+        G = nx.MultiDiGraph()
+        self.init_nodes = ["w0", "w1", "w2", "w3", "w4", "w5", "w6", "w_infinity"]
+        self.init_edges = [
+            ("w0", "w1", "a"),
+            ("w0", "w2", "b"),
+            ("w0", "w3", "c"),
+            ("w1", "w4", "d"),
+            ("w1", "w4", "e"),
+            ("w2", "w4", "f"),
+            ("w2", "w5", "g"),
+            ("w2", "w6", "h"),
+            ("w3", "w6", "i"),
+            ("w4", "w_infinity", "j"),
+            ("w4", "w5", "k"),
+            ("w5", "w_infinity", "l"),
+            ("w6", "w_infinity", "m"),
+            ("w2", "w3", "n"),
+            ("w5", "w6", "o"),
+        ]
+        G.add_nodes_from(self.init_nodes)
+        G.add_edges_from(self.init_edges)
+        G.root = "w0"
+        G.ahc_output = {
+            "Merged Situations": [("s1", "s2")],
+            "Loglikelihood": 1234.5678,
+        }
+        H = ChainEventGraph(G, generate=False)
+        self.rceg = ChainEventGraphReducer(H)
+        self.rceg._ceg.edges["w0", "w1", "a"]["probability"] = 0.3
+        self.rceg._ceg.edges["w0", "w2", "b"]["probability"] = 0.4
+        self.rceg._ceg.edges["w0", "w3", "c"]["probability"] = 0.3
+        self.rceg._ceg.edges["w1", "w4", "d"]["probability"] = 0.35
+        self.rceg._ceg.edges["w1", "w4", "e"]["probability"] = 0.65
+        self.rceg._ceg.edges["w2", "w4", "f"]["probability"] = 0.25
+        self.rceg._ceg.edges["w2", "w5", "g"]["probability"] = 0.5
+        self.rceg._ceg.edges["w2", "w6", "h"]["probability"] = 0.15
+        self.rceg._ceg.edges["w3", "w6", "i"]["probability"] = 1.0
+        self.rceg._ceg.edges["w4", "w_infinity", "j"]["probability"] = 0.8
+        self.rceg._ceg.edges["w4", "w5", "k"]["probability"] = 0.2
+        self.rceg._ceg.edges["w5", "w_infinity", "l"]["probability"] = 0.7
+        self.rceg._ceg.edges["w6", "w_infinity", "m"]["probability"] = 1.0
+        self.rceg._ceg.edges["w2", "w3", "n"]["probability"] = 0.1
+        self.rceg._ceg.edges["w5", "w6", "o"]["probability"] = 0.3
+        try:
+            self.rceg._ceg.create_figure("out/test_propation_pre.pdf")
+        except InvocationException:
+            pass
+
+    def test_propagation_one(self):
+        uncertain_nodes = {"w4", "w5"}
+        self.rceg.add_uncertain_node_set(uncertain_nodes)
+
+        rceg_out = self.rceg.graph
+        try:
+            rceg_out.create_figure("out/prop_two_test_propagation_one.pdf")
+        except InvocationException:
+            pass
+        assert rceg_out.edges["w0", "w1", "a"]["probability"] == (
+            pytest.approx(0.46, abs=0.01)
+        )
+        assert rceg_out.edges["w0", "w2", "b"]["probability"] == (
+            pytest.approx(0.54, abs=0.01)
+        )
+        assert rceg_out.edges["w1", "w4", "d"]["probability"] == (
+            pytest.approx(0.35, abs=0.01)
+        )
+        assert rceg_out.edges["w1", "w4", "e"]["probability"] == (
+            pytest.approx(0.65, abs=0.01)
+        )
+        assert rceg_out.edges["w2", "w4", "f"]["probability"] == (
+            pytest.approx(0.29, abs=0.01)
+        )
+        assert rceg_out.edges["w2", "w5", "g"]["probability"] == (
+            pytest.approx(0.71, abs=0.01)
+        )
+        assert rceg_out.edges["w4", "w_infinity", "j"]["probability"] == (
+            pytest.approx(1.00, abs=0.01)
+        )
+        assert rceg_out.edges["w5", "w_infinity", "l"]["probability"] == (
+            pytest.approx(0.70, abs=0.01)
+        )
+        assert rceg_out.edges["w6", "w_infinity", "m"]["probability"] == (
+            pytest.approx(1.00, abs=0.01)
+        )
+        assert rceg_out.edges["w5", "w6", "o"]["probability"] == (
+            pytest.approx(0.30, abs=0.01)
+        )
+
+    def test_propagation_two(self):
+        uncertain_nodes = [
+            {"w1", "w2"},
+            {"w5", "w6"},
+        ]
+
+        self.rceg.add_uncertain_node_set_list(uncertain_nodes)
+        rceg_out = self.rceg.graph
+        try:
+            rceg_out.create_figure("out/prop_two_test_propagation_two.pdf")
+        except InvocationException:
+            pass
+
+        assert rceg_out.edges["w0", "w1", "a"]["probability"] == (
+            pytest.approx(0.14, abs=0.01)
+        )
+        assert rceg_out.edges["w0", "w2", "b"]["probability"] == (
+            pytest.approx(0.86, abs=0.01)
+        )
+        assert rceg_out.edges["w1", "w4", "d"]["probability"] == (
+            pytest.approx(0.35, abs=0.01)
+        )
+        assert rceg_out.edges["w1", "w4", "e"]["probability"] == (
+            pytest.approx(0.65, abs=0.01)
+        )
+        assert rceg_out.edges["w2", "w4", "f"]["probability"] == (
+            pytest.approx(0.06, abs=0.01)
+        )
+        assert rceg_out.edges["w2", "w5", "g"]["probability"] == (
+            pytest.approx(0.55, abs=0.01)
+        )
+        assert rceg_out.edges["w2", "w6", "h"]["probability"] == (
+            pytest.approx(0.24, abs=0.01)
+        )
+        assert rceg_out.edges["w3", "w6", "i"]["probability"] == (
+            pytest.approx(1.0, abs=0.01)
+        )
+        assert rceg_out.edges["w4", "w5", "k"]["probability"] == (
+            pytest.approx(1.0, abs=0.01)
+        )
+        assert rceg_out.edges["w5", "w_infinity", "l"]["probability"] == (
+            pytest.approx(1.0, abs=0.01)
+        )
+        assert rceg_out.edges["w6", "w_infinity", "m"]["probability"] == (
+            pytest.approx(1.0, abs=0.01)
+        )
+        assert rceg_out.edges["w2", "w3", "n"]["probability"] == (
+            pytest.approx(0.16, abs=0.01)
+        )
+
+    def test_propagation_three(self):
+        uncertain_edges = {
+            ("w4", "w5", "k"),
+            ("w5", "w_infinity", "l"),
+        }
+
+        self.rceg.add_uncertain_edge_set(uncertain_edges)
+        rceg_out = self.rceg.graph
+        try:
+            rceg_out.create_figure("out/prop_two_test_propagation_three.pdf")
+        except InvocationException:
+            pass
diff --git a/src/tests/test_event.py b/src/tests/test_event.py
new file mode 100644
index 0000000..717a352
--- /dev/null
+++ b/src/tests/test_event.py
@@ -0,0 +1,614 @@
+"""Tests cegpy.EventTree"""
+import unittest
+from collections import defaultdict
+import re
+from pathlib import Path
+import numpy as np
+import pytest
+import pandas as pd
+from pydotplus.graphviz import InvocationException
+from cegpy import EventTree
+
+
+class TestEventTreeAPI(unittest.TestCase):
+    """Tests API of EventTree"""
+
+    def setUp(self):
+        df_path = (
+            Path(__file__)
+            .resolve()
+            .parent.parent.joinpath("../data/medical_dm_modified.xlsx")
+        )
+        self.dataframe = pd.read_excel(df_path)
+
+    def test_required_argument_missing_fails(self):
+        """No arguments raises TypeError"""
+        self.assertRaises(TypeError, EventTree)
+
+    def test_required_argument_wrong_type_fails(self):
+        """Wrong argument type raises ValueError"""
+        dataframe = 5
+        self.assertRaises(ValueError, EventTree, dataframe=dataframe)
+
+    def test_incorrect_sampling_zero_fails(self):
+        """Incorrect Sampling Zero argument fails with ValueError"""
+        szp = [("edge_1"), ("edge_1", "edge_2")]
+        self.assertRaises(
+            ValueError, EventTree, dataframe=self.dataframe, sampling_zero_paths=szp
+        )
+
+
+class TestEventTree:
+    def setup(self):
+        df_path = (
+            Path(__file__)
+            .resolve()
+            .parent.parent.joinpath("../data/medical_dm_modified.xlsx")
+        )
+
+        self.df = pd.read_excel(df_path)
+        self.et = EventTree(dataframe=self.df)
+        self.reordered_et = EventTree(
+            dataframe=self.df, var_order=self.df.columns[::-1]
+        )
+        self.node_format = re.compile("^s\\d\\d*$")
+
+    def test_check_sampling_zero_paths_param(self) -> None:
+        """Tests the function that is checking the sampling zero paths param"""
+        szp = [("Medium",), ("Medium", "High")]
+        assert self.et._check_sampling_zero_paths_param(szp) == szp
+        szp = [1, 2, 3, 4]
+        assert self.et._check_sampling_zero_paths_param(szp) is None
+
+        szp = [("path", "to"), (123, "something"), "path/to"]
+        assert self.et._check_sampling_zero_paths_param(szp) is None
+
+    def test_check_sampling_zero_get_and_set(self) -> None:
+        """Tests the functions that set and get the sampling zeros"""
+        assert self.et.sampling_zeros is None
+
+        szp = [("Medium",), ("Medium", "High")]
+        self.et.sampling_zeros = szp
+        assert self.et.sampling_zeros == szp
+
+    def test_order_of_columns(self) -> None:
+        assert self.reordered_et.variables == list(self.df.columns[::-1])
+
+    def test_create_node_list_from_paths(self) -> None:
+        paths = defaultdict(int)
+        paths[("path",)] += 1
+        paths[("path", "to")] += 1
+        paths[("path", "away")] += 1
+        paths[("road",)] += 1
+        paths[("road", "to")] += 1
+        paths[("road", "away")] += 1
+
+        # code being tested:
+        node_list = self.et._create_node_list_from_paths(paths)
+
+        print(node_list)
+        assert len(list(paths.keys())) + 1 == len(node_list)
+        assert node_list[0] == "s0"
+        assert node_list[-1] == "s%d" % (len(node_list) - 1)
+
+    def test_construct_event_tree(self) -> None:
+        """Tests the construction of an event tree from a set of paths,
+        nodes, and"""
+        EXPECTED_NODE_COUNT = 45
+        assert len(self.et) == EXPECTED_NODE_COUNT
+        assert len(self.et.edges) == EXPECTED_NODE_COUNT - 1
+        edge_counts = self.et.edge_counts
+
+        assert len(edge_counts) == EXPECTED_NODE_COUNT - 1
+        for _, count in edge_counts.items():
+            assert isinstance(count, int)
+
+    def test_get_functions_producing_expected_data(self) -> None:
+        edges = list(self.et.edges)
+        assert isinstance(edges, list)
+        for edge in edges:
+            assert isinstance(edge, tuple)
+            assert len(edge) == 3
+            assert isinstance(edge[0], str)
+            assert isinstance(edge[1], str)
+            assert isinstance(edge[2], str)
+
+        check_list_contains_strings(list(self.et))
+        check_list_contains_strings(self.et.situations)
+        check_list_contains_strings(self.et.leaves)
+
+        edge_counts = self.et.edge_counts
+        print(edge_counts)
+        assert isinstance(edge_counts, dict)
+        for edge, count in edge_counts.items():
+            assert isinstance(edge, tuple)
+            for node in edge:
+                assert isinstance(node, str)
+            assert isinstance(count, int)
+
+    def test_dataframe_with_numeric_values(self) -> None:
+        """Ensures figure can be produced from dataframes with
+        numeric values"""
+        self.df["NewColumn"] = [1] * len(self.df)
+        new_et = EventTree(self.df)
+        try:
+            new_et.create_figure("out/test_dataframe_with_numeric_values.pdf")
+        except InvocationException:
+            pass
+        except Exception as err:
+            raise AssertionError("Could not create figure with numeric data") from err
+        return None
+
+    def test_figure_with_wrong_edge_attribute(
+        self,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        """Ensures a warning is raised when a non-existent
+        attribute is passed for the edge_info argument"""
+        _ = self.et.create_figure(filename=None, edge_info="probability")
+        msg = (
+            r"edge_info 'probability' does not exist for the "
+            r"EventTree class. Using the default of 'count' values "
+            r"on edges instead. For more information, see the "
+            r"documentation."
+        )
+        assert msg in caplog.text, "Expected log message not logged."
+
+    def test_node_colours(self) -> None:
+        """Ensures that all nodes in the event tree dot graph object
+        are coloured in lightgrey"""
+        dot_nodes = self.et.dot_event_graph().get_nodes()
+        event_node_colours = [n.obj_dict["attributes"]["fillcolor"] for n in dot_nodes]
+        assert len(set(event_node_colours)) == 1
+        assert event_node_colours[0] == "lightgrey"
+
+
+class TestIntegration:
+    def setup(self):
+        # stratified dataset
+        med_df_path = (
+            Path(__file__)
+            .resolve()
+            .parent.parent.joinpath("../data/medical_dm_modified.xlsx")
+        )
+        self.med_s_z_paths = None
+        self.med_df = pd.read_excel(med_df_path)
+        self.med_et = EventTree(
+            dataframe=self.med_df, sampling_zero_paths=self.med_s_z_paths
+        )
+
+        # non-stratified dataset
+        fall_df_path = (
+            Path(__file__).resolve().parent.parent.joinpath("../data/Falls_Data.xlsx")
+        )
+        self.fall_s_z_paths = None
+        self.fall_df = pd.read_excel(fall_df_path)
+        self.fall_et = EventTree(
+            dataframe=self.fall_df, sampling_zero_paths=self.fall_s_z_paths
+        )
+
+    def test_categories_per_variable(self) -> None:
+        expected_med_cats_per_var = {
+            "Classification": 2,
+            "Group": 3,
+            "Difficulty": 2,
+            "Response": 2,
+        }
+        actual_med_cats_per_var = self.med_et.categories_per_variable
+        assert expected_med_cats_per_var == actual_med_cats_per_var
+
+        expected_fall_cats_per_var = {
+            "HousingAssessment": 4,
+            "Risk": 2,
+            "Treatment": 3,
+            "Fall": 2,
+        }
+        actual_fall_cats_per_var = self.fall_et.categories_per_variable
+        assert expected_fall_cats_per_var == actual_fall_cats_per_var
+
+
+def check_list_contains_strings(str_list) -> bool:
+    assert isinstance(str_list, list)
+    for elem in str_list:
+        assert isinstance(elem, str)
+
+
+class TestUsecase:
+    def setup(self):
+        # stratified dataset
+        med_df_path = (
+            Path(__file__)
+            .resolve()
+            .parent.parent.joinpath("../data/medical_dm_modified.xlsx")
+        )
+        self.med_s_z_paths = None
+        self.med_df = pd.read_excel(med_df_path)
+        self.med_et = EventTree(
+            dataframe=self.med_df, sampling_zero_paths=self.med_s_z_paths
+        )
+
+        # non-stratified dataset
+        fall_df_path = (
+            Path(__file__).resolve().parent.parent.joinpath("../data/Falls_Data.xlsx")
+        )
+        self.fall_s_z_paths = None
+        self.fall_df = pd.read_excel(fall_df_path)
+        self.fall_et = EventTree(dataframe=self.fall_df)
+
+    def test_fall_cats_per_var(self):
+        expected_fall_cats_per_var = {
+            "HousingAssessment": 4,
+            "Risk": 2,
+            "Treatment": 3,
+            "Fall": 2,
+        }
+        actual_fall_cats_per_var = self.fall_et.categories_per_variable
+        assert expected_fall_cats_per_var == actual_fall_cats_per_var
+
+
+class TestChangingDataFrame:
+    def setup(self):
+        # stratified dataset
+        med_df_path = (
+            Path(__file__)
+            .resolve()
+            .parent.parent.joinpath("../data/medical_dm_modified.xlsx")
+        )
+        self.med_s_z_paths = None
+        self.med_df = pd.read_excel(med_df_path)
+        self.med_et = EventTree(
+            dataframe=self.med_df, sampling_zero_paths=self.med_s_z_paths
+        )
+
+        # non-stratified dataset
+        fall_df_path = (
+            Path(__file__).resolve().parent.parent.joinpath("../data/Falls_Data.xlsx")
+        )
+        self.fall_s_z_paths = None
+        self.fall_df = pd.read_excel(fall_df_path)
+        self.fall_et = EventTree(
+            dataframe=self.fall_df, sampling_zero_paths=self.fall_s_z_paths
+        )
+
+    def test_add_empty_column(self) -> None:
+        # adding empty column
+        med_empty_column_df = self.med_df
+        med_empty_column_df["extra"] = ""
+        med_empty_column_et = EventTree(dataframe=med_empty_column_df)
+        assert med_empty_column_et.adj == self.med_et.adj
+
+        fall_empty_column_df = self.fall_df
+        fall_empty_column_df["extra"] = ""
+        fall_empty_column_et = EventTree(dataframe=fall_empty_column_df)
+        assert fall_empty_column_et.adj == self.fall_et.adj
+
+    def test_add_NA_column(self) -> None:
+        # adding NA column
+        med_add_NA_df = self.med_df
+        med_add_NA_df["extra"] = np.nan
+        med_add_NA_et = EventTree(dataframe=med_add_NA_df)
+        assert med_add_NA_et.adj == self.med_et.adj
+
+        fall_add_NA_df = self.fall_df
+        fall_add_NA_df["extra"] = np.nan
+        fall_add_NA_et = EventTree(dataframe=fall_add_NA_df)
+        assert fall_add_NA_et.adj == self.fall_et.adj
+
+    def test_add_same_column(self) -> None:
+        # adding column with no more information
+        med_add_same_df = self.med_df
+        med_add_same_df["extra"] = "same for all"
+        med_add_same_et = EventTree(dataframe=med_add_same_df)
+        assert len(med_add_same_et.leaves) == len(self.med_et.leaves)
+
+        fall_add_same_df = self.fall_df
+        fall_add_same_df["extra"] = "same for all"
+        fall_add_same_et = EventTree(dataframe=fall_add_same_df)
+        assert len(fall_add_same_et.leaves) == len(self.fall_et.leaves)
+
+    def test_add_same_column_int(self) -> None:
+        # adding column with no more information
+        med_add_same_df = self.med_df
+        med_add_same_df["extra"] = 1
+        med_add_same_et = EventTree(dataframe=med_add_same_df)
+        try:
+            med_add_same_et.create_figure("et_fig_path.pdf")
+        except InvocationException:
+            pass
+        assert len(med_add_same_et.leaves) == len(self.med_et.leaves)
+
+        fall_add_same_df = self.fall_df
+        fall_add_same_df["extra"] = 1
+        fall_add_same_et = EventTree(dataframe=fall_add_same_df)
+        try:
+            fall_add_same_et.create_figure("et_fig_path.pdf")
+        except InvocationException:
+            pass
+        assert len(fall_add_same_et.leaves) == len(self.fall_et.leaves)
+
+
+class TestMissingLabels:
+    def setup(self):
+        array = [
+            np.array(["1", "NotANum", "Recover"]),
+            np.array(["1", "Trt1", "NotANum"]),
+            np.array(["2", "Struct", "Recover"]),
+            np.array(["2", "Struct", "Dont Recover"]),
+            np.array(["1", "Trt1", "Recover"]),
+            np.array(["1", "Trt2", "Recover"]),
+            np.array(["1", "Trt2", "Dont Recover"]),
+            np.array(["1", "Trt1", "Dont Recover"]),
+        ]
+
+        self.df = pd.DataFrame(array)
+
+        array_with_empty = [
+            np.array(["1", "NotANum", "Recover"]),
+            np.array(["1", "Trt1", "NotANum"]),
+            np.array(["2", "Struct", "Recover"]),
+            np.array(["2", "Struct", "Dont Recover"]),
+            np.array(["1", "Trt1", "Recover"]),
+            np.array(["1", "Trt2", "Recover"]),
+            np.array(["1", "Trt2", "Dont Recover"]),
+            np.array(["1", "Trt1", "Dont Recover"]),
+            np.array(["1", "", "Dont Recover"]),
+        ]
+
+        self.df_with_empty = pd.DataFrame(array_with_empty)
+
+        array_with_na = [
+            np.array(["1", "NotANum", "Recover"]),
+            np.array(["1", "Trt1", "NotANum"]),
+            np.array(["2", "Struct", "Recover"]),
+            np.array(["2", "Struct", "Dont Recover"]),
+            np.array(["1", "Trt1", "Recover"]),
+            np.array(["1", "Trt2", "Recover"]),
+            np.array(["1", "Trt2", "Dont Recover"]),
+            np.array(["1", "Trt1", "Dont Recover"]),
+            np.array(["1", "Trt1", np.nan]),
+        ]
+
+        self.df_with_na = pd.DataFrame(array_with_na)
+
+        array_with_both = [
+            np.array(["1", "NotANum", "Recover"]),
+            np.array(["1", "Trt1", "NotANum"]),
+            np.array(["2", "Struct", "Recover"]),
+            np.array(["2", "Struct", "Dont Recover"]),
+            np.array(["1", "Trt1", "Recover"]),
+            np.array(["1", "Trt2", "Recover"]),
+            np.array(["1", "Trt2", "Dont Recover"]),
+            np.array(["1", "Trt1", "Dont Recover"]),
+            np.array(["1", "Trt1", np.nan]),
+            np.array(["", "Trt1", "Dont Recover"]),
+        ]
+
+        self.df_with_both = pd.DataFrame(array_with_both)
+
+    def test_structural_label_string(self) -> None:
+        with pytest.raises(ValueError):
+            df_et = EventTree(dataframe=self.df, struct_missing_label=5)
+
+    def test_structural_label_reduction(self) -> None:
+        df_et = EventTree(dataframe=self.df, struct_missing_label="Struct")
+        expected_df = pd.DataFrame(
+            [
+                np.array(["1", "NotANum", "Recover"]),
+                np.array(["1", "Trt1", "NotANum"]),
+                np.array(["2", "", "Recover"]),
+                np.array(["2", "", "Dont Recover"]),
+                np.array(["1", "Trt1", "Recover"]),
+                np.array(["1", "Trt2", "Recover"]),
+                np.array(["1", "Trt2", "Dont Recover"]),
+                np.array(["1", "Trt1", "Dont Recover"]),
+            ]
+        )
+        assert df_et.dataframe.equals(expected_df) is True
+
+        df_with_empty = EventTree(dataframe=self.df_with_empty)
+
+    def test_missing_label_string(self) -> None:
+        with pytest.raises(ValueError):
+            df_et = EventTree(dataframe=self.df, missing_label=15)
+
+    def test_missing_label_reduction(self) -> None:
+        df_et = EventTree(dataframe=self.df, missing_label="NotANum")
+        expected_df = pd.DataFrame(
+            [
+                np.array(["1", "missing", "Recover"]),
+                np.array(["1", "Trt1", "missing"]),
+                np.array(["2", "Struct", "Recover"]),
+                np.array(["2", "Struct", "Dont Recover"]),
+                np.array(["1", "Trt1", "Recover"]),
+                np.array(["1", "Trt2", "Recover"]),
+                np.array(["1", "Trt2", "Dont Recover"]),
+                np.array(["1", "Trt1", "Dont Recover"]),
+            ]
+        )
+        assert df_et.dataframe.equals(expected_df) is True
+
+    def test_struct_missing_label_reduction(self) -> None:
+        df_et = EventTree(
+            dataframe=self.df, struct_missing_label="Struct", missing_label="NotANum"
+        )
+        expected_df = pd.DataFrame(
+            [
+                np.array(["1", "missing", "Recover"]),
+                np.array(["1", "Trt1", "missing"]),
+                np.array(["2", "", "Recover"]),
+                np.array(["2", "", "Dont Recover"]),
+                np.array(["1", "Trt1", "Recover"]),
+                np.array(["1", "Trt2", "Recover"]),
+                np.array(["1", "Trt2", "Dont Recover"]),
+                np.array(["1", "Trt1", "Dont Recover"]),
+            ]
+        )
+        assert df_et.dataframe.equals(expected_df) is True
+
+    def test_complete_case_bool(self) -> None:
+        with pytest.raises(ValueError):
+            df_et = EventTree(dataframe=self.df, complete_case="Yes")
+
+    def test_complete_case_missing_reduction(self) -> None:
+        """complete case missing_label reduces"""
+
+        # Case 1
+        df_et = EventTree(
+            dataframe=self.df, missing_label="NotANum", complete_case=True
+        )
+        expected_df = pd.DataFrame(
+            [
+                np.array(["2", "Struct", "Recover"]),
+                np.array(["2", "Struct", "Dont Recover"]),
+                np.array(["1", "Trt1", "Recover"]),
+                np.array(["1", "Trt2", "Recover"]),
+                np.array(["1", "Trt2", "Dont Recover"]),
+                np.array(["1", "Trt1", "Dont Recover"]),
+            ]
+        )
+        assert df_et.dataframe.equals(expected_df) is True
+
+        # Case 2
+        df_with_empty = EventTree(
+            dataframe=self.df_with_empty, missing_label="NotANum", complete_case=True
+        )
+        expected_df_with_empty = pd.DataFrame(
+            [
+                np.array(["2", "Struct", "Recover"]),
+                np.array(["2", "Struct", "Dont Recover"]),
+                np.array(["1", "Trt1", "Recover"]),
+                np.array(["1", "Trt2", "Recover"]),
+                np.array(["1", "Trt2", "Dont Recover"]),
+                np.array(["1", "Trt1", "Dont Recover"]),
+                np.array(["1", "", "Dont Recover"]),
+            ]
+        )
+        assert df_with_empty.dataframe.equals(expected_df_with_empty) is True
+
+        # Case 3
+        df_with_na = EventTree(
+            dataframe=self.df_with_na, missing_label="NotANum", complete_case=True
+        )
+        expected_df_with_na = pd.DataFrame(
+            [
+                np.array(["2", "Struct", "Recover"]),
+                np.array(["2", "Struct", "Dont Recover"]),
+                np.array(["1", "Trt1", "Recover"]),
+                np.array(["1", "Trt2", "Recover"]),
+                np.array(["1", "Trt2", "Dont Recover"]),
+                np.array(["1", "Trt1", "Dont Recover"]),
+                np.array(["1", "Trt1", np.nan]),
+            ]
+        )
+        assert df_with_na.dataframe.equals(expected_df_with_na) is True
+
+        # Case 4
+        df_with_both = EventTree(
+            dataframe=self.df_with_both, missing_label="NotANum", complete_case=True
+        )
+        expected_df_with_both = pd.DataFrame(
+            [
+                np.array(["2", "Struct", "Recover"]),
+                np.array(["2", "Struct", "Dont Recover"]),
+                np.array(["1", "Trt1", "Recover"]),
+                np.array(["1", "Trt2", "Recover"]),
+                np.array(["1", "Trt2", "Dont Recover"]),
+                np.array(["1", "Trt1", "Dont Recover"]),
+                np.array(["1", "Trt1", np.nan]),
+                np.array(["", "Trt1", "Dont Recover"]),
+            ]
+        )
+        assert df_with_both.dataframe.equals(expected_df_with_both) is True
+
+    def test_complete_case_reduction(self) -> None:
+        df_et = EventTree(
+            dataframe=self.df,
+            struct_missing_label="Struct",
+            missing_label="NotANum",
+            complete_case=True,
+        )
+        expected_df = pd.DataFrame(
+            [
+                np.array(["2", "", "Recover"]),
+                np.array(["2", "", "Dont Recover"]),
+                np.array(["1", "Trt1", "Recover"]),
+                np.array(["1", "Trt2", "Recover"]),
+                np.array(["1", "Trt2", "Dont Recover"]),
+                np.array(["1", "Trt1", "Dont Recover"]),
+            ]
+        )
+        assert df_et.dataframe.equals(expected_df) is True
+
+
+class TestVariablesFiltered:
+    def setup(self):
+        array = [
+            np.array(["1", "NotANum", "Recover"]),
+            np.array(["1", "Trt1", "NotANum"]),
+            np.array(["2", "Struct", "Recover"]),
+            np.array(["2", "Struct", "Dont Recover"]),
+            np.array(["1", "Trt1", "Recover"]),
+            np.array(["1", "Trt2", "Recover"]),
+            np.array(["1", "Trt2", "Dont Recover"]),
+            np.array(["1", np.NaN, "Dont Recover"]),
+        ]
+
+        self.df = pd.DataFrame(array)
+
+    def test_pd_nans_filtered(self) -> None:
+        """Checking that the pd.nan filter works"""
+        df_et = EventTree(dataframe=self.df)
+
+        expected_categories = {0: 2, 1: 4, 2: 3}
+
+        assert df_et.categories_per_variable == expected_categories
+
+    def test_pd_nans_filtered_with_missing(self) -> None:
+        """Checking that the pd.nan filter works"""
+        df_et = EventTree(
+            dataframe=self.df,
+            struct_missing_label="Struct",
+            missing_label="NotANum",
+            complete_case=True,
+        )
+
+        expected_categories = {0: 2, 1: 3, 2: 2}
+
+        assert df_et.categories_per_variable == expected_categories
+
+
+class TestStageColours:
+    def setup(self):
+        array = [
+            np.array(["1", "NotANum", "Recover"]),
+            np.array(["1", "Trt1", "NotANum"]),
+            np.array(["2", "Struct", "Recover"]),
+            np.array(["2", "Struct", "Dont Recover"]),
+            np.array(["1", "Trt1", "Recover"]),
+            np.array(["1", "Trt2", "Recover"]),
+            np.array(["1", "Trt2", "Dont Recover"]),
+            np.array(["1", np.NaN, "Dont Recover"]),
+        ]
+
+        self.df = pd.DataFrame(array)
+
+    def test_pd_nans_filtered(self) -> None:
+        """Checking that the pd.nan filter works"""
+        df_et = EventTree(dataframe=self.df)
+
+        expected_categories = {0: 2, 1: 4, 2: 3}
+
+        assert df_et.categories_per_variable == expected_categories
+
+    def test_pd_nans_filtered_with_missing(self) -> None:
+        """Checking that the pd.nan filter works"""
+        df_et = EventTree(
+            dataframe=self.df,
+            struct_missing_label="Struct",
+            missing_label="NotANum",
+            complete_case=True,
+        )
+
+        expected_categories = {0: 2, 1: 3, 2: 2}
+
+        assert df_et.categories_per_variable == expected_categories
diff --git a/src/tests/test_staged.py b/src/tests/test_staged.py
new file mode 100644
index 0000000..31d4135
--- /dev/null
+++ b/src/tests/test_staged.py
@@ -0,0 +1,1105 @@
+"""Tests StagedTree class."""
+# pylint: disable=too-many-lines,protected-access
+from pathlib import Path
+import unittest
+from fractions import Fraction
+import pytest
+import pandas as pd
+from pydotplus.graphviz import InvocationException
+import numpy as np
+
+# import xlsxwriter
+from cegpy import StagedTree
+from cegpy.trees._staged import _calculate_mean_posterior_probs
+
+
+class TestLogging:
+    """Tests logging in stagedtree"""
+
+    def setup(self):
+        """Test setup"""
+        # pylint: disable=attribute-defined-outside-init
+        # stratified dataset
+        med_df_path = (
+            Path(__file__)
+            .resolve()
+            .parent.parent.joinpath("../data/medical_dm_modified.xlsx")
+        )
+        self.med_s_z_paths = None
+        self.med_df = pd.read_excel(med_df_path)
+        self.med_st = StagedTree(
+            dataframe=self.med_df, sampling_zero_paths=self.med_s_z_paths
+        )
+
+        # non-stratified dataset
+        fall_df_path = (
+            Path(__file__).resolve().parent.parent.joinpath("../data/Falls_Data.xlsx")
+        )
+        self.fall_s_z_paths = None
+        self.fall_df = pd.read_excel(fall_df_path)
+        self.fall_st = StagedTree(
+            dataframe=self.fall_df,
+            sampling_zero_paths=self.fall_s_z_paths,
+        )
+
+    def test_figure_with_wrong_edge_attribute(
+        self, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        """Ensures a warning is raised when a non-existent
+        attribute is passed for the edge_info argument"""
+        msg = (
+            r"edge_info 'prob' does not exist for the "
+            r"StagedTree class. Using the default of 'count' values "
+            r"on edges instead. For more information, see the "
+            r"documentation."
+        )
+
+        # stratified medical dataset
+        self.med_st.calculate_AHC_transitions()
+        _ = self.med_st.create_figure(filename=None, edge_info="prob")
+        assert msg in caplog.text, "Expected log message not logged."
+
+        # non-stratified dataset
+        self.fall_st.calculate_AHC_transitions()
+        _ = self.fall_st.create_figure(filename=None, edge_info="prob")
+        assert msg in caplog.text, "Expected log message not logged."
+
+    def test_run_ahc_before_figure(self, caplog) -> None:
+        """Tests expected error message is in the log when running without
+        running AHC"""
+        try:
+            self.med_st.create_figure()
+            assert "PLEASE RUN AHC" in caplog.text
+        except InvocationException:
+            pass
+
+
+class TestStagedTrees(unittest.TestCase):
+    """Tests for staged trees with a stratified and non stratified dataset"""
+
+    def setUp(self):
+        # stratified dataset
+        med_df_path = (
+            Path(__file__)
+            .resolve()
+            .parent.parent.joinpath("../data/medical_dm_modified.xlsx")
+        )
+        self.med_s_z_paths = None
+        self.med_df = pd.read_excel(med_df_path)
+        self.med_st = StagedTree(
+            dataframe=self.med_df, sampling_zero_paths=self.med_s_z_paths
+        )
+
+        # non-stratified dataset
+        fall_df_path = (
+            Path(__file__).resolve().parent.parent.joinpath("../data/Falls_Data.xlsx")
+        )
+        self.fall_s_z_paths = None
+        self.fall_df = pd.read_excel(fall_df_path)
+        self.fall_st = StagedTree(
+            dataframe=self.fall_df,
+            sampling_zero_paths=self.fall_s_z_paths,
+        )
+
+    def test_check_hyperstage(self) -> None:
+        # stratified medical dataset
+        med_hyperstage_less = [
+            [
+                "s0",
+                "s9",
+                "s10",
+                "s11",
+                "s12",
+                "s13",
+                "s14",
+                "s15",
+                "s16",
+                "s17",
+                "s19",
+                "s20",
+            ],
+            ["s1", "s2"],
+            ["s3", "s4", "s6", "s7", "s8"],
+        ]
+
+        med_hyperstage_more = [
+            [
+                "s0",
+                "s9",
+                "s10",
+                "s11",
+                "s12",
+                "s13",
+                "s14",
+                "s15",
+                "s16",
+                "s17",
+                "s18",
+                "s19",
+                "s20",
+            ],
+            ["s1", "s2", "s50"],
+            ["s3", "s4", "s5", "s6", "s7", "s8"],
+        ]
+
+        med_hyperstage_unequal = [
+            [
+                "s0",
+                "s1",
+                "s9",
+                "s10",
+                "s11",
+                "s12",
+                "s13",
+                "s14",
+                "s15",
+                "s16",
+                "s17",
+                "s18",
+                "s19",
+                "s20",
+            ],
+            ["s2", "s4", "s50"],
+            ["s3", "s5", "s6", "s7", "s8"],
+        ]
+
+        with pytest.raises(ValueError):
+            self.med_st._check_hyperstages(med_hyperstage_less)
+
+        with pytest.raises(ValueError):
+            self.med_st._check_hyperstages(med_hyperstage_more)
+
+        with pytest.raises(ValueError):
+            self.med_st._check_hyperstages(med_hyperstage_unequal)
+
+        # non-stratified falls dataset
+        fall_hyperstage_less = [
+            ["s1", "s2", "s3", "s4"],
+            ["s5", "s9"],
+            ["s6", "s10"],
+            [
+                "s7",
+                "s8",
+                "s11",
+                "s12",
+                "s13",
+                "s14",
+                "s15",
+                "s16",
+                "s17",
+                "s22",
+                "s23",
+                "s24",
+                "s25",
+                "s26",
+            ],
+        ]
+
+        fall_hyperstage_more = [
+            ["s0", "v10"],
+            ["s1", "s2", "s3", "s4"],
+            ["s5", "s9", "pg"],
+            ["s6", "s10"],
+            [
+                "s7",
+                "s8",
+                "s11",
+                "s12",
+                "s13",
+                "s14",
+                "s15",
+                "s16",
+                "s17",
+                "s22",
+                "s23",
+                "s24",
+                "s25",
+                "s26",
+            ],
+        ]
+
+        fall_hyperstage_unequal = [
+            ["s0", "s2"],
+            ["s1", "s3", "s4"],
+            ["s5", "s9", "s8"],
+            ["s6", "s10"],
+            [
+                "s7",
+                "s11",
+                "s12",
+                "s13",
+                "s14",
+                "s15",
+                "s16",
+                "s17",
+                "s22",
+                "s23",
+                "s24",
+                "s25",
+                "s26",
+            ],
+        ]
+
+        with pytest.raises(ValueError):
+            self.fall_st._check_hyperstages(fall_hyperstage_less)
+
+        with pytest.raises(ValueError):
+            self.fall_st._check_hyperstages(fall_hyperstage_more)
+
+        with pytest.raises(ValueError):
+            self.fall_st._check_hyperstages(fall_hyperstage_unequal)
+
+    def test_check_prior(self) -> None:
+        # stratified medical dataset
+        med_prior_incorrect_length = [
+            [Fraction(3, 2), Fraction(3, 2)],
+            [Fraction(1, 2), Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+        ]
+
+        med_prior_unequal = [
+            [Fraction(3, 2), Fraction(3, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 4), Fraction(1, 4), Fraction(1, 2)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+        ]
+
+        med_prior_negative = [
+            [Fraction(-3, 2), Fraction(3, 2)],
+            [Fraction(1, 2), Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 4), Fraction(-1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(-1, 8), Fraction(-1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+        ]
+
+        with pytest.raises(ValueError):
+            self.med_st._check_prior(med_prior_incorrect_length)
+
+        with pytest.raises(ValueError):
+            self.med_st._check_prior(med_prior_unequal)
+
+        with pytest.raises(ValueError):
+            self.med_st._check_prior(med_prior_negative)
+
+        # non-stratified falls dataset
+        fall_prior_incorrect_length = [
+            [Fraction(1, 1), Fraction(1, 1), Fraction(1, 1), Fraction(1, 1)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 6), Fraction(1, 6), Fraction(1, 6)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+        ]
+
+        fall_prior_unequal = [
+            [Fraction(1, 1), Fraction(1, 1)],
+            [Fraction(1, 2), Fraction(1, 1), Fraction(1, 1), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 6), Fraction(1, 6), Fraction(1, 6)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 6), Fraction(1, 6), Fraction(1, 6)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8), Fraction(1, 8)],
+        ]
+
+        fall_prior_negative = [
+            [Fraction(1, 1), Fraction(1, 1), Fraction(1, 1), Fraction(1, 1)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 6), Fraction(-1, 6), Fraction(1, 6)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(-1, 4)],
+            [Fraction(1, 6), Fraction(1, 6), Fraction(1, 6)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 12), Fraction(-1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+        ]
+
+        with pytest.raises(ValueError):
+            self.fall_st._check_prior(fall_prior_incorrect_length)
+
+        with pytest.raises(ValueError):
+            self.fall_st._check_prior(fall_prior_unequal)
+
+        with pytest.raises(ValueError):
+            self.fall_st._check_prior(fall_prior_negative)
+
+    def test_check_prior_assigned(self) -> None:
+        # stratified medical dataset
+        med_prior_noerror = [
+            [Fraction(34, 2), Fraction(3, 2)],
+            [Fraction(1, 2), Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(51, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(11, 42)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(21, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+        ]
+        self.med_st.calculate_AHC_transitions(prior=med_prior_noerror)
+        assert self.med_st.prior_list == med_prior_noerror
+
+        # non-stratified falls dataset
+        fall_prior_noerror = [
+            [Fraction(1, 1), Fraction(1, 1), Fraction(1, 1), Fraction(1, 1)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(21, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 6), Fraction(1, 6), Fraction(1, 6)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(31, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 6), Fraction(15, 6), Fraction(1, 6)],
+            [Fraction(12, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+        ]
+
+        self.fall_st.calculate_AHC_transitions(prior=fall_prior_noerror)
+        assert self.fall_st.prior_list == fall_prior_noerror
+
+    def test_check_hyperstage_assigned(self) -> None:
+        # stratified medical dataset
+        med_hyperstage_noerror = [
+            ["s0", "s9", "s10", "s13", "s14", "s15", "s16", "s17", "s18", "s19", "s20"],
+            ["s1", "s2"],
+            ["s3", "s4", "s5", "s11", "s12", "s6", "s7", "s8"],
+        ]
+
+        self.med_st.calculate_AHC_transitions(hyperstage=med_hyperstage_noerror)
+        assert self.med_st.hyperstage == med_hyperstage_noerror
+
+        # non-stratified falls dataset
+        fall_hyperstage_noerror = [
+            ["s0"],
+            ["s2", "s3", "s4"],
+            ["s5", "s9"],
+            ["s1", "s6", "s10"],
+            [
+                "s7",
+                "s8",
+                "s11",
+                "s12",
+                "s13",
+                "s14",
+                "s15",
+                "s16",
+                "s17",
+                "s22",
+                "s23",
+                "s24",
+                "s25",
+                "s26",
+            ],
+        ]
+
+        self.fall_st.calculate_AHC_transitions(hyperstage=fall_hyperstage_noerror)
+        assert self.fall_st.hyperstage == fall_hyperstage_noerror
+
+    def test_hyperstage_noerror(self) -> None:
+        # stratified medical dataset
+        med_hyperstage_noerror = [
+            ["s0", "s9", "s10", "s13", "s14", "s15", "s16", "s17", "s18", "s19", "s20"],
+            ["s1", "s2"],
+            ["s3", "s4", "s5", "s11", "s12", "s6", "s7", "s8"],
+        ]
+
+        try:
+            self.med_st._check_hyperstages(med_hyperstage_noerror)
+        except ValueError as exc:
+            self.fail(f"'check_hyperstages' raised an exception {exc}")
+
+        # non-stratified falls dataset
+        fall_hyperstage_noerror = [
+            ["s0"],
+            ["s2", "s3", "s4"],
+            ["s5", "s9"],
+            ["s1", "s6", "s10"],
+            [
+                "s7",
+                "s8",
+                "s11",
+                "s12",
+                "s13",
+                "s14",
+                "s15",
+                "s16",
+                "s17",
+                "s22",
+                "s23",
+                "s24",
+                "s25",
+                "s26",
+            ],
+        ]
+
+        try:
+            self.fall_st._check_hyperstages(fall_hyperstage_noerror)
+        except ValueError as exc:
+            self.fail(f"'check_hyperstages' raised an exception {exc}")
+
+    def test_create_default_prior(self) -> None:
+        # stratified medical dataset
+        # Expected prior calculated by hand for default alpha of 3
+        med_expected_prior = [
+            [Fraction(3, 2), Fraction(3, 2)],
+            [Fraction(1, 2), Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+        ]
+
+        fall_expected_prior = [
+            [Fraction(1, 1), Fraction(1, 1), Fraction(1, 1), Fraction(1, 1)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 6), Fraction(1, 6), Fraction(1, 6)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 6), Fraction(1, 6), Fraction(1, 6)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 12), Fraction(1, 12)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+        ]
+        assert len(med_expected_prior) == 21
+        assert len(fall_expected_prior) == 23
+
+        alpha = 3
+        # check that prior is correct.
+        prior = self.med_st._create_default_prior(alpha)
+        assert med_expected_prior == prior
+
+        alpha = 4
+        prior = self.fall_st._create_default_prior(alpha)
+        assert fall_expected_prior == prior
+
+    def test_create_default_hyperstage(self) -> None:
+        med_expected_hyperstage = [
+            [
+                "s0",
+                "s9",
+                "s10",
+                "s11",
+                "s12",
+                "s13",
+                "s14",
+                "s15",
+                "s16",
+                "s17",
+                "s18",
+                "s19",
+                "s20",
+            ],
+            ["s1", "s2"],
+            ["s3", "s4", "s5", "s6", "s7", "s8"],
+        ]
+        fall_expected_hyperstage = [
+            ["s0"],
+            ["s1", "s2", "s3", "s4"],
+            ["s5", "s9"],
+            ["s6", "s10"],
+            [
+                "s7",
+                "s8",
+                "s11",
+                "s12",
+                "s13",
+                "s14",
+                "s15",
+                "s16",
+                "s17",
+                "s22",
+                "s23",
+                "s24",
+                "s25",
+                "s26",
+            ],
+        ]
+        med_hyperstage = self.med_st._create_default_hyperstage()
+
+        fall_hyperstage = self.fall_st._create_default_hyperstage()
+        assert med_hyperstage == med_expected_hyperstage
+        assert fall_hyperstage == fall_expected_hyperstage
+
+    def test_calculate_posterior(self) -> None:
+        def calculate_posterior(
+            staged_tree: StagedTree, expected_countset, alpha, expected_likelihood
+        ):
+            actual_countset = staged_tree._create_edge_countset()
+            assert actual_countset == expected_countset
+
+            prior = staged_tree._create_default_prior(alpha)
+            staged_tree.prior = prior
+            expected_posterior = []
+            for idx, countset in enumerate(actual_countset):
+                p_elem = []
+                for jdx, count in enumerate(countset):
+                    p_elem.append(count + prior[idx][jdx])
+
+                expected_posterior.append(p_elem)
+
+            actual_posterior = staged_tree.posterior_list
+            assert actual_posterior == expected_posterior
+            actual_likelihood = staged_tree._calculate_initial_loglikelihood(
+                staged_tree.prior_list,
+                staged_tree.posterior_list,
+            )
+            actual_likelihood = round(actual_likelihood, 2)
+            assert actual_likelihood == expected_likelihood
+
+        falls_expected_edge_countset = [
+            [379, 1539, 2871, 45211],
+            [235, 144],
+            [436, 1103],
+            [1449, 1422],
+            [5375, 39836],
+            [53, 52, 130],
+            [130, 14],
+            [82, 354],
+            [870, 233],
+            [319, 315, 815],
+            [1266, 156],
+            [1103, 4272],
+            [30769, 9067],
+            [11, 42],
+            [28, 24],
+            [58, 72],
+            [102, 28],
+            [12, 2],
+            [65, 254],
+            [151, 164],
+            [379, 436],
+            [974, 292],
+            [133, 23],
+        ]
+        alpha = 4
+        expected_likelihood = -68721.50  # Calculated manually
+        assert len(falls_expected_edge_countset) == 23
+        calculate_posterior(
+            self.fall_st, falls_expected_edge_countset, alpha, expected_likelihood
+        )
+
+        med_expected_edge_countset = [
+            [5491, 5493],
+            [798, 1000, 3693],
+            [799, 998, 3696],
+            [399, 399],
+            [500, 500],
+            [1846, 1847],
+            [400, 399],
+            [500, 498],
+            [1849, 1847],
+            [393, 6],
+            [347, 52],
+            [480, 20],
+            [433, 67],
+            [1482, 364],
+            [1305, 542],
+            [5, 395],
+            [92, 307],
+            [70, 430],
+            [202, 296],
+            [338, 1511],
+            [716, 1131],
+        ]
+        alpha = float(3)
+        expected_likelihood = -30134.07
+        assert len(med_expected_edge_countset) == 21
+        calculate_posterior(
+            self.med_st, med_expected_edge_countset, alpha, expected_likelihood
+        )
+
+    def test_prior_alpha_conflict(self) -> None:
+        prior = [
+            [Fraction(3, 2), Fraction(3, 2)],
+            [Fraction(1, 2), Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 2), Fraction(1, 2), Fraction(1, 2)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 4), Fraction(1, 4)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+            [Fraction(1, 8), Fraction(1, 8)],
+        ]
+        alpha = 4
+        self.med_st.calculate_AHC_transitions(prior=prior, alpha=alpha)
+        assert self.med_st.alpha is None
+
+    def test_alpha_format(self) -> None:
+        with pytest.raises(TypeError):
+            self.med_st.calculate_AHC_transitions(alpha={"5"})
+
+    def test_merged_leaves_med(self) -> None:
+        # check that no leaves have been merged
+        self.med_st.calculate_AHC_transitions()
+        for stage in self.med_st.ahc_output["Merged Situations"]:
+            for situ in stage:
+                assert situ not in self.med_st.leaves
+
+    def test_merged_leaves_fall(self) -> None:
+        self.fall_st.calculate_AHC_transitions()
+        for stage in self.fall_st.ahc_output["Merged Situations"]:
+            for situ in stage:
+                assert situ not in self.fall_st.leaves
+
+    def test_independent_hyperstage_generator(self) -> None:
+        """generator correctly establishes which subsets have
+        cross-over, and returns correct number of subsets"""
+        example_hyperstage = [
+            ["s1", "s2", "s3"],
+            ["s3", "s4"],
+            ["s5", "s6"],
+            ["s7", "s8"],
+            ["s6", "s4"],
+        ]
+        first_hyperstage = set()
+        first_hyperstage.add(frozenset(example_hyperstage[0]))
+        first_hyperstage.add(frozenset(example_hyperstage[1]))
+        first_hyperstage.add(frozenset(example_hyperstage[2]))
+        first_hyperstage.add(frozenset(example_hyperstage[4]))
+        second_hyperstage = set()
+        second_hyperstage.add(frozenset(example_hyperstage[3]))
+        expected_hyperstages = set()
+        expected_hyperstages.add(frozenset(first_hyperstage))
+        expected_hyperstages.add(frozenset(second_hyperstage))
+
+        actual_hyperstages = set()
+        for hyperstage in self.fall_st._independent_hyperstage_generator(
+            hyperstages=example_hyperstage
+        ):
+            actual_hyperstages.add(
+                frozenset([frozenset(sublist) for sublist in hyperstage])
+            )
+
+        assert actual_hyperstages == expected_hyperstages
+
+    def test_node_colours(self) -> None:
+        """Ensures that all nodes in the event tree dot graph object
+        are coloured in lightgrey and the nodes in the staged tree
+        agree with the result of AHC"""
+        dot_event_nodes = self.med_st.dot_event_graph().get_nodes()
+        self.med_st.calculate_AHC_transitions()
+        dot_staged_nodes = self.med_st.dot_staged_graph().get_nodes()
+        event_node_colours = [
+            n.obj_dict["attributes"]["fillcolor"] for n in dot_event_nodes
+        ]
+        staged_node_colours = [
+            n.obj_dict["attributes"]["fillcolor"] for n in dot_staged_nodes
+        ]
+        assert len(set(event_node_colours)) == 1
+        assert event_node_colours[0] == "lightgrey"
+        assert len(set(staged_node_colours)) > 1
+        g = self.med_st.dot_staged_graph()
+        for node in self.med_st.nodes():
+            assert (
+                g.get_node(node)[0].obj_dict["attributes"]["fillcolor"]
+                == self.med_st.nodes[node]["colour"]
+            )
+
+    def test_new_colours(self) -> None:
+        colours = ["#8dd3c7", "#ffffb3", "#bebada", "#fb8072", "#80b1d3", "#fdb462"]
+        self.fall_st.calculate_AHC_transitions(colour_list=colours)
+        dot_staged_nodes = self.fall_st.dot_staged_graph().get_nodes()
+        staged_node_colours = [
+            n.obj_dict["attributes"]["fillcolor"] for n in dot_staged_nodes
+        ]
+        assert set(colours + ["lightgrey"]) == set(staged_node_colours)
+
+    def test_new_colours_length(self) -> None:
+        colours = ["#8dd3c7", "#ffffb3", "#bebada", "#fb8072"]
+        with pytest.raises(IndexError):
+            self.fall_st.calculate_AHC_transitions(colour_list=colours)
+
+
+class TestChangingDataFrame(unittest.TestCase):
+    """Tests changing DataFrame"""
+
+    def setUp(self):
+        # stratified dataset
+        med_df_path = (
+            Path(__file__)
+            .resolve()
+            .parent.parent.joinpath("../data/medical_dm_modified.xlsx")
+        )
+        self.med_s_z_paths = None
+        self.med_df = pd.read_excel(med_df_path)
+        self.med_st = StagedTree(
+            dataframe=self.med_df, sampling_zero_paths=self.med_s_z_paths
+        )
+        self.med_st.calculate_AHC_transitions()
+        # non-stratified dataset
+        fall_df_path = (
+            Path(__file__).resolve().parent.parent.joinpath("../data/Falls_Data.xlsx")
+        )
+        self.fall_s_z_paths = None
+        self.fall_df = pd.read_excel(fall_df_path)
+        self.fall_st = StagedTree(
+            dataframe=self.fall_df, sampling_zero_paths=self.fall_s_z_paths
+        )
+        self.fall_st.calculate_AHC_transitions()
+
+    def test_add_empty_column(self) -> None:
+        # adding empty column
+        med_empty_column_df = self.med_df
+        med_empty_column_df["extra"] = ""
+        med_empty_column_st = StagedTree(dataframe=med_empty_column_df)
+        med_empty_column_st.calculate_AHC_transitions()
+        assert med_empty_column_st.ahc_output == self.med_st.ahc_output
+
+        fall_empty_column_df = self.fall_df
+        fall_empty_column_df["extra"] = ""
+        fall_empty_column_st = StagedTree(dataframe=fall_empty_column_df)
+        fall_empty_column_st.calculate_AHC_transitions()
+        assert fall_empty_column_st.ahc_output == self.fall_st.ahc_output
+
+    def test_add_NA_column(self) -> None:
+        # adding NA column
+        med_add_NA_df = self.med_df
+        med_add_NA_df["extra"] = np.nan
+        med_add_NA_st = StagedTree(dataframe=med_add_NA_df)
+        med_add_NA_st.calculate_AHC_transitions()
+        assert med_add_NA_st.ahc_output == self.med_st.ahc_output
+
+        fall_add_NA_df = self.fall_df
+        fall_add_NA_df["extra"] = np.nan
+        fall_add_NA_st = StagedTree(dataframe=fall_add_NA_df)
+        fall_add_NA_st.calculate_AHC_transitions()
+        assert fall_add_NA_st.ahc_output == self.fall_st.ahc_output
+
+    def test_add_same_column_med(self) -> None:
+        # adding column with no more information
+        med_add_same_df = self.med_df
+        med_add_same_df["extra"] = "same for all"
+        med_add_same_st = StagedTree(dataframe=med_add_same_df)
+        med_add_same_st.calculate_AHC_transitions()
+        try:
+            med_add_same_st.create_figure("out/test_add_same_column_med_fig.pdf")
+        except InvocationException:
+            pass
+
+        first_set = set(tuple(x) for x in self.med_st.ahc_output["Merged Situations"])
+        second_set = set(
+            tuple(x) for x in med_add_same_st.ahc_output["Merged Situations"]
+        )
+        assert first_set.issubset(second_set)
+
+    def test_add_same_column_fall(self) -> None:
+        fall_add_same_df = self.fall_df
+        fall_add_same_df["extra"] = "same for all"
+        fall_add_same_st = StagedTree(dataframe=fall_add_same_df)
+        fall_add_same_st.calculate_AHC_transitions()
+        try:
+            fall_add_same_st.create_figure("out/test_add_same_column_fall_fig.pdf")
+        except InvocationException:
+            pass
+
+        first_set = set(tuple(x) for x in self.fall_st.ahc_output["Merged Situations"])
+        second_set = set(
+            tuple(x) for x in fall_add_same_st.ahc_output["Merged Situations"]
+        )
+        assert first_set.issubset(second_set)
+
+    def test_add_same_column_int_med(self) -> None:
+        # adding column with no more information
+        med_add_same_df = self.med_df
+        med_add_same_df["extra"] = 1
+        med_add_same_st = StagedTree(dataframe=med_add_same_df)
+        med_add_same_st.calculate_AHC_transitions()
+        try:
+            med_add_same_st.create_figure("out/test_add_same_column_int_med_fig.pdf")
+        except InvocationException:
+            pass
+
+        first_set = set(tuple(x) for x in self.med_st.ahc_output["Merged Situations"])
+        second_set = set(
+            tuple(x) for x in med_add_same_st.ahc_output["Merged Situations"]
+        )
+        assert first_set.issubset(second_set)
+
+    def test_add_same_column_int_fall(self) -> None:
+        fall_add_same_df = self.fall_df
+        fall_add_same_df["extra"] = 1
+        fall_add_same_st = StagedTree(dataframe=fall_add_same_df)
+        fall_add_same_st.calculate_AHC_transitions()
+        try:
+            fall_add_same_st.create_figure("out/test_add_same_column_int_fall_fig.pdf")
+        except InvocationException:
+            pass
+
+        first_set = set(tuple(x) for x in self.fall_st.ahc_output["Merged Situations"])
+        second_set = set(
+            tuple(x) for x in fall_add_same_st.ahc_output["Merged Situations"]
+        )
+        assert first_set.issubset(second_set)
+
+
+class TestWithDynamicDataset(unittest.TestCase):
+    """Dynamic test set test"""
+
+    def setUp(self):
+        self.dataframe = pd.read_excel("data/Falls_Dynamic_Data.xlsx")
+
+    def test_single_floret_stages(self):
+        """Single floret stages are marked as same stage"""
+        dataframe = self.dataframe[
+            ["Residence", "Risk", "Treatment", "Fall", "Outcome"]
+        ]
+        s_tree = StagedTree(dataframe)
+        ahc_out = s_tree.calculate_AHC_transitions()
+        merged_sits = [set(merged_sits) for merged_sits in ahc_out["Merged Situations"]]
+        assert {"s9", "s13", "s14"} in merged_sits
+
+
+class TestNumericalDataset(unittest.TestCase):
+    """Test case for purely numerical dataset."""
+
+    def setUp(self):
+        """setup for tests"""
+        self.data = pd.read_csv("data/Asym.csv")
+
+    def test_string_missing_paths(self):
+        """Missing paths are provided as strings, there's no error."""
+        missing_paths = [
+            ("0", "1", "1", "1"),
+            ("1", "0", "1", "1"),
+            ("0", "1", "0", "1"),
+        ]
+        try:
+            _ = StagedTree(self.data, sampling_zero_paths=missing_paths)
+        # pylint: disable=broad-except
+        except Exception as err:
+            self.fail(f"There was an error when using string missing_paths:\n{err}")
+
+    def test_numerical_missing_paths(self):
+        """Missing paths are provided as numerical data, there's no error."""
+        missing_paths = [(0, 1, 1, 1), (1, 0, 1, 1), (0, 1, 0, 1)]
+        try:
+            _ = StagedTree(self.data, sampling_zero_paths=missing_paths)
+        # pylint: disable=broad-except
+        except Exception as err:
+            self.fail(f"There was an error when using string missing_paths:\n{err}")
+
+
+class TestPosteriorProbabilityCalculations(unittest.TestCase):
+    """Tests the _calculate_and_apply_mean_posterior_probs() functions"""
+
+    def setUp(self) -> None:
+        self.dataframe = pd.DataFrame(
+            [
+                np.array(["1", "Trt1", "Recover"]),
+                np.array(["1", "Trt1", "Dont Recover"]),
+                np.array(["2", "Trt1", "Recover"]),
+                np.array(["2", "Trt1", "Dont Recover"]),
+                np.array(["1", "Trt2", "Recover"]),
+                np.array(["1", "Trt2", "Dont Recover"]),
+                np.array(["2", "Trt2", "Recover"]),
+                np.array(["2", "Trt2", "Dont Recover"]),
+            ]
+        )
+        self.staged = StagedTree(self.dataframe)
+        self.merged_situations = [
+            ("s0", "s1"),
+            ("s3", "s5", "s6"),
+            ("s2",),
+            ("s4",),
+        ]
+        self.probs = [
+            [50, 70],
+            [0],
+            [125, 310],
+            [0],
+            [12, 82],
+            [0],
+            [55, 352],
+        ]
+
+    def test_merged_probabilities(self):
+        actual_mpp = _calculate_mean_posterior_probs(
+            self.staged.situations, self.merged_situations, self.probs
+        )
+        expected_mpp = [
+            [0.417, 0.583],
+            [0.135, 0.865],
+            [0.287, 0.713],
+            [0.128, 0.872],
+        ]
+        self.assertEqual(actual_mpp, expected_mpp)
+
+    def test_merged_probabilities_are_applied(self):
+        """Merged probabilites are applied to the StagedTree."""
+        expected_mpp = [
+            [0.417, 0.583],
+            [0.135, 0.865],
+            [0.287, 0.713],
+            [0.128, 0.872],
+        ]
+        self.staged._apply_mean_posterior_probs(self.merged_situations, expected_mpp)
+        edges = {
+            ("s0", "s1", "1"): 0.417,
+            ("s0", "s2", "2"): 0.583,
+            ("s1", "s3", "Trt1"): 0.417,
+            ("s1", "s4", "Trt2"): 0.583,
+            ("s2", "s5", "Trt1"): 0.287,
+            ("s2", "s6", "Trt2"): 0.713,
+            ("s3", "s7", "Dont Recover"): 0.135,
+            ("s3", "s8", "Recover"): 0.865,
+            ("s4", "s9", "Dont Recover"): 0.128,
+            ("s4", "s10", "Recover"): 0.872,
+            ("s5", "s11", "Dont Recover"): 0.135,
+            ("s5", "s12", "Recover"): 0.865,
+            ("s6", "s13", "Dont Recover"): 0.135,
+            ("s6", "s14", "Recover"): 0.865,
+        }
+        for edge, probability in edges.items():
+            self.assertEqual(self.staged.edges[edge]["probability"], probability)
diff --git a/src/tests/test_util.py b/src/tests/test_util.py
new file mode 100644
index 0000000..d8ed8d0
--- /dev/null
+++ b/src/tests/test_util.py
@@ -0,0 +1,57 @@
+"""Tests cegpy.utilities"""
+from collections import defaultdict
+import unittest
+from cegpy.utilities import (
+    check_list_contains_strings,
+    check_tuple_contains_strings,
+    create_sampling_zeros,
+)
+
+
+class TestCegUtil(unittest.TestCase):
+    """Tests util functions"""
+
+    def setUp(self):
+        self.sampling_zero_list = [("path", "to"), ("road", "from", "nowhere")]
+        self.sample_path_dict = defaultdict(int)
+        self.sample_path_dict[("path",)] += 3
+        self.sample_path_dict[("path", "away")] += 2
+        self.sample_path_dict[("road",)] += 3
+        self.sample_path_dict[("road", "from")] += 1
+        self.sample_path_dict[("road", "to")] += 2
+        self.sample_path_dict[("road", "from", "somewhere")] += 1
+        self.sample_path_dict[("road", "to", "somewhere")] += 1
+        self.sample_path_dict[("road", "to", "nowhere")] += 1
+
+    def test_check_list_contains_strings(self) -> None:
+        """Tests function check_list_contains_strings().
+        Passes two lists. One that only contains strings, and another
+        that contains a multitude of types."""
+
+        str_list = ["blah", "path/to/node", "123", "~$%^"]
+        non_str_list = [123, "path/to/node", True]
+
+        self.assertTrue(check_list_contains_strings(str_list))
+        self.assertFalse(check_list_contains_strings(non_str_list))
+
+    def test_create_sampling_zeros(self) -> None:
+        """Tests function create_sampling_zeros().
+        Passes the new sampling zero list of tree path tuples, and a sample
+        dictionary that does not contain those paths. It then checks that the
+        keys exist, and that their values are zero."""
+        test_path_dict = create_sampling_zeros(
+            self.sampling_zero_list, self.sample_path_dict
+        )
+
+        for key in self.sampling_zero_list:
+            self.assertIn(key, list(test_path_dict.keys()))
+            self.assertEqual(test_path_dict[key], 0)
+
+    def test_check_tuple_contains_strings(self) -> None:
+        """Tests function check_tuple_contains_strings().
+        Passes various tuples that contain different types,
+        and verifies output"""
+        self.assertTrue(check_tuple_contains_strings(("string thing",)))
+        self.assertFalse(check_tuple_contains_strings(("string thing")))
+        self.assertTrue(check_tuple_contains_strings(("string one", "string two")))
+        self.assertFalse(check_tuple_contains_strings((1, "2")))
diff --git a/tests/__init__.py b/tests/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/tests/test_ceg.py b/tests/test_ceg.py
deleted file mode 100644
index f2000d4..0000000
--- a/tests/test_ceg.py
+++ /dev/null
@@ -1,236 +0,0 @@
-import networkx as nx
-import pandas as pd
-from src.cegpy import StagedTree, Evidence, ChainEventGraph
-from pathlib import Path
-
-
-class TestUnitCEG(object):
-    def setup(self):
-        self.node_prefix = 'w'
-        self.sink_suffix = '&infin;'
-        df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/medical_dm_modified.xlsx')
-
-        self.st = StagedTree(
-            dataframe=pd.read_excel(df_path),
-            name="medical_staged"
-        )
-        self.st.calculate_AHC_transitions()
-        self.ceg = ChainEventGraph(self.st)
-
-    def test_node_name_generation(self):
-        prefix = self.ceg.node_prefix
-        largest = 20
-        node_names = [
-            self.ceg._ChainEventGraph__get_next_node_name()
-            for _ in range(0, largest)
-        ]
-        assert (prefix + '1') == node_names[0]
-        assert (prefix + str(largest)) == node_names[largest - 1]
-
-    def test_trim_leaves_from_graph(self) -> None:
-        self.ceg._ChainEventGraph__trim_leaves_from_graph()
-        for leaf in self.st.leaves:
-            try:
-                self.ceg.nodes[leaf]
-                leaf_removed = False
-
-            except KeyError:
-                leaf_removed = True
-
-            assert leaf_removed
-
-            for edge_list_key in self.ceg.edges.keys():
-                assert edge_list_key[1] != leaf
-
-    def test_update_distances_of_nodes_to_sink(self) -> None:
-        def check_distances():
-            actual_node_dists = \
-                nx.get_node_attributes(self.ceg, 'max_dist_to_sink')
-            for node, distance in actual_node_dists.items():
-                assert distance == expected_node_dists[node]
-        root_node = self.node_prefix + '0'
-        sink_node = self.node_prefix + self.sink_suffix
-        expected_node_dists = {
-            root_node: 4,
-            's1': 3,
-            's2': 3,
-            's3': 2,
-            's4': 2,
-            's5': 2,
-            's6': 2,
-            's7': 2,
-            's8': 2,
-            's9': 1,
-            's10': 1,
-            's11': 1,
-            's12': 1,
-            's13': 1,
-            's14': 1,
-            's15': 1,
-            's16': 1,
-            's17': 1,
-            's18': 1,
-            's19': 1,
-            's20': 1,
-            sink_node: 0
-        }
-        nx.relabel_nodes(self.ceg, {'s0': self.ceg.root_node}, copy=False)
-        self.ceg._ChainEventGraph__trim_leaves_from_graph()
-        self.ceg._ChainEventGraph__update_distances_of_nodes_to_sink_node()
-        check_distances()
-
-        # Add another edge to the dictionary, to show that the path is max,
-        # and not min distance to sink
-        self.ceg.add_edge('s3', self.ceg.sink_node)
-        self.ceg.add_edge('s1', self.ceg.sink_node)
-        self.ceg.add_edge('s2', 's10')
-        self.ceg._ChainEventGraph__update_distances_of_nodes_to_sink_node()
-        check_distances()
-
-    def test_gen_nodes_with_increasing_distance(self) -> None:
-        expected_nodes = {
-            0: ['s21', 's22', 's23', 's24', 's25', 's26', 's27', 's28', 's29',
-                's30', 's31', 's32', 's33', 's34', 's35', 's36', 's37', 's38',
-                's39', 's40', 's41', 's42', 's43', 's44'],
-            1: ['s9', 's10', 's11', 's12', 's13', 's14', 's15', 's16', 's17',
-                's18', 's19', 's20'],
-            2: ['s3', 's4', 's5', 's6', 's7', 's8'],
-            3: ['s1', 's2']
-        }
-        nx.relabel_nodes(self.ceg, {'s0': self.ceg.root_node}, copy=False)
-        self.ceg._ChainEventGraph__trim_leaves_from_graph()
-        self.ceg._ChainEventGraph__update_distances_of_nodes_to_sink_node()
-        nodes_gen = self.ceg.\
-            _ChainEventGraph__gen_nodes_with_increasing_distance(
-                start=0
-            )
-
-        for nodes in range(len(expected_nodes)):
-            expected_node_list = expected_nodes[nodes]
-            actual_node_list = next(nodes_gen)
-            assert actual_node_list.sort() == expected_node_list.sort()
-
-    def test_adding_evidence(self) -> None:
-        certain_edges = [
-            ('w1', 'w4', 'Experienced'),
-            ('w4', 'w10', 'Hard'),
-        ]
-        for (u, v, k) in certain_edges:
-            self.ceg.evidence.add_edge(u, v, k, Evidence.CERTAIN)
-            assert (u, v, k) in self.ceg.evidence.certain_edges
-
-        certain_vertices = {
-            'w4', 'w10'
-        }
-        for vertex in certain_vertices:
-            self.ceg.evidence.add_node(vertex, Evidence.CERTAIN)
-            assert vertex in self.ceg.evidence.certain_vertices
-
-        uncertain_edges = [
-            ('w2', 'w5', 'Experienced'),
-            ('w2', 'w6', 'Novice')
-        ]
-        for (u, v, k) in uncertain_edges:
-            self.ceg.evidence.add_edge(u, v, k, Evidence.UNCERTAIN)
-            assert (u, v, k) in self.ceg.evidence.uncertain_edges
-
-        uncertain_vertices = {
-            'w2', 'w13'
-        }
-        for vertex in uncertain_vertices:
-            self.ceg.evidence.add_node(vertex, Evidence.UNCERTAIN)
-            assert vertex in self.ceg.evidence.uncertain_vertices
-
-    def test_propagation(self) -> None:
-        self.ceg.generate()
-        uncertain_edges = [
-            ('w2', 'w5', 'Experienced'),
-            ('w2', 'w6', 'Novice')
-        ]
-        certain_nodes = {
-            'w12'
-        }
-        self.ceg.evidence.add_edges_from(uncertain_edges, Evidence.UNCERTAIN)
-        self.ceg.evidence.add_nodes_from(certain_nodes, Evidence.CERTAIN)
-        self.ceg.reduced
-
-        self.ceg.clear_evidence()
-
-
-class TestEvidence(object):
-    def setup(self):
-        G = nx.MultiDiGraph()
-        nodes = ['w0', 'w1', 'w2', 'w3', 'w4',
-                 'w5', 'w6', 'w7', 'w8', 'w&infin;']
-        edges = [
-            ('w0', 'w1', 'a'),
-            ('w0', 'w1', 'b'),
-            ('w1', 'w2', 'c'),
-            ('w1', 'w3', 'd'),
-            ('w2', 'w4', 'e'),
-            ('w2', 'w4', 'f'),
-            ('w3', 'w4', 'g'),
-            ('w3', 'w8', 'h'),
-            ('w4', 'w8', 'i'),
-            ('w0', 'w5', 'j'),
-            ('w0', 'w7', 'k'),
-            ('w0', 'w6', 'l'),
-            ('w5', 'w6', 'm'),
-            ('w5', 'w7', 'n'),
-            ('w6', 'w7', 'o'),
-            ('w7', 'w8', 'p'),
-            ('w8', 'w&infin;', 'q'),
-            ('w4', 'w&infin;', 'r'),
-            ('w0', 'w3', 's')
-        ]
-        G.add_nodes_from(nodes)
-        G.add_edges_from(edges)
-        H = ChainEventGraph(G)
-        self.evidence = Evidence(H)
-        print(self.evidence)
-
-    def test_edge_add_and_remove(self):
-        certain_edges = [
-            ('s1', 's2', 'Experienced'),
-            ('s1', 's2', 'Other'),
-            ('s1', 's2', 'Novice'),
-        ]
-        for edge in certain_edges:
-            self.evidence.add_edge(edge[0], edge[1], edge[2], certain=True)
-        assert certain_edges == self.evidence.certain_edges
-        self.evidence.remove_edge('s1', 's2', 'Other', certain=True)
-        certain_edges.remove(('s1', 's2', 'Other'))
-        assert certain_edges == self.evidence.certain_edges
-
-        uncertain_edges = [
-            ('s1', 's2', 'Experienced'),
-            ('s1', 's2', 'Other'),
-            ('s1', 's2', 'Novice'),
-        ]
-        for edge in uncertain_edges:
-            self.evidence.add_edge(edge[0], edge[1], edge[2], certain=False)
-        assert uncertain_edges == self.evidence.uncertain_edges
-        self.evidence.remove_edge('s1', 's2', 'Other', certain=False)
-        uncertain_edges.remove(('s1', 's2', 'Other'))
-        assert uncertain_edges == self.evidence.uncertain_edges
-
-    def test_add_vertex_add_and_remove(self):
-        uncertain_vertices = {'s1', 's2', 's3', 's45'}
-        for vertex in uncertain_vertices:
-            self.evidence.add_node(vertex, certain=False)
-        assert uncertain_vertices == self.evidence.uncertain_vertices
-
-        self.evidence.remove_node('s2', certain=False)
-        uncertain_vertices.remove('s2')
-        assert uncertain_vertices == self.evidence.uncertain_vertices
-
-        certain_vertices = {'s1', 's2', 's3', 's45'}
-        for vertex in certain_vertices:
-            self.evidence.add_node(vertex, certain=True)
-        assert certain_vertices == self.evidence.certain_vertices
-
-        self.evidence.remove_node('s2', certain=True)
-        certain_vertices.remove('s2')
-        assert certain_vertices == self.evidence.certain_vertices
diff --git a/tests/test_ctceg.py b/tests/test_ctceg.py
deleted file mode 100644
index 9f76914..0000000
--- a/tests/test_ctceg.py
+++ /dev/null
@@ -1,5 +0,0 @@
-# from ct_chain_event_graph import CTChainEventGraph
-
-
-def test_some_func() -> None:
-    assert 1 == 1
diff --git a/tests/test_event.py b/tests/test_event.py
deleted file mode 100644
index ed5d86b..0000000
--- a/tests/test_event.py
+++ /dev/null
@@ -1,295 +0,0 @@
-import pandas as pd
-import re
-import numpy as np
-import pytest
-from collections import defaultdict
-from pathlib import Path
-from pytest_mock import MockerFixture
-from src.cegpy import EventTree
-
-
-class TestEventTreeAPI():
-    def setup(self):
-        df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/medical_dm_modified.xlsx')
-        self.df = pd.read_excel(df_path)
-
-    def test_required_argument_missing_fails(self):
-        pytest.raises(TypeError, EventTree)
-
-    def test_required_argument_wrong_type_fails(self):
-        dataframe = 5
-        pytest.raises(ValueError, EventTree, dataframe=dataframe)
-
-    def test_incorrect_sampling_zero_fails(self):
-        szp = [('edge_1'), ('edge_1', 'edge_2')]
-        pytest.raises(
-            ValueError,
-            EventTree,
-            dataframe=self.df,
-            sampling_zero_paths=szp)
-
-
-class TestEventTree():
-    def setup(self):
-        df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/medical_dm_modified.xlsx')
-
-        self.df = pd.read_excel(df_path)
-        self.et = EventTree(dataframe=self.df)
-        self.reordered_et = EventTree(
-            dataframe=self.df, 
-            var_order=self.df.columns[::-1]
-        )
-        self.node_format = re.compile('^s\\d\\d*$')
-
-    def test_check_sampling_zero_paths_param(self) -> None:
-        """Tests the function that is checking the sampling zero paths param"""
-        szp = [('Medium',), ('Medium', 'High')]
-        assert self.et._EventTree__check_sampling_zero_paths_param(szp) == szp
-        szp = [1, 2, 3, 4]
-        assert self.et._EventTree__check_sampling_zero_paths_param(szp) is None
-
-        szp = [('path', 'to'), (123, 'something'), 'path/to']
-        assert self.et._EventTree__check_sampling_zero_paths_param(szp) is None
-
-    def test_check_sampling_zero_get_and_set(self) -> None:
-        """Tests the functions that set and get the sampling zeros"""
-        assert self.et.sampling_zeros is None
-
-        szp = [('Medium',), ('Medium', 'High')]
-        self.et.sampling_zeros = szp
-        assert self.et.sampling_zeros == szp
-    
-    def test_order_of_columns(self) -> None:
-        assert self.reordered_et.variables == list(self.df.columns[::-1])
-
-    def test_create_node_list_from_paths(self) -> None:
-        paths = defaultdict(int)
-        paths[('path',)] += 1
-        paths[('path', 'to')] += 1
-        paths[('path', 'away')] += 1
-        paths[('road',)] += 1
-        paths[('road', 'to')] += 1
-        paths[('road', 'away')] += 1
-
-        # code being tested:
-        node_list = self.et._EventTree__create_node_list_from_paths(paths)
-
-        print(node_list)
-        assert len(list(paths.keys())) + 1 == len(node_list)
-        assert node_list[0] == 's0'
-        assert node_list[-1] == 's%d' % (len(node_list) - 1)
-
-    def test_construct_event_tree(self) -> None:
-        """Tests the construction of an event tree from a set of paths,
-        nodes, and """
-        EXPECTED_NODE_COUNT = 45
-        assert len(self.et) == EXPECTED_NODE_COUNT
-        assert len(self.et.edges) == EXPECTED_NODE_COUNT - 1
-        edge_counts = self.et.edge_counts
-
-        assert len(edge_counts) == EXPECTED_NODE_COUNT - 1
-        for _, count in edge_counts.items():
-            assert isinstance(count, int)
-
-    def test_get_functions_producing_expected_data(self) -> None:
-        edges = list(self.et.edges)
-        assert isinstance(edges, list)
-        for edge in edges:
-            assert isinstance(edge, tuple)
-            assert len(edge) == 3
-            assert isinstance(edge[0], str)
-            assert isinstance(edge[1], str)
-            assert isinstance(edge[2], str)
-
-        check_list_contains_strings(list(self.et))
-        check_list_contains_strings(self.et.situations)
-        check_list_contains_strings(self.et.leaves)
-
-        edge_counts = self.et.edge_counts
-        print(edge_counts)
-        assert isinstance(edge_counts, dict)
-        for edge, count in edge_counts.items():
-            assert isinstance(edge, tuple)
-            for node in edge:
-                assert isinstance(node, str)
-            assert isinstance(count, int)
-
-
-class TestIntegration():
-    def setup(self):
-        # stratified dataset
-        med_df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/medical_dm_modified.xlsx')
-        self.med_s_z_paths = None
-        self.med_df = pd.read_excel(med_df_path)
-        self.med_et = EventTree(
-            dataframe=self.med_df,
-            sampling_zero_paths=self.med_s_z_paths
-        )
-
-        # non-stratified dataset
-        fall_df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/Falls_Data.xlsx')
-        self.fall_s_z_paths = None
-        self.fall_df = pd.read_excel(fall_df_path)
-        self.fall_et = EventTree(
-            dataframe=self.fall_df,
-            sampling_zero_path=self.fall_s_z_paths
-        )
-
-    def test_categories_per_variable(self) -> None:
-        expected_med_cats_per_var = {
-            "Classification": 2,
-            "Group": 3,
-            "Difficulty": 2,
-            "Response": 2,
-        }
-        actual_med_cats_per_var = self.med_et.categories_per_variable
-        assert expected_med_cats_per_var == actual_med_cats_per_var
-
-        expected_fall_cats_per_var = {
-            "HousingAssessment": 4,
-            "Risk": 2,
-            "Treatment": 3,
-            "Fall": 2,
-        }
-        actual_fall_cats_per_var = self.fall_et.categories_per_variable
-        assert expected_fall_cats_per_var == actual_fall_cats_per_var
-
-
-def check_list_contains_strings(str_list) -> bool:
-    assert isinstance(str_list, list)
-    for elem in str_list:
-        assert isinstance(elem, str)
-
-
-class TestUsecase():
-    def setup(self):
-        # stratified dataset
-        med_df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/medical_dm_modified.xlsx')
-        self.med_s_z_paths = None
-        self.med_df = pd.read_excel(med_df_path)
-        self.med_et = EventTree(
-            dataframe=self.med_df,
-            sampling_zero_paths=self.med_s_z_paths
-        )
-
-        # non-stratified dataset
-        fall_df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/Falls_Data.xlsx')
-        self.fall_s_z_paths = None
-        self.fall_df = pd.read_excel(fall_df_path)
-        self.fall_et = EventTree(
-            dataframe=self.fall_df
-        )
-
-    def test_fall_cats_per_var(self):
-        expected_fall_cats_per_var = {
-            "HousingAssessment": 4,
-            "Risk": 2,
-            "Treatment": 3,
-            "Fall": 2,
-        }
-        actual_fall_cats_per_var = self.fall_et.categories_per_variable
-        assert expected_fall_cats_per_var == actual_fall_cats_per_var
-
-
-class TestChangingDataFrame():
-    def setup(self):
-        # stratified dataset
-        med_df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/medical_dm_modified.xlsx')
-        self.med_s_z_paths = None
-        self.med_df = pd.read_excel(med_df_path)
-        self.med_et = EventTree(
-            dataframe=self.med_df,
-            sampling_zero_paths=self.med_s_z_paths
-        )
-
-        # non-stratified dataset
-        fall_df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/Falls_Data.xlsx')
-        self.fall_s_z_paths = None
-        self.fall_df = pd.read_excel(fall_df_path)
-        self.fall_et = EventTree(
-            dataframe=self.fall_df,
-            sampling_zero_path=self.fall_s_z_paths
-        )
-
-    def test_add_empty_column(self) -> None:
-        # adding empty column
-        med_empty_column_df = self.med_df
-        med_empty_column_df["extra"] = ""
-        med_empty_column_et = EventTree(
-            dataframe=med_empty_column_df
-        )
-        assert med_empty_column_et.adj == self.med_et.adj
-
-        fall_empty_column_df = self.fall_df
-        fall_empty_column_df["extra"] = ""
-        fall_empty_column_et = EventTree(
-            dataframe=fall_empty_column_df
-        )
-        assert fall_empty_column_et.adj == self.fall_et.adj
-
-    def test_add_NA_column(self) -> None:
-        # adding NA column
-        med_add_NA_df = self.med_df
-        med_add_NA_df["extra"] = np.nan
-        med_add_NA_et = EventTree(
-            dataframe=med_add_NA_df
-        )
-        assert med_add_NA_et.adj == self.med_et.adj
-
-        fall_add_NA_df = self.fall_df
-        fall_add_NA_df["extra"] = np.nan
-        fall_add_NA_et = EventTree(
-            dataframe=fall_add_NA_df
-        )
-        assert fall_add_NA_et.adj == self.fall_et.adj
-
-    def test_add_same_column(self) -> None:
-        # adding column with no more information
-        med_add_same_df = self.med_df
-        med_add_same_df["extra"] = "same for all"
-        med_add_same_et = EventTree(
-            dataframe=med_add_same_df
-        )
-        assert len(med_add_same_et.leaves) == len(self.med_et.leaves)
-
-        fall_add_same_df = self.fall_df
-        fall_add_same_df["extra"] = "same for all"
-        fall_add_same_et = EventTree(
-            dataframe=fall_add_same_df
-        )
-        assert len(fall_add_same_et.leaves) == len(self.fall_et.leaves)
-
-    def test_add_same_column_int(self, mocker: MockerFixture) -> None:
-        mocker.patch('pydotplus.Dot.write')
-        # adding column with no more information
-        med_add_same_df = self.med_df
-        med_add_same_df["extra"] = 1
-        med_add_same_et = EventTree(
-            dataframe=med_add_same_df
-        )
-        med_add_same_et.create_figure("et_fig_path.pdf")
-        assert len(med_add_same_et.leaves) == len(self.med_et.leaves)
-
-        fall_add_same_df = self.fall_df
-        fall_add_same_df["extra"] = 1
-        fall_add_same_et = EventTree(
-            dataframe=fall_add_same_df
-        )
-        fall_add_same_et.create_figure("et_fig_path.pdf")
-        assert len(fall_add_same_et.leaves) == len(self.fall_et.leaves)
diff --git a/tests/test_staged.py b/tests/test_staged.py
deleted file mode 100644
index 686d174..0000000
--- a/tests/test_staged.py
+++ /dev/null
@@ -1,406 +0,0 @@
-from src.cegpy import StagedTree
-import pandas as pd
-from pathlib import Path
-from fractions import Fraction as frac
-from pytest_mock import MockerFixture
-# import xlsxwriter
-import numpy as np
-
-
-class TestStagedTrees():
-    def setup(self):
-        # stratified dataset
-        med_df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/medical_dm_modified.xlsx')
-        self.med_s_z_paths = None
-        self.med_df = pd.read_excel(med_df_path)
-        self.med_st = StagedTree(
-            dataframe=self.med_df,
-            sampling_zero_paths=self.med_s_z_paths
-        )
-
-        # non-stratified dataset
-        fall_df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/Falls_Data.xlsx')
-        self.fall_s_z_paths = None
-        self.fall_df = pd.read_excel(fall_df_path)
-        self.fall_st = StagedTree(
-            dataframe=self.fall_df,
-            sampling_zero_paths=self.fall_s_z_paths,
-        )
-
-    def test_create_default_prior(self) -> None:
-        # stratified medical dataset
-        # Expected prior calculated by hand for default alpha of 3
-        med_expected_prior = [
-            [frac(3, 2), frac(3, 2)],
-            [frac(1, 2), frac(1, 2), frac(1, 2)],
-            [frac(1, 2), frac(1, 2), frac(1, 2)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-        ]
-
-        fall_expected_prior = [
-            [frac(1, 1), frac(1, 1), frac(1, 1), frac(1, 1)],
-            [frac(1, 2), frac(1, 2)],
-            [frac(1, 2), frac(1, 2)],
-            [frac(1, 2), frac(1, 2)],
-            [frac(1, 2), frac(1, 2)],
-            [frac(1, 6), frac(1, 6), frac(1, 6)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 6), frac(1, 6), frac(1, 6)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 4), frac(1, 4)],
-            [frac(1, 12), frac(1, 12)],
-            [frac(1, 12), frac(1, 12)],
-            [frac(1, 12), frac(1, 12)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 12), frac(1, 12)],
-            [frac(1, 12), frac(1, 12)],
-            [frac(1, 12), frac(1, 12)],
-            [frac(1, 8), frac(1, 8)],
-            [frac(1, 8), frac(1, 8)]
-        ]
-        assert len(med_expected_prior) == 21
-        assert len(fall_expected_prior) == 23
-
-        alpha = 3
-        # check that prior is correct.
-        prior = self.med_st._StagedTree__create_default_prior(alpha)
-        assert med_expected_prior == prior
-
-        alpha = 4
-        prior = self.fall_st._StagedTree__create_default_prior(alpha)
-        assert fall_expected_prior == prior
-
-    def test_create_default_hyperstage(self) -> None:
-        med_expected_hyperstage = [
-            ["s0", "s9", "s10", "s11", "s12", "s13", "s14", "s15", "s16",
-             "s17", "s18", "s19", "s20"],
-            ["s1", "s2"],
-            ["s3", "s4", "s5", "s6", "s7", "s8"]
-        ]
-        fall_expected_hyperstage = [
-            ["s0"],
-            ["s1", "s2", "s3", "s4"],
-            ["s5", "s9"],
-            ["s6", "s10"],
-            ["s7", "s8", "s11", "s12", "s13", "s14", "s15", "s16", "s17",
-             "s22", "s23", "s24", "s25", "s26"]
-        ]
-        med_hyperstage = self.med_st._StagedTree__create_default_hyperstage()
-        print(med_expected_hyperstage)
-        print(med_hyperstage)
-
-        fall_hyperstage = self.fall_st._StagedTree__create_default_hyperstage()
-        print(fall_expected_hyperstage)
-        print(fall_hyperstage)
-        assert med_hyperstage == med_expected_hyperstage
-        assert fall_hyperstage == fall_expected_hyperstage
-
-    def test_calculate_posterior(self) -> None:
-        def calculate_posterior(staged_tree: StagedTree, expected_countset,
-                                alpha, expected_likelihood):
-            actual_countset = staged_tree._StagedTree__create_edge_countset()
-            assert actual_countset == expected_countset
-
-            prior = staged_tree._StagedTree__create_default_prior(alpha)
-            staged_tree.prior = prior
-            expected_posterior = []
-            for idx, countset in enumerate(actual_countset):
-                p_elem = []
-                for jdx, count in enumerate(countset):
-                    p_elem.append(count + prior[idx][jdx])
-
-                expected_posterior.append(p_elem)
-
-            actual_posterior = staged_tree.posterior_list
-            assert actual_posterior == expected_posterior
-            actual_likelihood = staged_tree._calculate_initial_loglikelihood(
-                staged_tree.prior_list,
-                staged_tree.posterior_list,
-            )
-            actual_likelihood = round(actual_likelihood, 2)
-            assert actual_likelihood == expected_likelihood
-
-        falls_expected_edge_countset = [
-            [379, 1539, 2871, 45211],
-            [235, 144],
-            [436, 1103],
-            [1449, 1422],
-            [5375, 39836],
-            [53, 52, 130],
-            [130, 14],
-            [82, 354],
-            [870, 233],
-            [319, 315, 815],
-            [1266, 156],
-            [1103, 4272],
-            [30769, 9067],
-            [11, 42],
-            [28, 24],
-            [58, 72],
-            [102, 28],
-            [12, 2],
-            [65, 254],
-            [151, 164],
-            [379, 436],
-            [974, 292],
-            [133, 23]
-        ]
-        alpha = 4
-        expected_likelihood = -68721.50  # Calculated manually
-        assert len(falls_expected_edge_countset) == 23
-        calculate_posterior(
-            self.fall_st,
-            falls_expected_edge_countset,
-            alpha,
-            expected_likelihood
-        )
-
-        med_expected_edge_countset = [
-            [5491, 5493],
-            [798, 1000, 3693],
-            [799, 998, 3696],
-            [399, 399],
-            [500, 500],
-            [1846, 1847],
-            [400, 399],
-            [500, 498],
-            [1849, 1847],
-            [393, 6],
-            [347, 52],
-            [480, 20],
-            [433, 67],
-            [1482, 364],
-            [1305, 542],
-            [5, 395],
-            [92, 307],
-            [70, 430],
-            [202, 296],
-            [338, 1511],
-            [716, 1131]
-        ]
-        alpha = 3
-        expected_likelihood = -30134.07
-        assert len(med_expected_edge_countset) == 21
-        calculate_posterior(
-            self.med_st,
-            med_expected_edge_countset,
-            alpha,
-            expected_likelihood
-        )
-
-    def test_merged_leaves_med(self) -> None:
-        # check that no leaves have been merged
-        self.med_st.calculate_AHC_transitions()
-        for stage in self.med_st.ahc_output['Merged Situations']:
-            for situ in stage:
-                assert situ not in self.med_st.leaves
-
-    def test_merged_leaves_fall(self) -> None:
-        self.fall_st.calculate_AHC_transitions()
-        for stage in self.fall_st.ahc_output['Merged Situations']:
-            for situ in stage:
-                assert situ not in self.fall_st.leaves
-
-    def test_independent_hyperstage_generator(self) -> None:
-        """generator correctly establishes which subsets have
-        cross-over, and returns correct number of subsets"""
-        example_hyperstage = [
-            ['s1', 's2', 's3'],
-            ['s3', 's4'],
-            ['s5', 's6'],
-            ['s7', 's8'],
-            ['s6', 's4']]
-        first_hyperstage = set()
-        first_hyperstage.add(frozenset(example_hyperstage[0]))
-        first_hyperstage.add(frozenset(example_hyperstage[1]))
-        first_hyperstage.add(frozenset(example_hyperstage[2]))
-        first_hyperstage.add(frozenset(example_hyperstage[4]))
-        second_hyperstage = set()
-        second_hyperstage.add(frozenset(example_hyperstage[3]))
-        expected_hyperstages = set()
-        expected_hyperstages.add(frozenset(first_hyperstage))
-        expected_hyperstages.add(frozenset(second_hyperstage))
-
-        actual_hyperstages = set()
-        for hyperstage in self.fall_st._independent_hyperstage_generator(
-                hyperstage=example_hyperstage):
-            actual_hyperstages.add(
-                frozenset([frozenset(sublist) for sublist in hyperstage]))
-
-        assert actual_hyperstages == expected_hyperstages
-
-
-class TestChangingDataFrame():
-    def setup(self):
-        # stratified dataset
-        med_df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/medical_dm_modified.xlsx')
-        self.med_s_z_paths = None
-        self.med_df = pd.read_excel(med_df_path)
-        self.med_st = StagedTree(
-            dataframe=self.med_df,
-            sampling_zero_paths=self.med_s_z_paths
-        )
-        self.med_st.calculate_AHC_transitions()
-        # non-stratified dataset
-        fall_df_path = Path(__file__).resolve(
-            ).parent.parent.joinpath(
-            'data/Falls_Data.xlsx')
-        self.fall_s_z_paths = None
-        self.fall_df = pd.read_excel(fall_df_path)
-        self.fall_st = StagedTree(
-            dataframe=self.fall_df,
-            sampling_zero_path=self.fall_s_z_paths
-        )
-        self.fall_st.calculate_AHC_transitions()
-
-    def test_add_empty_column(self) -> None:
-        # adding empty column
-        med_empty_column_df = self.med_df
-        med_empty_column_df["extra"] = ""
-        med_empty_column_st = StagedTree(
-            dataframe=med_empty_column_df
-        )
-        med_empty_column_st.calculate_AHC_transitions()
-        assert med_empty_column_st.ahc_output == self.med_st.ahc_output
-
-        fall_empty_column_df = self.fall_df
-        fall_empty_column_df["extra"] = ""
-        fall_empty_column_st = StagedTree(
-            dataframe=fall_empty_column_df
-        )
-        fall_empty_column_st.calculate_AHC_transitions()
-        assert fall_empty_column_st.ahc_output == self.fall_st.ahc_output
-
-    def test_add_NA_column(self) -> None:
-        # adding NA column
-        med_add_NA_df = self.med_df
-        med_add_NA_df["extra"] = np.nan
-        med_add_NA_st = StagedTree(
-            dataframe=med_add_NA_df
-        )
-        med_add_NA_st.calculate_AHC_transitions()
-        assert med_add_NA_st.ahc_output == self.med_st.ahc_output
-
-        fall_add_NA_df = self.fall_df
-        fall_add_NA_df["extra"] = np.nan
-        fall_add_NA_st = StagedTree(
-            dataframe=fall_add_NA_df
-        )
-        fall_add_NA_st.calculate_AHC_transitions()
-        assert fall_add_NA_st.ahc_output == self.fall_st.ahc_output
-
-    def test_add_same_column_med(self, mocker: MockerFixture) -> None:
-        mocker.patch('pydotplus.Dot.write')
-        # adding column with no more information
-        med_add_same_df = self.med_df
-        med_add_same_df["extra"] = "same for all"
-        med_add_same_st = StagedTree(
-            dataframe=med_add_same_df
-        )
-        med_add_same_st.calculate_AHC_transitions()
-        med_add_same_st.create_figure("out/test_add_same_column_med_fig.pdf")
-
-        first_set = set(
-            tuple(x) for x in self.med_st.ahc_output['Merged Situations']
-        )
-        second_set = set(
-            tuple(x) for x in med_add_same_st.ahc_output['Merged Situations']
-        )
-        assert first_set.issubset(second_set)
-
-    def test_add_same_column_fall(self, mocker: MockerFixture) -> None:
-        mocker.patch('pydotplus.Dot.write')
-        fall_add_same_df = self.fall_df
-        fall_add_same_df["extra"] = "same for all"
-        fall_add_same_st = StagedTree(
-            dataframe=fall_add_same_df
-        )
-        fall_add_same_st.calculate_AHC_transitions()
-        fall_add_same_st.create_figure("out/test_add_same_column_fall_fig.pdf")
-
-        first_set = set(
-            tuple(x) for x in self.fall_st.ahc_output['Merged Situations']
-        )
-        second_set = set(
-            tuple(x) for x in fall_add_same_st.ahc_output['Merged Situations']
-        )
-        assert first_set.issubset(second_set)
-
-    def test_add_same_column_int_med(self, mocker: MockerFixture) -> None:
-        mocker.patch('pydotplus.Dot.write')
-        # adding column with no more information
-        med_add_same_df = self.med_df
-        med_add_same_df["extra"] = 1
-        med_add_same_st = StagedTree(
-            dataframe=med_add_same_df
-        )
-        med_add_same_st.calculate_AHC_transitions()
-        med_add_same_st.create_figure(
-            "out/test_add_same_column_int_med_fig.pdf")
-
-        first_set = set(
-            tuple(x) for x in self.med_st.ahc_output['Merged Situations']
-        )
-        second_set = set(
-            tuple(x) for x in med_add_same_st.ahc_output['Merged Situations']
-        )
-        assert first_set.issubset(second_set)
-
-    def test_add_same_column_int_fall(self, mocker: MockerFixture) -> None:
-        mocker.patch('pydotplus.Dot.write')
-        fall_add_same_df = self.fall_df
-        fall_add_same_df["extra"] = 1
-        fall_add_same_st = StagedTree(
-            dataframe=fall_add_same_df
-        )
-        fall_add_same_st.calculate_AHC_transitions()
-        fall_add_same_st.create_figure(
-            "out/test_add_same_column_int_fall_fig.pdf")
-
-        first_set = set(
-            tuple(x) for x in self.fall_st.ahc_output['Merged Situations']
-        )
-        second_set = set(
-            tuple(x) for x in fall_add_same_st.ahc_output['Merged Situations']
-        )
-        assert first_set.issubset(second_set)
-
-
-class TestWithDynamicDataset():
-    def setup(self):
-        self.df = pd.read_excel("data/Falls_Dynamic_Data.xlsx")
-
-    def test_single_floret_stages(self):
-        """Single floret stages are marked as same stage"""
-        df = self.df[["Residence", "Risk", "Treatment", "Fall", "Outcome"]]
-        st = StagedTree(df)
-        ahc_out = st.calculate_AHC_transitions()
-        ms = [set(merged_sits) for merged_sits in ahc_out["Merged Situations"]]
-        assert {"s9", "s13", "s14"} in ms
diff --git a/tests/test_util.py b/tests/test_util.py
deleted file mode 100644
index 26757b2..0000000
--- a/tests/test_util.py
+++ /dev/null
@@ -1,62 +0,0 @@
-from collections import defaultdict
-from src.cegpy import Util
-
-
-class TestCegUtil(object):
-
-    def setup(self):
-        self.sampling_zero_list = [('path', 'to'), ('road', 'from', 'nowhere')]
-        self.sample_path_dict = defaultdict(int)
-        self.sample_path_dict[('path',)] += 3
-        self.sample_path_dict[('path', 'away')] += 2
-        self.sample_path_dict[('road',)] += 3
-        self.sample_path_dict[('road', 'from')] += 1
-        self.sample_path_dict[('road', 'to')] += 2
-        self.sample_path_dict[('road', 'from', 'somewhere')] += 1
-        self.sample_path_dict[('road', 'to', 'somewhere')] += 1
-        self.sample_path_dict[('road', 'to', 'nowhere')] += 1
-
-    def test_check_list_contains_strings(self) -> None:
-        """Tests function util.check_list_contains_strings().
-        Passes two lists. One that only contains strings, and another
-        that contains a multitude of types."""
-
-        str_list = ['blah', 'path/to/node', '123', '~$%^']
-        non_str_list = [123, 'path/to/node', True]
-
-        assert Util.check_list_contains_strings(str_list) is True
-        assert Util.check_list_contains_strings(non_str_list) is False
-
-    def test_create_sampling_zeros(self) -> None:
-        """Tests function Util.create_sampling_zeros().
-        Passes the new sampling zero list of tree path tuples, and a sample
-        dictionary that does not contain those paths. It then checks that the
-        keys exist, and that their values are zero."""
-        test_path_dict = Util.create_sampling_zeros(
-            self.sampling_zero_list,
-            self.sample_path_dict
-            )
-
-        for key in self.sampling_zero_list:
-            assert key in list(test_path_dict.keys())
-            assert test_path_dict[key] == 0
-
-    def test_check_tuple_contains_strings(self) -> None:
-        """Tests function Util.check_tuple_contains_strings().
-        Passes various tuples that contain different types,
-        and verifies output"""
-        tup = ('string thing',)
-        assert Util.check_tuple_contains_strings(tup) is True
-        tup = ('string thing')  # Not a tuple, missing comma
-        assert Util.check_tuple_contains_strings(tup) is False
-        tup = ('string one', 'string two')
-        assert Util.check_tuple_contains_strings(tup) is True
-        tup = (1, '2')
-        assert Util.check_tuple_contains_strings(tup) is False
-
-    # def test_generate_colours(self) -> None:
-    #     high = 6
-    #     for cnt in range(1, high+1):
-    #         print("%d:" % cnt)
-    #         print(Util.generate_colours(cnt))
-    #     assert 1==2