add pow support for jax

Author: enigne

pinnicle/physics/physics.py

@@ -3,7 +3,7 @@
 from ..parameter import PhysicsParameter
 from . import EquationBase
 import itertools
-from ..utils import slice_column, jacobian
+from ..utils import slice_column, jacobian, ppow
 
 class Physics:
     """ All the physics in used as constraint in the PINN
@@ -80,7 +80,7 @@ def vel_mag(self, nn_input_var, nn_output_var, X):
         vid = self.output_var.index('v')
         u = slice_column(nn_output_var, uid)
         v = slice_column(nn_output_var, vid) 
-        vel = bkd.pow((bkd.square(u) + bkd.square(v) + 1.0e-30), 0.5)
+        vel = ppow((bkd.square(u) + bkd.square(v) + 1.0e-30), 0.5)
         return vel
 
     def surf_x(self, nn_input_var, nn_output_var, X):

pinnicle/utils/__init__.py

@@ -1,7 +1,7 @@
 from .helper import *
-from .backends_specified import slice_column, jacobian, slice_function_jax, matmul
+from .backends_specified import slice_column, jacobian, slice_function_jax, matmul, ppow
 from .history import History
 from .data_misfit import get
 from .plotting import plot_solutions, plot_dict_data, plot_data, plot_nn, tripcolor_similarity, tripcolor_residuals, diffplot, resplot, plot_tracks
 from .plotmodel import *
-from .data_interpolation import *
+from .data_interpolation import *

pinnicle/utils/backends_specified.py

@@ -45,9 +45,11 @@ def matmul_jax(A, B):
     slice_column = slice_column_tf
     jacobian = jacobian_tf
     matmul = matmul_tf
+    ppow = bkd.pow
 elif backend_name == "jax":
     slice_column = slice_column_jax
     jacobian = jacobian_jax
     matmul = matmul_jax
+    ppow = jax.numpy.pow
 else:
     raise ValueError(f"{backend_name} is not supported by PINNICLE")

tests/test_utils.py

@@ -83,6 +83,12 @@ def test_slice_column():
     assert c.shape == (1,)
     assert c[0] == 2
 
+def test_ppow():
+    a = backend.as_tensor([2.0])
+    c = pinnicle.utils.backends_specified.ppow(a, 2.0)
+    assert c == 4.0
+
+
 def test_interpfrombedmachine():
     x = np.array([300025,301025,302025])
     y = np.array([-2579975, -2578975, -2577975])