docstrings for io functions

Tractables · Feb 14, 2024 · 7fa3009 · 7fa3009
1 parent 187258a
commit 7fa3009
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 2 deletions.
diff --git a/docs/source/python-api/pyjuice.rst b/docs/source/python-api/pyjuice.rst
@@ -34,4 +34,14 @@ PC Compilation
     :toctree: generated
     :nosignatures:
 
-    compile
+    compile
+
+IO
+--
+
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    load
+    save
diff --git a/src/pyjuice/io/io.py b/src/pyjuice/io/io.py
@@ -11,6 +11,15 @@
 
 
 def save(fname: str, model: Union[CircuitNodes,TensorCircuit]):
+    """
+    Save a PC to file. The file should end with ".jpc".
+
+    :param fname: file name
+    :type fname: str
+
+    :param model: an uncompiled (`CircuitNodes`) or compiled (`TensorCircuit`) PC
+    :type model: Union[CircuitNodes,TensorCircuit]
+    """
     if isinstance(model, TensorCircuit):
         model.update_parameters()
         root_ns = model.root_ns
@@ -25,7 +34,15 @@ def save(fname: str, model: Union[CircuitNodes,TensorCircuit]):
         raise ValueError(f"Unknown file type `.{fname.split('.')[-1]}`.")
 
 
-def load(fname: str):
+def load(fname: str) -> CircuitNodes:
+    """
+    Load a PC from file. The file should end with ".jpc".
+
+    :param fname: file name
+    :type fname: str
+
+    :returns: a PC node
+    """
     if fname.endswith(".jpc"):
         with open(fname, "rb") as f:
             sel_nodes = pickle.load(f)