sleipnirgroup-jormungandr 0.2.1__cp312-abi3-macosx_14_0_universal2.whl

jormungandr/__init__.py

@@ -0,0 +1,8 @@
+"""
+A linearity-exploiting sparse nonlinear constrained optimization problem solver
+that uses the interior-point method.
+"""
+
+from ._jormungandr import *
+
+__version__ = "0.2.1"

jormungandr/__init__.pyi

@@ -0,0 +1 @@
+from . import autodiff as autodiff, optimization as optimization

jormungandr/autodiff/__init__.py

@@ -0,0 +1,58 @@
+import scipy.sparse
+
+from .._jormungandr.autodiff import *
+
+
+def spy(mat: scipy.sparse.csc_matrix):
+    """
+    Plot the sparsity pattern of a sparse matrix.
+
+    Green points represent positive values and red points represent negative
+    values.
+
+    Parameter ``mat``:
+        The sparse matrix.
+    """
+    import matplotlib.pyplot as plt
+    import matplotlib.ticker as mticker
+    from matplotlib.colors import ListedColormap
+
+    xs = []
+    ys = []
+    vs = []
+
+    cmat = mat.tocoo()
+    for row, col, value in zip(cmat.row, cmat.col, cmat.data):
+        xs.append(col)
+        ys.append(row)
+        if value > 0.0:
+            vs.append(1.0)
+        elif value < 0.0:
+            vs.append(-1.0)
+        else:
+            vs.append(0.0)
+
+    fig = plt.figure()
+    ax = fig.add_subplot()
+
+    # Display scatter plot
+    cmap = ListedColormap(["red", "green"])
+    sc = ax.scatter(xs, ys, s=1, c=vs, marker=".", cmap=cmap, vmin=-1, vmax=1)
+
+    # Display colorbar
+    colorbar = fig.colorbar(sc)
+    ticklabels = ["Negative", "Positive"]
+    colorbar.ax.yaxis.set_major_locator(mticker.FixedLocator(ticklabels))
+    colorbar.ax.set_yticks([-1, 1])
+    colorbar.ax.set_yticklabels(ticklabels)
+
+    ax.set_title("Sparsity")
+    ax.set_xlabel("Cols")
+    ax.set_ylabel("Rows")
+
+    ax.set_xlim([0, mat.shape[1]])
+    ax.set_ylim([0, mat.shape[0]])
+    ax.invert_yaxis()
+    ax.set_aspect(1.0)
+
+    plt.show()