Merge pull request #106 from ChristopherMayes/lazy_matplotlib

lazy import matplotlib

xopt/utils.py

@@ -8,7 +8,6 @@
 import pandas as pd
 import torch
 import yaml
-from matplotlib import pyplot as plt
 
 from .pydantic import get_descriptions_defaults
 from .vocs import VOCS
@@ -173,13 +172,20 @@ def read_xopt_csv(*files):
     return pd.concat(dfs)
 
 
-def visualize_model(generator, data):
+def visualize_model(generator, data, axes=None):
     test_x = torch.linspace(*torch.tensor(generator.vocs.bounds.flatten()), 100)
     generator.add_data(data)
     model = generator.train_model()
 
-    fig, ax = plt.subplots(2, 1, sharex="all")
-    fig.set_size_inches(6, 6)
+    if axes is None:
+        # Lazy import
+        from matplotlib import pyplot as plt
+
+        fig, ax = plt.subplots(2, 1, sharex="all")
+        fig.set_size_inches(6, 6)
+    else:
+        ax = axes
+
     with torch.no_grad():
         post = model.posterior(test_x.reshape(-1, 1, 1).double())