Port to jax

README.rst

@@ -70,6 +70,12 @@ You should then be able to run the examples and tutorials in the examples folder
 .. _PyPI: https://pypi.org/project/SiPANN/
 
 
+JAX support
+===========
+
+SiPANN provides optional support for [JAX](https://github.com/google/jax). You need to have JAX installed in order to let SiPANN use `jax.numpy`(which is faster) instead of `numpy`. For installation instructions, see the [JAX documentation](https://github.com/google/jax#installation).
+
+
 References
 ==========
 

SiPANN/comp.py

@@ -2,7 +2,10 @@
 from abc import ABC, abstractmethod
 
 import gdspy
-import numpy as np
+try:
+    import jax.numpy as np
+except ImportError:
+    import numpy as np
 import pkg_resources
 import scipy.integrate as integrate
 import scipy.special as special

SiPANN/import_nn.py

@@ -1,7 +1,10 @@
 import pickle
 from itertools import combinations_with_replacement as comb_w_r
 
-import numpy as np
+try:
+    import jax.numpy as np
+except ImportError:
+    import numpy as np
 import tensorflow as tf
 
 

SiPANN/nn.py

@@ -17,7 +17,10 @@
 # ---------------------------------------------------------------------------- #
 # Import libraries
 # ---------------------------------------------------------------------------- #
-import numpy as np
+try:
+    import jax.numpy as np
+except ImportError:
+    import numpy as np
 import pkg_resources
 import skrf as rf
 from scipy.interpolate import UnivariateSpline

SiPANN/scee.py

@@ -2,7 +2,10 @@
 from abc import ABC, abstractmethod
 
 import gdspy
-import numpy as np
+try:
+    import jax.numpy as np
+except ImportError:
+    import numpy as np
 import pkg_resources
 from scipy import special
 from scipy.integrate import quad

SiPANN/scee_int.py

@@ -1,4 +1,7 @@
-import numpy as np
+try:
+    import jax.numpy as np
+except ImportError:
+    import numpy as np
 from simphony import Model
 from simphony.layout import Circuit
 from simphony.pins import Pin, PinList

SiPANN/scee_opt.py

@@ -2,7 +2,10 @@
 import os
 
 import matplotlib.pyplot as plt
-import numpy as np
+try:
+    import jax.numpy as np
+except ImportError:
+    import numpy as np
 from numba import njit, vectorize
 from numba.extending import get_cython_function_address
 from scipy import special