isort

jaxleyverse · Dec 6, 2023 · afae93f · afae93f
1 parent 5407e5f
commit afae93f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/jaxley/modules/network.py b/jaxley/modules/network.py
@@ -1,8 +1,9 @@
 import itertools
 from copy import deepcopy
 from typing import Callable, Dict, List, Optional, Tuple, Union
-import matplotlib.pyplot as plt
+
 import jax.numpy as jnp
+import matplotlib.pyplot as plt
 import networkx as nx
 import numpy as np
 import pandas as pd

diff --git a/jaxley/utils/plot_utils.py b/jaxley/utils/plot_utils.py
@@ -68,7 +68,8 @@ def plot_morph(
     counter_highlight_branches = 0
     lines = []
 
-    fig, ax = plt.subplots(1, 1, figsize=figsize)
+    if fig is None or ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=figsize)
 
     for b in range(cell.total_nbranches):
         if parents[b] > -1: