kkt-hardnet 1.0.1__tar.gz

kkt_hardnet-1.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 MM Faruque Hasan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

kkt_hardnet-1.0.1/PKG-INFO

@@ -0,0 +1,156 @@
+Metadata-Version: 2.4
+Name: kkt-hardnet
+Version: 1.0.1
+Summary: KKT-HardNet: Physics-constrained neural networks with hard nonlinear equality and inequality constraint satisfaction guarantees
+Author-email: Rahul Golder <rahulgolder8420@tamu.edu>, Bimol Nath Roy <bimolnathroy@tamu.edu>, Ashfaq Iftakher <iftakher@tamu.edu>, MM Faruque Hasan <hasan@tamu.edu>
+License: MIT
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax==0.4.30
+Requires-Dist: jaxlib==0.4.30
+Requires-Dist: flax==0.8.5
+Requires-Dist: optax==0.2.4
+Requires-Dist: chex==0.1.90
+Requires-Dist: numpy==1.26.4
+Requires-Dist: scipy==1.13.1
+Requires-Dist: pandas==2.2.3
+Requires-Dist: matplotlib==3.9.4
+Provides-Extra: cuda12
+Requires-Dist: jax[cuda12]==0.4.30; extra == "cuda12"
+Requires-Dist: jaxlib==0.4.30; extra == "cuda12"
+Dynamic: license-file
+
+# kkt-hardnet
+
+`kkt-hardnet` is the publishable Python package for KKT-HardNet.
+
+## Install
+
+Editable install from this repository:
+
+```bash
+pip install -e kkthn
+```
+
+Editable install with CUDA 12 support:
+
+```bash
+pip install -e "kkthn[cuda12]"
+```
+
+Or use pip after publishing:
+
+```bash
+pip install kkt-hardnet
+```
+
+```bash
+pip install "kkt-hardnet[cuda12]"
+```
+
+## Import
+
+```python
+from kkthn import KKTHardNet
+```
+
+## Core Methods
+
+- `model()` for supervised surrogate learning from `parameters.csv` and `variables.csv`
+- `optimize()` for unsupervised optimization from `parameters.csv`
+- `estimate()` for inverse parameter estimation from `parameters.csv` and `variables.csv`
+- `load(metadata_path)` to reload a trained run
+- `predict(x)` to infer variables for new parameter values
+
+## Packaging
+
+The PyPI project name is `kkt-hardnet`, while the import path remains `kkthn`. The default install is CPU-oriented; CUDA support is selected explicitly with the `cuda12` extra instead of automatic device detection.
+
+# Modeling Workflow
+
+Use `KKTHardNet` to define symbolic constrained problems with:
+
+- named parameters, decision variables, and optional inverse parameters
+- equality and inequality constraints written as Python expressions
+- optional objectives for optimization tasks
+- hard KKT projection during training and inference
+- saved run artifacts including weights, predictions, history, summary, and metadata
+
+Available methods:
+
+- `dataset(parameters=..., variables=...)` to attach training data
+- `model()` for supervised surrogate learning
+- `optimize()` for unsupervised optimization
+- `estimate()` for inverse parameter estimation
+- `load(metadata_path)` to reload a trained run
+- `predict(x)` to evaluate a trained or loaded model on new inputs
+
+The example below shows one complete workflow. The CSV column names must match the declared parameter and variable names. Each run writes `<model_name>_<timestamp>/` with `parameters.csv`, optional `variables.csv`, `history.csv`, `predictions.csv`, `model_weights.npz`, `summary.json`, and `metadata.json`.
+
+```python
+from kkthn import KKTHardNet
+
+TRAIN = {
+    "epochs": 1200,
+    "batch_size": 32,
+    "learning_rate": 1e-3,
+    "train_frac": 0.8,
+    "hidden_size": 64,
+    "hidden_layers": 2,
+    "seed": 42,
+    "dtype": "float64",
+    "print_every": 1,
+    "newton_step_length": 0.5,
+    "newton_tol": 1e-6,
+    "newton_reg_factor": 1e-3,
+    "max_newton_iter": 30,
+    "max_backtrack_iter": 10,
+}
+
+# Build a symbolic problem.
+model = KKTHardNet(name="demo_model", train=TRAIN)
+x = model.add_parameter(["x1", "x2"])
+theta = model.add_inverse_parameter(["a0", "a1"], init_value=[10.0, -10.0])
+y = model.add_variable(["y1", "y2", "y3"])
+
+# Objectives are optional for surrogate modeling and inverse estimation,
+# but required for optimize().
+model.objective = 0.5 * (y.y1**2 + y.y2**2 + y.y3**2)
+model.constraints.add(
+    theta.a0 * y.y1 + y.y2 - x.x1 == 0,
+    y.y2 - theta.a1 * y.y3 - x.x2 == 0,
+    y.y1**2 + y.y3**2 <= 2.0,
+    y.y1 >= 0,
+)
+
+# Attach data.
+# For model() or estimate(): provide both parameters and variables.
+# For optimize(): provide only parameters.
+model.dataset(parameters="parameters.csv", variables="variables.csv")
+
+# Choose one training mode.
+surrogate_result = model.model()
+# optimize_result = model.optimize()
+# estimate_result = model.estimate()
+
+# Reload a saved run later and predict on a new parameter vector.
+reloaded = KKTHardNet()
+reloaded.load("demo_model_20260414_120000/metadata.json")
+prediction = reloaded.predict([1.0, 2.0])
+```
+⚠️ Please cite our work if you use this code in your research.
+Citation formats are provided below.
+
+arXiv Preprint: https://arxiv.org/pdf/2507.08124
+Journal: https://doi.org/10.1016/j.compchemeng.2025.109418
+
+```bibtex
+@article{iftakher2025physics,
+  title={Physics-informed neural networks with hard nonlinear equality and inequality constraints},
+  author={Iftakher, Ashfaq and Golder, Rahul and Nath Roy, Bimol and Hasan, MM Faruque},
+  journal={Computers \& Chemical Engineering},
+  pages={109418},
+  year={2025},
+  publisher={Elsevier}
+}

kkt_hardnet-1.0.1/README.md

@@ -0,0 +1,133 @@
+# kkt-hardnet
+
+`kkt-hardnet` is the publishable Python package for KKT-HardNet.
+
+## Install
+
+Editable install from this repository:
+
+```bash
+pip install -e kkthn
+```
+
+Editable install with CUDA 12 support:
+
+```bash
+pip install -e "kkthn[cuda12]"
+```
+
+Or use pip after publishing:
+
+```bash
+pip install kkt-hardnet
+```
+
+```bash
+pip install "kkt-hardnet[cuda12]"
+```
+
+## Import
+
+```python
+from kkthn import KKTHardNet
+```
+
+## Core Methods
+
+- `model()` for supervised surrogate learning from `parameters.csv` and `variables.csv`
+- `optimize()` for unsupervised optimization from `parameters.csv`
+- `estimate()` for inverse parameter estimation from `parameters.csv` and `variables.csv`
+- `load(metadata_path)` to reload a trained run
+- `predict(x)` to infer variables for new parameter values
+
+## Packaging
+
+The PyPI project name is `kkt-hardnet`, while the import path remains `kkthn`. The default install is CPU-oriented; CUDA support is selected explicitly with the `cuda12` extra instead of automatic device detection.
+
+# Modeling Workflow
+
+Use `KKTHardNet` to define symbolic constrained problems with:
+
+- named parameters, decision variables, and optional inverse parameters
+- equality and inequality constraints written as Python expressions
+- optional objectives for optimization tasks
+- hard KKT projection during training and inference
+- saved run artifacts including weights, predictions, history, summary, and metadata
+
+Available methods:
+
+- `dataset(parameters=..., variables=...)` to attach training data
+- `model()` for supervised surrogate learning
+- `optimize()` for unsupervised optimization
+- `estimate()` for inverse parameter estimation
+- `load(metadata_path)` to reload a trained run
+- `predict(x)` to evaluate a trained or loaded model on new inputs
+
+The example below shows one complete workflow. The CSV column names must match the declared parameter and variable names. Each run writes `<model_name>_<timestamp>/` with `parameters.csv`, optional `variables.csv`, `history.csv`, `predictions.csv`, `model_weights.npz`, `summary.json`, and `metadata.json`.
+
+```python
+from kkthn import KKTHardNet
+
+TRAIN = {
+    "epochs": 1200,
+    "batch_size": 32,
+    "learning_rate": 1e-3,
+    "train_frac": 0.8,
+    "hidden_size": 64,
+    "hidden_layers": 2,
+    "seed": 42,
+    "dtype": "float64",
+    "print_every": 1,
+    "newton_step_length": 0.5,
+    "newton_tol": 1e-6,
+    "newton_reg_factor": 1e-3,
+    "max_newton_iter": 30,
+    "max_backtrack_iter": 10,
+}
+
+# Build a symbolic problem.
+model = KKTHardNet(name="demo_model", train=TRAIN)
+x = model.add_parameter(["x1", "x2"])
+theta = model.add_inverse_parameter(["a0", "a1"], init_value=[10.0, -10.0])
+y = model.add_variable(["y1", "y2", "y3"])
+
+# Objectives are optional for surrogate modeling and inverse estimation,
+# but required for optimize().
+model.objective = 0.5 * (y.y1**2 + y.y2**2 + y.y3**2)
+model.constraints.add(
+    theta.a0 * y.y1 + y.y2 - x.x1 == 0,
+    y.y2 - theta.a1 * y.y3 - x.x2 == 0,
+    y.y1**2 + y.y3**2 <= 2.0,
+    y.y1 >= 0,
+)
+
+# Attach data.
+# For model() or estimate(): provide both parameters and variables.
+# For optimize(): provide only parameters.
+model.dataset(parameters="parameters.csv", variables="variables.csv")
+
+# Choose one training mode.
+surrogate_result = model.model()
+# optimize_result = model.optimize()
+# estimate_result = model.estimate()
+
+# Reload a saved run later and predict on a new parameter vector.
+reloaded = KKTHardNet()
+reloaded.load("demo_model_20260414_120000/metadata.json")
+prediction = reloaded.predict([1.0, 2.0])
+```
+⚠️ Please cite our work if you use this code in your research.
+Citation formats are provided below.
+
+arXiv Preprint: https://arxiv.org/pdf/2507.08124
+Journal: https://doi.org/10.1016/j.compchemeng.2025.109418
+
+```bibtex
+@article{iftakher2025physics,
+  title={Physics-informed neural networks with hard nonlinear equality and inequality constraints},
+  author={Iftakher, Ashfaq and Golder, Rahul and Nath Roy, Bimol and Hasan, MM Faruque},
+  journal={Computers \& Chemical Engineering},
+  pages={109418},
+  year={2025},
+  publisher={Elsevier}
+}

kkt_hardnet-1.0.1/pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "kkt-hardnet"
+version = "1.0.1"
+description = "KKT-HardNet: Physics-constrained neural networks with hard nonlinear equality and inequality constraint satisfaction guarantees"
+readme = "README.md"
+requires-python = ">=3.9"
+license = {text = "MIT"}
+authors = [{name = "Rahul Golder", email = "rahulgolder8420@tamu.edu"},
+    {name = "Bimol Nath Roy", email = "bimolnathroy@tamu.edu"},
+    {name = "Ashfaq Iftakher", email = "iftakher@tamu.edu"},
+    {name = "MM Faruque Hasan", email = "hasan@tamu.edu"}
+]
+dependencies = [
+    "jax==0.4.30",
+    "jaxlib==0.4.30",
+    "flax==0.8.5",
+    "optax==0.2.4",
+    "chex==0.1.90",
+    "numpy==1.26.4",
+    "scipy==1.13.1",
+    "pandas==2.2.3",
+    "matplotlib==3.9.4",
+]
+
+[project.optional-dependencies]
+cuda12 = [
+    "jax[cuda12]==0.4.30",
+    "jaxlib==0.4.30",
+]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]

kkt_hardnet-1.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

kkt_hardnet-1.0.1/src/jaxmodel/__init__.py

@@ -0,0 +1,52 @@
+from .variables import VariableBuilder, VariableSpec
+from .parameters import ParameterSpec
+from .bounds import BoundSpec
+from .builder import NLPBuilder
+from .highlevel import HighLevelNLPBuilder
+from .objective import QuadraticObjective
+from .constraints import (
+    ConstraintEntry,
+    eq_scalar,
+    ineq_scalar,
+    eq_block,
+    ineq_block,
+    quadratic_eq_scalar,
+    quadratic_ineq_scalar,
+)
+from .blocks import (
+    affine_eq_from_parts,
+    affine_ineq_from_parts,
+    nonlinear_eq_from_parts,
+    nonlinear_ineq_from_parts,
+)
+from .approximations import (
+    LinearizationData,
+    QuadraticObjectiveData,
+    SQPSubproblemData,
+)
+from .model import JaxNLPModel
+
+__all__ = [
+    "VariableBuilder",
+    "VariableSpec",
+    "ParameterSpec",
+    "BoundSpec",
+    "NLPBuilder",
+    "HighLevelNLPBuilder",
+    "QuadraticObjective",
+    "ConstraintEntry",
+    "eq_scalar",
+    "ineq_scalar",
+    "eq_block",
+    "ineq_block",
+    "quadratic_eq_scalar",
+    "quadratic_ineq_scalar",
+    "affine_eq_from_parts",
+    "affine_ineq_from_parts",
+    "nonlinear_eq_from_parts",
+    "nonlinear_ineq_from_parts",
+    "LinearizationData",
+    "QuadraticObjectiveData",
+    "SQPSubproblemData",
+    "JaxNLPModel",
+]

kkt_hardnet-1.0.1/src/jaxmodel/approximations.py

@@ -0,0 +1,181 @@
+import jax.numpy as jnp
+from typing import NamedTuple, Optional
+
+
+class LinearizationData(NamedTuple):
+    A: jnp.ndarray
+    b: jnp.ndarray
+    C: jnp.ndarray
+    d: jnp.ndarray
+    h_val: jnp.ndarray
+    g_val: jnp.ndarray
+
+
+class QuadraticObjectiveData(NamedTuple):
+    grad_f: jnp.ndarray
+    Q: jnp.ndarray
+    Q_diag: jnp.ndarray
+    c: jnp.ndarray
+
+
+class SQPSubproblemData(NamedTuple):
+    objective: QuadraticObjectiveData
+    constraints: LinearizationData
+    l: Optional[jnp.ndarray]
+    u: Optional[jnp.ndarray]
+
+
+def linearize_constraints_data(eq_fun, ineq_fun, jac_eq, jac_ineq, params, y):
+    A = jac_eq(y, params)
+    h = eq_fun(params, y)
+    b = A @ y - h
+
+    C = jac_ineq(y, params)
+    g = ineq_fun(params, y)
+    d = C @ y - g
+
+    return LinearizationData(A=A, b=b, C=C, d=d, h_val=h, g_val=g)
+
+
+def linearize_constraints(eq_fun, ineq_fun, jac_eq, jac_ineq, params, y):
+    data = linearize_constraints_data(
+        eq_fun=eq_fun,
+        ineq_fun=ineq_fun,
+        jac_eq=jac_eq,
+        jac_ineq=jac_ineq,
+        params=params,
+        y=y,
+    )
+    return dict(data._asdict())
+
+
+def quadraticize_objective_data(
+    grad_fun,
+    hess_fun,
+    diag_hess_fun,
+    params,
+    y,
+    rho: float = 1.0,
+    use_diagonal_hessian: bool = True,
+    diag_floor: Optional[float] = None,
+):
+    grad = grad_fun(y, params)
+
+    if use_diagonal_hessian:
+        qdiag = diag_hess_fun(params, y)
+        if diag_floor is not None:
+            qdiag = jnp.maximum(qdiag, diag_floor)
+        Q = jnp.diag(rho * qdiag)
+    else:
+        Q = rho * hess_fun(y, params)
+
+    c = grad - Q @ y
+
+    return QuadraticObjectiveData(grad_f=grad, Q=Q, Q_diag=jnp.diag(Q), c=c)
+
+
+def quadraticize_objective(
+    grad_fun,
+    hess_fun,
+    diag_hess_fun,
+    params,
+    y,
+    rho: float = 1.0,
+    use_diagonal_hessian: bool = True,
+    diag_floor: Optional[float] = None,
+):
+    data = quadraticize_objective_data(
+        grad_fun=grad_fun,
+        hess_fun=hess_fun,
+        diag_hess_fun=diag_hess_fun,
+        params=params,
+        y=y,
+        rho=rho,
+        use_diagonal_hessian=use_diagonal_hessian,
+        diag_floor=diag_floor,
+    )
+    return dict(data._asdict())
+
+
+def build_sqp_subproblem_data(
+    eq_fun,
+    ineq_fun,
+    jac_eq,
+    jac_ineq,
+    grad_fun,
+    hess_fun,
+    diag_hess_fun,
+    lower_fun,
+    upper_fun,
+    params,
+    y,
+    rho: float = 1.0,
+    use_diagonal_hessian: bool = True,
+    diag_floor: Optional[float] = None,
+):
+    qdata = quadraticize_objective_data(
+        grad_fun=grad_fun,
+        hess_fun=hess_fun,
+        diag_hess_fun=diag_hess_fun,
+        params=params,
+        y=y,
+        rho=rho,
+        use_diagonal_hessian=use_diagonal_hessian,
+        diag_floor=diag_floor,
+    )
+
+    cdata = linearize_constraints_data(
+        eq_fun=eq_fun,
+        ineq_fun=ineq_fun,
+        jac_eq=jac_eq,
+        jac_ineq=jac_ineq,
+        params=params,
+        y=y,
+    )
+
+    return SQPSubproblemData(
+        objective=qdata,
+        constraints=cdata,
+        l=None if lower_fun is None else lower_fun(params),
+        u=None if upper_fun is None else upper_fun(params),
+    )
+
+
+def build_sqp_data(
+    eq_fun,
+    ineq_fun,
+    jac_eq,
+    jac_ineq,
+    grad_fun,
+    hess_fun,
+    diag_hess_fun,
+    lower_fun,
+    upper_fun,
+    params,
+    y,
+    rho: float = 1.0,
+    use_diagonal_hessian: bool = True,
+    diag_floor: Optional[float] = None,
+):
+    data = build_sqp_subproblem_data(
+        eq_fun=eq_fun,
+        ineq_fun=ineq_fun,
+        jac_eq=jac_eq,
+        jac_ineq=jac_ineq,
+        grad_fun=grad_fun,
+        hess_fun=hess_fun,
+        diag_hess_fun=diag_hess_fun,
+        lower_fun=lower_fun,
+        upper_fun=upper_fun,
+        params=params,
+        y=y,
+        rho=rho,
+        use_diagonal_hessian=use_diagonal_hessian,
+        diag_floor=diag_floor,
+    )
+    return {
+        **dict(data.objective._asdict()),
+        **dict(data.constraints._asdict()),
+        "l": data.l,
+        "u": data.u,
+    }

kkt_hardnet-1.0.1/src/jaxmodel/autodiff.py

@@ -0,0 +1,24 @@
+import jax
+import jax.numpy as jnp
+
+
+def build_grad_y_objective(obj_flat):
+    return jax.grad(lambda y, p: obj_flat(p, y), argnums=0)
+
+
+def build_hess_y_objective(obj_flat):
+    return jax.hessian(lambda y, p: obj_flat(p, y), argnums=0)
+
+
+def build_diag_hess_y_objective(grad_fun):
+    def diag_hess(params, y):
+        n = y.shape[0]
+        basis = jnp.eye(n, dtype=y.dtype)
+        _, jvp_fun = jax.linearize(lambda yy: grad_fun(yy, params), y)
+        cols = jax.vmap(jvp_fun)(basis)
+        return jnp.diag(cols.T)
+    return diag_hess
+
+
+def build_jac_y_vector_fun(vec_fun):
+    return jax.jacrev(lambda y, p: vec_fun(p, y), argnums=0)