Fix unit-tests

Signed-off-by: Adam Li <[email protected]>
py-why · Jul 8, 2024 · aecebeb · aecebeb
1 parent f5359b2
commit aecebeb
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 10 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -117,11 +117,26 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
           architecture: "x64"
+
+      - name: Setup torch for pgmpy
+        if: "matrix.os == 'ubuntu'"
+        shell: bash
+        run: |
+          sudo apt-get update
+          sudo apt-get install nvidia-cuda-toolkit nvidia-cuda-toolkit-gcc
+
       - name: Install packages via pip
         run: |
           pip install --upgrade pip
           pip install numpy scipy networkx statsmodels
           pip install .[test]
+
+      - name: Install DoDiscover (main)
+        run: |
+          git clone https://github.com/py-why/dodiscover.git
+          cd dodiscover
+          pip install .
+
       - name: Install Networkx (main)
         if: "matrix.networkx == 'main'"
         run: |
@@ -131,16 +146,10 @@ jobs:
           pip install .[default]
           # pip install --progress-bar off git+https://github.com/networkx/networkx
 
-      - name: Setup torch for pgmpy
-        if: "matrix.os == 'ubuntu'"
-        shell: bash
-        run: |
-          sudo apt-get update
-          sudo apt-get install nvidia-cuda-toolkit nvidia-cuda-toolkit-gcc
-
       - name: Run pytest # headless via Xvfb on linux
         run: |
           pytest --cov pywhy_graphs ./pywhy_graphs
+
       - name: Upload coverage stats to codecov
         if: ${{ matrix.os == 'ubuntu' && matrix.python-version == '3.11' && matrix.networkx == 'stable' }}
         uses: codecov/codecov-action@v4

diff --git a/pywhy_graphs/algorithms/pag.py b/pywhy_graphs/algorithms/pag.py
@@ -5,9 +5,6 @@
 
 import networkx as nx
 import numpy as np
-from dodiscover import FCI, make_context
-from dodiscover.ci import Oracle
-from dodiscover.constraint.utils import dummy_sample
 
 from pywhy_graphs import ADMG, CPDAG, PAG, StationaryTimeSeriesPAG
 from pywhy_graphs.algorithms.generic import (
@@ -1293,6 +1290,14 @@ def mag_to_pag(G: PAG):
     pag : PAG
         The PAG constructed from the MAG.
     """
+    try:
+        from dodiscover import FCI, make_context
+        from dodiscover.ci import Oracle
+        from dodiscover.constraint.utils import dummy_sample
+    except ImportError as e:
+        raise ImportError(
+            "The 'dodiscover' package is required to convert a MAG to a PAG."
+        )
 
     data = dummy_sample(G)
     oracle = Oracle(G)