daehn 2.0.0__tar.gz

daehn-2.0.0/.github/scripts/api_contract.py

@@ -0,0 +1,95 @@
+"""Freeze public Python signatures; additions are allowed, removals/changes need review."""
+
+import argparse
+import ast
+import json
+from pathlib import Path
+
+
+def surface(root):
+    entries = {}
+    for path in sorted((root / "src").rglob("*.py")):
+        module = ".".join(path.relative_to(root / "src").with_suffix("").parts)
+        tree = ast.parse(path.read_text(encoding="utf-8"))
+
+        def collect(body, prefix=""):
+            for node in body:
+                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                    if node.name.startswith("_") and node.name not in {
+                        "__init__",
+                        "__call__",
+                    }:
+                        continue
+                    entries[f"{module}:{prefix}{node.name}"] = (
+                        ast.unparse(node.args)
+                        + " | decorators="
+                        + ",".join(ast.unparse(d) for d in node.decorator_list)
+                    )
+                elif isinstance(node, ast.ClassDef) and not node.name.startswith("_"):
+                    collect(node.body, prefix + node.name + ".")
+                    for field in node.body:
+                        if isinstance(field, ast.AnnAssign) and isinstance(
+                            field.target, ast.Name
+                        ):
+                            entries[
+                                f"{module}:{prefix}{node.name}.{field.target.id}"
+                            ] = ast.unparse(field.annotation) + (
+                                "=" + ast.unparse(field.value)
+                                if field.value is not None
+                                else ""
+                            )
+                elif isinstance(node, ast.Assign):
+                    names = [t.id for t in node.targets if isinstance(t, ast.Name)]
+                    if "__all__" in names and isinstance(
+                        node.value, (ast.List, ast.Tuple)
+                    ):
+                        for item in node.value.elts:
+                            if isinstance(item, ast.Constant) and isinstance(
+                                item.value, str
+                            ):
+                                entries[f"{module}:export:{item.value}"] = "declared"
+                    if "_EXPORTS" in names and isinstance(node.value, ast.Dict):
+                        for key, value in zip(node.value.keys, node.value.values):
+                            if isinstance(key, ast.Constant):
+                                entries[f"{module}:export:{key.value}"] = ast.unparse(
+                                    value
+                                )
+
+        collect(tree.body)
+    return entries
+
+
+def compare(current, baseline):
+    return {
+        name: {"expected": value, "actual": current.get(name)}
+        for name, value in baseline.items()
+        if current.get(name) != value
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--write-baseline", action="store_true")
+    args = parser.parse_args()
+    root = Path(__file__).resolve().parents[2]
+    path = root / ".github/validation/api.json"
+    current = surface(root)
+    if args.write_baseline:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(
+            json.dumps(current, indent=2, sort_keys=True) + "\n", encoding="utf-8"
+        )
+    else:
+        changes = compare(current, json.loads(path.read_text()))
+        if changes:
+            raise SystemExit(
+                "Public API changed; review compatibility before updating the baseline:\n"
+                + json.dumps(changes, indent=2)
+            )
+        assert compare({"f": "x"}, {"f": "x"}) == {}
+        assert compare({}, {"f": "x"}) and compare({"f": "y"}, {"f": "x"})
+        print(f"Public API compatible: {len(current)} declarations checked.")
+
+
+if __name__ == "__main__":
+    main()

daehn-2.0.0/.github/scripts/coverage_report.py

@@ -0,0 +1,53 @@
+"""Publish a badge and a small report from coverage.py's measured totals."""
+
+import html
+import json
+import os
+from pathlib import Path
+import sys
+
+
+def main(source, destination):
+    report = json.loads(Path(source).read_text(encoding="utf-8"))
+    totals = report["totals"]
+    percent = float(totals["percent_covered"])
+    color = "4c1" if percent >= 90 else "97ca00" if percent >= 75 else "dfb317"
+    label = f"{percent:.1f}%"
+    svg = (
+        '<svg xmlns="http://www.w3.org/2000/svg" width="130" height="20" '
+        'role="img" aria-label="Coverage: ' + label + '">'
+        "<title>Coverage: " + label + "</title>"
+        '<rect width="75" height="20" fill="#555"/>'
+        '<rect x="75" width="55" height="20" fill="#' + color + '"/>'
+        '<g fill="#fff" text-anchor="middle" font-family="Verdana,sans-serif" '
+        'font-size="11"><text x="37.5" y="14">coverage</text>'
+        '<text x="102.5" y="14">' + html.escape(label) + "</text></g></svg>\n"
+    )
+    destination = Path(destination)
+    destination.mkdir(parents=True, exist_ok=True)
+    (destination / "coverage.svg").write_text(svg, encoding="utf-8")
+    summary = {
+        "commit": os.environ.get("GITHUB_SHA"),
+        "measurement": "statement and branch coverage, entire package",
+        "totals": totals,
+    }
+    (destination / "coverage-summary.json").write_text(
+        json.dumps(summary, indent=2) + "\n", encoding="utf-8"
+    )
+    message = (
+        f"## Coverage: {label}\n\n"
+        "Measured over the entire package, including unexecuted modules. "
+        "The percentage combines statements and branches.\n\n"
+        f"- Covered statements: {totals['covered_lines']}/{totals['num_statements']}\n"
+        f"- Covered branches: {totals['covered_branches']}/{totals['num_branches']}\n"
+        f"- Source commit: {os.environ.get('GITHUB_SHA', 'local')}\n"
+    )
+    (destination / "README.md").write_text(message, encoding="utf-8")
+    if os.environ.get("GITHUB_STEP_SUMMARY"):
+        with open(os.environ["GITHUB_STEP_SUMMARY"], "a", encoding="utf-8") as handle:
+            handle.write(message)
+    print(message)
+
+
+if __name__ == "__main__":
+    main(*sys.argv[1:])

daehn-2.0.0/.github/scripts/documentation.py

@@ -0,0 +1,59 @@
+"""Build the real Sphinx site and validate citation metadata."""
+
+import json
+import re
+from urllib.parse import unquote, urlsplit
+from pathlib import Path
+import subprocess
+import sys
+from pybtex.database import parse_file
+import yaml
+
+
+def main():
+    root = Path(__file__).resolve().parents[2]
+    for link in re.findall(
+        r"\]\(([^)]+)\)", (root / "README.md").read_text(encoding="utf-8")
+    ):
+        url = urlsplit(link)
+        if url.scheme or not url.path:
+            continue
+        target = (root / unquote(url.path)).resolve()
+        if not target.is_relative_to(root) or not target.exists():
+            raise ValueError("Broken README target: " + link)
+    citation = yaml.safe_load((root / "CITATION.cff").read_text(encoding="utf-8"))
+    assert (
+        citation["cff-version"] == "1.2.0" and citation["title"] and citation["authors"]
+    )
+    entries = parse_file(str(root / "CITATION.bib")).entries
+    assert entries and all(entry.fields.get("title") for entry in entries.values())
+    subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "sphinx",
+            "-b",
+            "html",
+            "-W",
+            "--keep-going",
+            str(root / "docs"),
+            str(root / "site"),
+        ],
+        check=True,
+    )
+    report = {
+        "source_pages": len(list((root / "docs").rglob("*.rst")))
+        + len(list((root / "docs").rglob("*.md"))),
+        "citation_entries": len(entries),
+        "html_build": "site/index.html",
+        "builder": "Sphinx / sphinx_rtd_theme",
+        "warnings_as_errors": True,
+    }
+    (root / "documentation-results.json").write_text(
+        json.dumps(report, indent=2), encoding="utf-8"
+    )
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()

daehn-2.0.0/.github/scripts/documentation_examples.py

@@ -0,0 +1,37 @@
+"""Run checked documentation examples in isolated working directories."""
+
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+
+ROOT = Path(__file__).resolve().parents[2]
+
+
+def main():
+    scripts = sorted((ROOT / "docs/examples").glob("*.py"))
+    if not scripts:
+        raise RuntimeError("No checked documentation examples")
+    environment = dict(
+        os.environ,
+        JAX_ENABLE_X64="true",
+        JAX_PLATFORMS="cpu",
+        MPLBACKEND="Agg",
+        PYTHONDONTWRITEBYTECODE="1",
+    )
+    for script in scripts:
+        with tempfile.TemporaryDirectory(prefix="hardnet-doc-example-") as temporary:
+            subprocess.run(
+                [sys.executable, str(script)],
+                cwd=temporary,
+                env=environment,
+                check=True,
+                timeout=240,
+            )
+    print(json.dumps({"examples_passed": len(scripts)}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

daehn-2.0.0/.github/scripts/environment.py

@@ -0,0 +1,46 @@
+"""Record only public runtime configuration for cross-platform CI artifacts."""
+
+from importlib import metadata
+import json
+from pathlib import Path
+import platform
+import sys
+
+import jax
+
+
+def main():
+    versions = {}
+    for name in (
+        "daehn",
+        "jax",
+        "jaxlib",
+        "numpy",
+        "scipy",
+        "ohagi",
+        "purin",
+        "flax",
+        "optax",
+        "orbax-checkpoint",
+    ):
+        try:
+            versions[name] = metadata.version(name)
+        except metadata.PackageNotFoundError:
+            versions[name] = None
+    report = {
+        "python": sys.version,
+        "system": platform.system(),
+        "release": platform.release(),
+        "machine": platform.machine(),
+        "versions": versions,
+        "devices": [str(device) for device in jax.devices()],
+        "x64_enabled": bool(jax.config.jax_enable_x64),
+    }
+    Path("ci-environment.json").write_text(
+        json.dumps(report, indent=2) + "\n", encoding="utf-8"
+    )
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()

daehn-2.0.0/.github/scripts/numerical_regression.py

@@ -0,0 +1,190 @@
+import argparse
+
+
+import json
+
+import os
+
+from pathlib import Path
+
+import subprocess
+
+import sys
+
+import tempfile
+
+import jax
+
+import jax.numpy as jnp
+
+import numpy as np
+
+
+def reload_check(path, inputs, expected, framework, full_metadata=None):
+    code = """
+import importlib.abc, json, sys
+import numpy as np
+class Block(importlib.abc.MetaPathFinder):
+ def find_spec(self, fullname, path=None, target=None):
+  if fullname.split('.')[0] in {'kkthn','daehn','nlpoptnet','flax','orbax'}:
+   raise AssertionError('Deployment imported '+fullname)
+sys.meta_path.insert(0,Block())
+from ohagi.inference import load
+predictor=load(sys.argv[1])
+x=json.loads(sys.argv[2]); expected=json.loads(sys.argv[3])
+for _ in range(25):
+ np.testing.assert_allclose(predictor.predict(x),expected,atol=1e-8,rtol=1e-7)
+ np.testing.assert_allclose(predictor.predict_jax(x),expected,atol=1e-8,rtol=1e-7)
+"""
+    subprocess.run(
+        [
+            sys.executable,
+            "-B",
+            "-c",
+            code,
+            str(path),
+            json.dumps(np.asarray(inputs).tolist()),
+            json.dumps(np.asarray(expected).tolist()),
+        ],
+        check=True,
+        timeout=180,
+    )
+    if framework == "kkthn":
+        from kkthn import KKTHardNet
+
+        facade = KKTHardNet().load(full_metadata or path)
+    elif framework == "daehn":
+        from daehn import DAEHardNet
+
+        facade = DAEHardNet().load(path)
+    else:
+        from nlpoptnet import NLPOptNet
+
+        facade = NLPOptNet().load(full_metadata or path)
+    np.testing.assert_allclose(
+        np.asarray(facade.predict(np.asarray(inputs)[None])).reshape(-1),
+        np.asarray(expected).reshape(-1),
+        atol=1e-8,
+        rtol=1e-7,
+    )
+
+
+def daehn(root):
+    from daehn.dae import (
+        Coordinate,
+        DifferentialState,
+        DerivativeSpec,
+        EquationBlock,
+        DAESystem,
+        TaylorNeighborhood,
+        ForwardDifference,
+        CentralDifference,
+        compile_dae,
+        DAETrainer,
+        DAETrainingConfig,
+    )
+    from daehn import MLPBackbone
+    from ohagi.inference import export, load
+
+    records = {}
+    for name, strategy in [
+        ("taylor", TaylorNeighborhood(deltas=(0.01,), order=1)),
+        ("forward", ForwardDifference(deltas=(0.01,))),
+        ("central", CentralDifference(deltas=(0.01,))),
+    ]:
+        state = DifferentialState("y")
+        system = DAESystem(
+            coordinates=(Coordinate("t"),),
+            differential_states=(state,),
+            derivatives=(DerivativeSpec(state, (1,)),),
+            differential_equations=EquationBlock(
+                lambda c, s, d, p: jnp.array([d[0] + s[0]]),
+                size=1,
+                affine_in_projection_variables=True,
+            ),
+            algebraic_equalities=EquationBlock(affine_in_projection_variables=True),
+            algebraic_inequalities=EquationBlock(affine_in_projection_variables=True),
+            name="decay",
+        )
+        compiled = compile_dae(system, discretization=strategy)
+        trainer = DAETrainer(
+            compiled=compiled,
+            prepared_projection=compiled.prepare_projection(
+                batch_size=2, method="linear"
+            ),
+            backbone=MLPBackbone(hidden_sizes=(4,)),
+            config=DAETrainingConfig(
+                epochs=2,
+                pretrain_epochs=1,
+                batch_size=2,
+                hidden_sizes=(4,),
+                seed=7,
+                print_every=0,
+            ),
+        )
+        x = np.linspace(0, 1, 8)[:, None]
+        p = np.empty((8, 0))
+        y = np.exp(-x)
+        result = trainer.fit(x[:4], p[:4], y[:4], -y[:4], x[4:], p[4:], y[4:], -y[4:])
+        decoded, *_ = trainer.model.forward(
+            result.parameters, jnp.asarray(x[:2]), jnp.asarray(p[:2])
+        )
+        np.testing.assert_allclose(
+            decoded["states"] + decoded["derivatives"], 0, atol=1e-7
+        )
+
+        def graph(inputs):
+            decoded, *_ = trainer.model.forward(
+                result.parameters, jnp.repeat(inputs, 2, axis=0), jnp.empty((2, 0))
+            )
+            return decoded["states"][:1]
+
+        path = root / name / "metadata.json"
+        export(path, graph, input_size=1, output_size=1, framework="daehn")
+        expected = load(path).predict([0.7])
+        reload_check(path, [0.7], expected, "daehn")
+        records[name] = {"epochs": len(result.history), "prediction": expected.tolist()}
+    return records
+
+
+def run_case(directory):
+    """Exercise the installed frontend through short training and fresh-process reload."""
+    jax.config.update("jax_enable_x64", True)
+    original = Path.cwd()
+    os.chdir(directory)
+    try:
+        return daehn(Path(directory))
+    finally:
+        os.chdir(original)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--write-baseline", action="store_true")
+    args = parser.parse_args()
+    root = Path(__file__).resolve().parents[2]
+    with tempfile.TemporaryDirectory() as directory:
+        result = run_case(directory)
+    baseline = root / ".github/validation/numerical.json"
+    if args.write_baseline:
+        baseline.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
+    else:
+        expected = json.loads(baseline.read_text())
+
+        def check(actual, reference):
+            if isinstance(reference, dict):
+                assert actual.keys() == reference.keys()
+                for key in reference:
+                    check(actual[key], reference[key])
+            else:
+                np.testing.assert_allclose(actual, reference, atol=2e-5, rtol=2e-5)
+
+        check(result, expected)
+    (root / "numerical-results.json").write_text(
+        json.dumps(result, indent=2) + "\n", encoding="utf-8"
+    )
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

daehn-2.0.0/.github/scripts/package_install.py

@@ -0,0 +1,105 @@
+"""Install exact stack wheels from PyPI or TestPyPI and verify dependencies."""
+
+import argparse
+from email.parser import BytesParser
+import json
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+import zipfile
+
+TARGET = "daehn"
+DEPENDENCIES = {"purin": "1.0.0", "ohagi": "1.0.0"}
+
+
+def wheel_identity(path):
+    with zipfile.ZipFile(path) as archive:
+        metadata = next(
+            name for name in archive.namelist() if name.endswith(".dist-info/METADATA")
+        )
+        fields = BytesParser().parsebytes(archive.read(metadata))
+    return fields["Name"], fields["Version"]
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--version", required=True)
+    parser.add_argument(
+        "--local-dir", type=Path, help="Use matching local wheels before downloading"
+    )
+    parser.add_argument(
+        "--testpypi", action="store_true", help="Use the isolated testing index"
+    )
+    args = parser.parse_args()
+    target = "kkthn" if args.testpypi and TARGET == "kkt-hardnet" else TARGET
+    index = (
+        "https://test.pypi.org/simple/" if args.testpypi else "https://pypi.org/simple/"
+    )
+    requirements = dict(DEPENDENCIES, **{target: args.version})
+    with tempfile.TemporaryDirectory(prefix="stack-wheels-") as directory:
+        downloads = Path(directory)
+        wheels = []
+        origins = {}
+        for name, version in requirements.items():
+            matches = (
+                []
+                if args.local_dir is None
+                else [
+                    path
+                    for path in args.local_dir.glob("*.whl")
+                    if wheel_identity(path) == (name, version)
+                ]
+            )
+            if len(matches) > 1:
+                raise ValueError("Ambiguous local wheel: " + name)
+            if matches:
+                wheel = matches[0].resolve()
+                origins[name] = "local build"
+            else:
+                subprocess.run(
+                    [
+                        sys.executable,
+                        "-m",
+                        "pip",
+                        "download",
+                        "--no-deps",
+                        "--only-binary=:all:",
+                        "--index-url",
+                        index,
+                        "--dest",
+                        str(downloads),
+                        name + "==" + version,
+                    ],
+                    check=True,
+                )
+                wheel = next(
+                    path
+                    for path in downloads.glob("*.whl")
+                    if wheel_identity(path) == (name, version)
+                )
+                origins[name] = "TestPyPI" if args.testpypi else "PyPI"
+            wheels.append(str(wheel))
+        subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "--index-url",
+                "https://pypi.org/simple/",
+                "--only-binary=:all:",
+                "--report",
+                "package-install-results.json",
+                *wheels,
+            ],
+            check=True,
+        )
+        subprocess.run([sys.executable, "-m", "pip", "check"], check=True)
+    print(
+        json.dumps({"versions": requirements, "stack_wheel_origins": origins}, indent=2)
+    )
+
+
+if __name__ == "__main__":
+    main()

daehn-2.0.0/.github/scripts/package_smoke.py

@@ -0,0 +1,68 @@
+"""Run installed distributions outside the checkout, including numerical and export checks."""
+
+import argparse
+import importlib
+from importlib.metadata import version
+import json
+import os
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+import tempfile
+
+PACKAGE = "daehn"
+DISTRIBUTION = "daehn"
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--version", required=True)
+    parser.add_argument("--distribution", default=DISTRIBUTION)
+    args = parser.parse_args()
+    root = Path(__file__).resolve().parents[2]
+    module = importlib.import_module(PACKAGE)
+    installed = Path(module.__file__).resolve()
+    if root in installed.parents:
+        raise RuntimeError(
+            "Smoke test imported the checkout instead of an installed wheel"
+        )
+    assert version(args.distribution) == args.version
+    assert module.__version__ == args.version
+    environment = dict(
+        os.environ,
+        JAX_PLATFORMS="cpu",
+        JAX_ENABLE_X64="true",
+        PYTHONNOUSERSITE="1",
+        PYTHONDONTWRITEBYTECODE="1",
+        MPLBACKEND="Agg",
+    )
+    environment.pop("PYTHONPATH", None)
+    scripts = sorted((root / "docs/examples").glob("*.py"))
+    assert scripts, "No runnable release examples found"
+    with tempfile.TemporaryDirectory(prefix=PACKAGE + "-smoke-") as directory:
+        work = Path(directory)
+        for script in scripts:
+            target = work / script.name
+            shutil.copyfile(script, target)
+            subprocess.run(
+                [sys.executable, "-I", str(target)],
+                cwd=work,
+                env=environment,
+                check=True,
+                timeout=300,
+            )
+    report = {
+        "distribution": args.distribution,
+        "version": version(args.distribution),
+        "installed_module": str(installed),
+        "examples": [path.name for path in scripts],
+    }
+    Path("package-smoke-results.json").write_text(
+        json.dumps(report, indent=2), encoding="utf-8"
+    )
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()