1 files changed, 101 insertions, 0 deletions
diff --git a/sumofsquares/init/initenergyfunction.py b/sumofsquares/init/initenergyfunction.py
new file mode 100644
index 0000000..8b2f1bf
--- /dev/null
+++ b/sumofsquares/init/initenergyfunction.py
@@ -0,0 +1,101 @@
+import polymatrix
+
+from sumofsquares.abc.sosexpr import SOSExpr
+from sumofsquares.impl.energyfunctionimpl import EnergyFunctionImpl
+from sumofsquares.init.initsosexpr import init_param_expr, init_param_expr_from_reference, init_sos_expr
+
+
+def init_energy_function(
+    name: str,
+    h: SOSExpr,
+    x_dot: SOSExpr,
+    positive_b: bool,
+    epsilon_min: SOSExpr | None = None,
+    roi: SOSExpr | polymatrix.Expression | None = None,
+    decrease_rate: SOSExpr | polymatrix.Expression | None = None,
+    w_outer: SOSExpr | polymatrix.Expression | None = None,
+    w_inner: SOSExpr | polymatrix.Expression | None = None,
+):
+    sos_constraints = tuple()
+    
+    nh = h.diff().T
+    h_dot = (nh.T @ x_dot).cache()
+
+    gamma_b = init_param_expr_from_reference(
+        name=f'gamma_b_{name}',
+        reference=h_dot, 
+        multiplicand=h,
+    )
+    if positive_b:
+        sos_constraints += gamma_b.sos_constraints
+
+    energy_expr = -h_dot
+    energy_expr -= gamma_b * h
+
+    if epsilon_min is None:
+        epsilon = None
+
+    else:
+        epsilon = init_param_expr(
+            name=f'epsilon_{name}',
+            variables=h.variables,
+        )
+
+        sos_constraints += (epsilon - epsilon_min).sos_constraints
+        energy_expr += epsilon
+
+    if roi is None:
+        gamma_roi = None,
+
+    else:
+        gamma_roi = init_param_expr_from_reference(
+            name=f'gamma_roi_{name}',
+            reference=h_dot, 
+            multiplicand=roi,
+        )
+        sos_constraints += gamma_roi.sos_constraints
+
+        energy_expr += gamma_roi * roi
+
+    if decrease_rate is not None:    
+        energy_expr -= decrease_rate
+
+    sos_constraints += energy_expr.sos_constraints
+
+    if w_outer is None:
+        gamma_w_outer = None
+
+    else:
+        gamma_w_outer = init_param_expr_from_reference(
+            name=f'gamma_w_outer_{name}',
+            reference=h, 
+            multiplicand=w_outer,
+        )
+
+        sos_constraints += gamma_w_outer.sos_constraints
+        sos_constraints += (h - gamma_w_outer * w_outer).sos_constraints
+
+    if w_inner is None:
+        gamma_w_inner = None
+
+    else:
+        gamma_w_inner = init_param_expr_from_reference(
+            name=f'gamma_w_inner_{name}',
+            reference=h, 
+            multiplicand=w_inner,
+        )
+
+        sos_constraints += gamma_w_inner.sos_constraints
+        sos_constraints += (gamma_w_inner * w_inner - h).sos_constraints
+
+    return EnergyFunctionImpl(
+        name=name,
+        h=h,
+        h_dot=h_dot,
+        epsilon=epsilon,
+        gamma_roi=gamma_roi,
+        gamma_b=gamma_b,
+        gamma_w_outer=gamma_w_outer,
+        gamma_w_inner=gamma_w_inner,
+        sos_constraints=sos_constraints,
+    )