From 0f4236c0428a0e825a611f8612531b70d80311d0 Mon Sep 17 00:00:00 2001
From: Mansa Krishna <mansakrishna23@gmail.com>
Date: Wed, 3 Jul 2024 13:08:11 -0400
Subject: [PATCH] surface slopes additional loss added

---
 pinnicle/physics/physics.py | 18 +++++++++++++++++-
 pinnicle/pinn.py            |  4 ++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/pinnicle/physics/physics.py b/pinnicle/physics/physics.py
index 284745e..8d5748d 100644
--- a/pinnicle/physics/physics.py
+++ b/pinnicle/physics/physics.py
@@ -1,7 +1,7 @@
 from ..parameter import PhysicsParameter
 from . import EquationBase
 import itertools
-
+import deepxde as dde
 
 class Physics:
     """ All the physics in used as constraint in the PINN
@@ -65,6 +65,22 @@ def vel_mag(self, nn_input_var, nn_output_var, X):
         vel = (nn_output_var[:,uid:uid+1]**2.0 + nn_output_var[:,vid:vid+1]**2.0) ** 0.5
         return vel
 
+    def surf_x(self, nn_input_var, nn_output_var, X):
+        """dsdx
+        """
+        sid = self.output_var.index('s')
+        xid = self.input_var.index('x')
+        dsdx = dde.grad.jacobian(nn_output_var, nn_input_var, i=sid, j=xid)
+        return dsdx
+
+    def surf_y(self, nn_input_var, nn_output_var, X):
+        """dsdy
+        """
+        sid = self.output_var.index('s')
+        yid = self.input_var.index('y')
+        dsdy = dde.grad.jacobian(nn_output_var, nn_input_var, i=sid, j=yid)
+        return dsdy
+
     def operator(self, pname):
         """ grab the pde operator
 
diff --git a/pinnicle/pinn.py b/pinnicle/pinn.py
index 43a9503..d70f4d9 100644
--- a/pinnicle/pinn.py
+++ b/pinnicle/pinn.py
@@ -233,6 +233,10 @@ def update_training_data(self, training_data):
                     # currently, only implement 'vel'
                     elif d == "vel":
                         training_temp.append(dde.icbc.PointSetOperatorBC(training_data.X[d], training_data.sol[d], self.physics.vel_mag))
+                    elif d == "sx":
+                        training_temp.append(dde.icbc.PointSetOperatorBC(training_data.X[d], training_data.sol[d], self.physics.surf_x))
+                    elif d == "sy":
+                        training_temp.append(dde.icbc.PointSetOperatorBC(training_data.X[d], training_data.sol[d], self.physics.surf_y))
                     else:
                         raise ValueError(f"{d} is not found in the output_variable of the nn, and not defined")