sleipnirgroup-jormungandr 0.6.7.dev129__tar.gz → 0.6.7.dev131__tar.gz

{sleipnirgroup_jormungandr-0.6.7.dev129 → sleipnirgroup_jormungandr-0.6.7.dev131}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleipnirgroup-jormungandr
-Version: 0.6.7.dev129
+Version: 0.6.7.dev131
 Summary: Reverse mode autodiff library and NLP solver DSL
 License-Expression: BSD-3-Clause
 License-File: LICENSE.txt

{sleipnirgroup_jormungandr-0.6.7.dev129 → sleipnirgroup_jormungandr-0.6.7.dev131}/python/cpp/autodiff/bind_variable.cpp

@@ -98,7 +98,8 @@ void bind_variable(nb::module_& autodiff, nb::class_<Variable<double>>& cls) {
   using namespace nb::literals;
 
   cls.def(nb::init<>(), DOC(slp, Variable, Variable));
-  cls.def(nb::init<double>(), "value"_a, DOC(slp, Variable, Variable, 2));
+  cls.def(nb::init_implicit<double>(), "value"_a,
+          DOC(slp, Variable, Variable, 2));
   cls.def("set_value", nb::overload_cast<double>(&Variable<double>::set_value),
           "value"_a, DOC(slp, Variable, set_value));
   cls.def("value", &Variable<double>::value, DOC(slp, Variable, value));

{sleipnirgroup_jormungandr-0.6.7.dev129 → sleipnirgroup_jormungandr-0.6.7.dev131}/python/cpp/autodiff/bind_variable_block.cpp

@@ -340,12 +340,11 @@ void bind_variable_block(
   cls.def(-nb::self);
 
   // Matrix-scalar/scalar-matrix operations
-  for_each_type<double, int, Variable<double>>([&]<typename T> {
-    cls.def(nb::self * T(), "rhs"_a);
-    cls.def(T() * nb::self, "lhs"_a);
-    cls.def(nb::self / T(), "rhs"_a);
-    cls.def(nb::self /= T(), "rhs"_a, DOC(slp, VariableBlock, operator, idiv));
-  });
+  cls.def(nb::self * Variable<double>(), "rhs"_a);
+  cls.def(Variable<double>() * nb::self, "lhs"_a);
+  cls.def(nb::self / Variable<double>(), "rhs"_a);
+  cls.def(nb::self /= Variable<double>(), "rhs"_a,
+          DOC(slp, VariableBlock, operator, idiv));
   for_each_type<double, int, const Variable<double>&,
                 nb::DRef<Eigen::MatrixXd>>([&]<typename T> {
     cls.def(

{sleipnirgroup_jormungandr-0.6.7.dev129 → sleipnirgroup_jormungandr-0.6.7.dev131}/python/cpp/autodiff/bind_variable_matrix.cpp

@@ -372,12 +372,11 @@ void bind_variable_matrix(nb::module_& autodiff,
   cls.def(-nb::self);
 
   // Matrix-scalar/scalar-matrix operations
-  for_each_type<double, int, Variable<double>>([&]<typename T> {
-    cls.def(nb::self * T(), "rhs"_a);
-    cls.def(T() * nb::self, "lhs"_a);
-    cls.def(nb::self / T(), "rhs"_a);
-    cls.def(nb::self /= T(), "rhs"_a, DOC(slp, VariableMatrix, operator, idiv));
-  });
+  cls.def(nb::self * Variable<double>(), "rhs"_a);
+  cls.def(Variable<double>() * nb::self, "lhs"_a);
+  cls.def(nb::self / Variable<double>(), "rhs"_a);
+  cls.def(nb::self /= Variable<double>(), "rhs"_a,
+          DOC(slp, VariableMatrix, operator, idiv));
   for_each_type<double, int, const Variable<double>&,
                 const VariableBlock<VariableMatrix<double>>,
                 nb::DRef<Eigen::MatrixXd>>([&]<typename T> {
@@ -529,6 +528,18 @@ void bind_variable_matrix(nb::module_& autodiff,
                nb::overload_cast<const VariableMatrix<double>&,
                                  const VariableMatrix<double>&>(&solve<double>),
                "A"_a, "B"_a, DOC(slp, solve));
+  autodiff.def(
+      "solve",
+      [](const VariableMatrix<double>& A, nb::DRef<Eigen::MatrixXd> B) {
+        return solve(A, VariableMatrix<double>(B));
+      },
+      "A"_a, "B"_a, DOC(slp, solve));
+  autodiff.def(
+      "solve",
+      [](nb::DRef<Eigen::MatrixXd> A, const VariableMatrix<double>& B) {
+        return solve(VariableMatrix<double>(A), B);
+      },
+      "A"_a, "B"_a, DOC(slp, solve));
 }
 
 }  // namespace slp

{sleipnirgroup_jormungandr-0.6.7.dev129 → sleipnirgroup_jormungandr-0.6.7.dev131}/python/postprocess_stubs.py

@@ -5,9 +5,11 @@ Fixes up type stubs generated by nanobind's stubgen.
 
 stubgen can't express some things the bindings support at runtime:
 
-1. VariableBlock implicitly converts to VariableMatrix, so every VariableMatrix
-   parameter also accepts a VariableBlock. Return types and callback arguments
-   are left alone since those are always VariableMatrix.
+1. Implicit conversions widen what parameters accept: VariableBlock converts
+   to VariableMatrix, and float converts to Variable. Every VariableMatrix
+   parameter also accepts a VariableBlock, and every Variable parameter also
+   accepts a float. Return types and callback arguments are left alone since
+   those are always VariableMatrix or Variable.
 2. Sparse matrices are annotated as scipy.sparse.csc_matrix[float], but
    scipy-stubs expects a NumPy numeric scalar type argument.
 """
@@ -15,6 +17,7 @@ stubgen can't express some things the bindings support at runtime:
 import argparse
 import ast
 import re
+from dataclasses import dataclass
 from pathlib import Path
 
 PREFIX = r"(?:sleipnir\._sleipnir\.autodiff\.)?"
@@ -26,14 +29,53 @@ PREFIX = r"(?:sleipnir\._sleipnir\.autodiff\.)?"
 #   -> x: VariableMatrix | VariableBlock | None
 VARIABLE_MATRIX = re.compile(
     rf"""
-    (?<![\w.])                                  # Not part of another name
-    ({PREFIX})                                  # Optional module prefix
-    VariableMatrix\b                            # Type name
-    (?!\ \|\ {PREFIX}VariableBlock\b)           # Not already widened
+    (?<![\w.])                         # Not part of another name
+    ({PREFIX})                         # Optional module prefix
+    VariableMatrix\b                   # Type name
+    (?!\ \|\ {PREFIX}VariableBlock\b)  # Not already widened
     """,
     re.VERBOSE,
 )
 
+# Variable, optionally module-qualified. The word boundary excludes
+# VariableMatrix and VariableBlock, and the lookahead prevents already-widened
+# annotations from matching again.
+#
+#   x: Sequence[Variable]
+#   -> x: Sequence[Variable | float]
+VARIABLE = re.compile(
+    rf"""
+    (?<![\w.])         # Not part of another name
+    ({PREFIX})         # Optional module prefix
+    Variable\b         # Type name
+    (?!\ \|\ float\b)  # Not already widened
+    """,
+    re.VERBOSE,
+)
+
+
+@dataclass(frozen=True)
+class Widening:
+    """A type that also accepts another type via implicit conversion."""
+
+    # Matches the narrow type
+    pattern: re.Pattern[str]
+
+    # Replaces the narrow type with the type that converts to it
+    source: str
+
+    # Replaces the narrow type with the widened type
+    widened: str
+
+
+# \1 is "sleipnir._sleipnir.autodiff." prefix or empty
+WIDENINGS = [
+    Widening(
+        VARIABLE_MATRIX, r"\1VariableBlock", r"\1VariableMatrix | \1VariableBlock"
+    ),
+    Widening(VARIABLE, "float", r"\1Variable | float"),
+]
+
 
 def is_callable(annotation: ast.expr) -> bool:
     """Returns True if an annotation contains a Callable type."""
@@ -44,6 +86,17 @@ def is_callable(annotation: ast.expr) -> bool:
     )
 
 
+def is_nested(annotation: ast.expr, pattern: re.Pattern[str]) -> bool:
+    """
+    Returns True if a type matching the pattern appears inside a subscript of
+    an annotation (e.g., Variable in Sequence[Variable]).
+    """
+    return any(
+        isinstance(node, ast.Subscript) and pattern.search(ast.unparse(node.slice))
+        for node in ast.walk(annotation)
+    )
+
+
 def parameters(node: ast.FunctionDef) -> list[ast.arg | None]:
     """Returns a function's parameters in signature order."""
     args = node.args
@@ -68,11 +121,13 @@ def signature(node: ast.FunctionDef) -> list[str | None]:
 def widen_signatures(stub: str) -> str:
     """
     Widens every VariableMatrix in function parameter annotations to
-    VariableMatrix | VariableBlock.
+    VariableMatrix | VariableBlock, and every Variable to Variable | float.
 
     A parameter isn't widened if a sibling overload is identical except it
-    takes VariableBlock there; widening it would make that overload
-    unreachable.
+    takes the source type there (e.g., float instead of Variable); widening it
+    would make that overload unreachable. Types nested in a subscript (e.g.,
+    Sequence[Variable]) are always widened, since the sibling overloads don't
+    accept containers that mix both types.
     """
     # Group overloads by name within each module or class body
     overload_sets: dict[tuple[int, str], list[ast.FunctionDef]] = {}
@@ -81,25 +136,35 @@ def widen_signatures(stub: str) -> str:
             if isinstance(child, ast.FunctionDef):
                 overload_sets.setdefault((id(node), child.name), []).append(child)
 
-    annotations = []
+    # Annotations to widen, and the widenings that apply to each
+    annotations: list[tuple[ast.expr, list[Widening]]] = []
     for overloads in overload_sets.values():
         signatures = [signature(overload) for overload in overloads]
         for overload, sig in zip(overloads, signatures):
             for i, arg in enumerate(parameters(overload)):
-                # Callbacks receive VariableMatrix, so their annotations are
-                # left alone
+                # Callbacks receive VariableMatrix or Variable, so their
+                # annotations are left alone
                 if arg is None or arg.annotation is None or is_callable(arg.annotation):
                     continue
 
-                # Skip if a sibling overload already takes VariableBlock here
-                block_sig = sig.copy()
-                block_sig[i] = VARIABLE_MATRIX.sub(
-                    r"\1VariableBlock", ast.unparse(arg.annotation)
-                )
-                if block_sig != sig and block_sig in signatures:
-                    continue
-
-                annotations.append(arg.annotation)
+                widenings = []
+                for widening in WIDENINGS:
+                    # Skip if a sibling overload already takes the source type
+                    # here, unless the type is nested in a subscript
+                    source_sig = sig.copy()
+                    source_sig[i] = widening.pattern.sub(
+                        widening.source, ast.unparse(arg.annotation)
+                    )
+                    if (
+                        source_sig != sig
+                        and source_sig in signatures
+                        and not is_nested(arg.annotation, widening.pattern)
+                    ):
+                        continue
+
+                    widenings.append(widening)
+
+                annotations.append((arg.annotation, widenings))
 
     # Encode stub string into UTF-8 bytes
     source_bytes = stub.encode()
@@ -110,8 +175,10 @@ def widen_signatures(stub: str) -> str:
         line_byte_offsets.append(line_byte_offsets[-1] + len(line))
 
     # Edit back to front so earlier offsets stay valid
-    for annotation in sorted(
-        annotations, key=lambda node: (node.lineno, node.col_offset), reverse=True
+    for annotation, widenings in sorted(
+        annotations,
+        key=lambda item: (item[0].lineno, item[0].col_offset),
+        reverse=True,
     ):
         # Assertions to satisfy mypy. They can't fail in practice because
         # ast.parse() always sets those fields.
@@ -123,11 +190,9 @@ def widen_signatures(stub: str) -> str:
         end = line_byte_offsets[annotation.end_lineno - 1] + annotation.end_col_offset
 
         # Splice widened type annotations into source_bytes
-        #
-        # \1 is "sleipnir._sleipnir.autodiff." prefix or empty
-        widened = VARIABLE_MATRIX.sub(
-            r"\1VariableMatrix | \1VariableBlock", source_bytes[start:end].decode()
-        )
+        widened = source_bytes[start:end].decode()
+        for widening in widenings:
+            widened = widening.pattern.sub(widening.widened, widened)
         source_bytes = source_bytes[:start] + widened.encode() + source_bytes[end:]
     return source_bytes.decode()
 

{sleipnirgroup_jormungandr-0.6.7.dev129 → sleipnirgroup_jormungandr-0.6.7.dev131}/python/src/sleipnir/__init__.py

@@ -2,4 +2,4 @@
 
 from ._sleipnir import *
 
-__version__ = "0.6.7.dev129"
+__version__ = "0.6.7.dev131"

{sleipnirgroup_jormungandr-0.6.7.dev129 → sleipnirgroup_jormungandr-0.6.7.dev131}/python/test/autodiff/variable_matrix_test.py

@@ -3,12 +3,26 @@ import sys
 
 import numpy as np
 from sleipnir import autodiff
-from sleipnir.autodiff import VariableMatrix
+from sleipnir.autodiff import ExpressionType, Variable, VariableMatrix
 
 # Machine epsilon for the binding's scalar type (double)
 ε = sys.float_info.epsilon
 
 
+def test_mixed_list_constructor() -> None:
+    x = Variable()
+    x.set_value(3.0)
+
+    # float and int elements are implicitly converted to constant Variables
+    mat = VariableMatrix([[x, 1.0], [2, x]])
+
+    assert mat.shape == (2, 2)
+    assert (mat.value() == np.array([[3.0, 1.0], [2.0, 3.0]])).all()
+    assert mat[0, 0].type() == ExpressionType.LINEAR
+    assert mat[0, 1].type() == ExpressionType.CONSTANT
+    assert mat[1, 0].type() == ExpressionType.CONSTANT
+
+
 def test_assignment_to_default() -> None:
     mat = VariableMatrix()
 
@@ -523,13 +537,16 @@ def test_block_free_function() -> None:
 def check_solve(A: VariableMatrix, B: VariableMatrix) -> None:
     print(f"solve(A, B) {A.shape[0]}x{A.shape[1]}")
 
-    X = autodiff.solve(A, B)
-
-    assert X.shape == (A.shape[1], B.shape[1])
+    for X in (
+        autodiff.solve(A, B),
+        autodiff.solve(A, B.value()),
+        autodiff.solve(A.value(), B),
+    ):
+        assert X.shape == (A.shape[1], B.shape[1])
 
-    # Backward error bound of a stable solve
-    residual = np.linalg.norm(A.value() @ X.value() - B.value())
-    assert residual <= 8 * ε * np.linalg.norm(A.value()) * np.linalg.norm(X.value())
+        # Backward error bound of a stable solve
+        residual = np.linalg.norm(A.value() @ X.value() - B.value())
+        assert residual <= 8 * ε * np.linalg.norm(A.value()) * np.linalg.norm(X.value())
 
 
 def test_solve_free_function() -> None: