Merge pull request #28 from btalamini/feature/move_jax_settings_to_init

Feature/specify global jax settings in a more standard way

Author: ralberd

examples/arch_bc/ArchBc.py

@@ -1,6 +1,5 @@
 from matplotlib import pyplot as plt
 
-from optimism.JaxConfig import *
 from optimism import EquationSolver as EqSolver
 from optimism import FunctionSpace
 from optimism import Interpolants

examples/arch_bc/ArchTraction.py

@@ -1,6 +1,5 @@
 from matplotlib import pyplot as plt
 
-from optimism.JaxConfig import *
 from optimism import EquationSolver as EqSolver
 from optimism import EquationSolverSubspace as SolverSubspace
 from optimism import FunctionSpace

examples/buckle_slide/Sliding.py

@@ -1,4 +1,6 @@
-from optimism.JaxConfig import *
+from functools import partial
+from jax import jit
+from jax import numpy as np
 
 from optimism import EquationSolver as EqSolver
 from optimism import FunctionSpace
@@ -20,9 +22,6 @@
 
 from optimism.test.MeshFixture import MeshFixture
 
-import os
-os.environ['XLA_FLAGS'] = "--xla_force_host_platform_device_count=2"
-
 props = {'elastic modulus': 10.0,
          'poisson ratio': 0.3}
 

examples/contact_mesh_on_mesh/TwoBodyContact.py

@@ -1,8 +1,7 @@
 from matplotlib import pyplot as plt
+from jax import numpy as np
 
-from optimism.JaxConfig import *
 from optimism import EquationSolver as EqSolver
-
 from optimism import Mesh
 from optimism import Mechanics
 from optimism import FunctionSpace

examples/friction/CornerSlide.py

@@ -1,5 +1,6 @@
-from optimism.JaxConfig import *
-
+from functools import partial
+import jax
+from jax import numpy as np
 
 from optimism.material import Neohookean
 from optimism import Mechanics
@@ -112,7 +113,7 @@ def constraint_func(Uu, p):
         self.constraint_func = constraint_func
 
 
-    @partial(jit, static_argnums=0)
+    @partial(jax.jit, static_argnums=0)
     def create_field(self, Uu, p):
         return self.dofManager.create_field(Uu, self.get_ubcs(p))
 
@@ -150,7 +151,7 @@ def run(self):
         penalty = 2.5
         kappa0 = penalty * np.ones(lam0.shape)
 
-        hess_func = jit(hessian(lambda Uu,p: self.energy_func(Uu,lam0,p)))
+        hess_func = jax.jit(jax.hessian(lambda Uu,p: self.energy_func(Uu,lam0,p)))
 
         objective = ConstrainedQuasiObjective(self.energy_func,
                                               self.constraint_func,

examples/hemisphere_cap/hemisphere_plastic_disp_control.py

@@ -1,4 +1,6 @@
-from optimism.JaxConfig import *
+import jax
+from jax import numpy as np
+
 from optimism import EquationSolver as EqSolver
 from optimism import FunctionSpace
 from optimism import Interpolants
@@ -55,7 +57,7 @@ def compute_energy_from_bcs(Uu, Ubc, p):
             strainEnergy = self.bvpFuncs.compute_strain_energy(U, internalVariables)
             return strainEnergy
 
-        self.compute_bc_reactions = jit(grad(compute_energy_from_bcs, 1))
+        self.compute_bc_reactions = jax.jit(jax.grad(compute_energy_from_bcs, 1))
 
         self.trSettings = EqSolver.get_settings(max_trust_iters=400, t1=0.4, t2=1.5, eta1=1e-8, eta2=0.2, eta3=0.8, over_iters=100)
 

examples/hemisphere_cap/hemisphere_plastic_load_control.py

@@ -1,4 +1,6 @@
-from optimism.JaxConfig import *
+import jax
+from jax import numpy as np
+
 from optimism import EquationSolver as EqSolver
 from optimism import FunctionSpace
 from optimism import Interpolants
@@ -31,11 +33,11 @@ def integrate_field_function_on_edge(quadratureRule, edge, mesh, func):
     jac = np.linalg.norm(edgeCoords[0,:] - edgeCoords[1,:])
     XGauss = edgeCoords[0] + np.outer(xigauss, edgeCoords[1] - edgeCoords[0])
     dX = jac*wgauss
-    return np.dot(vmap(func)(XGauss), dX)
+    return np.dot(jax.vmap(func)(XGauss), dX)
 
 
 def integrate_field_function_on_surface(quadratureRule, edges, mesh, func):
-    F = vmap(integrate_field_function_on_edge, (None,0,None,None))
+    F = jax.vmap(integrate_field_function_on_edge, (None,0,None,None))
     return np.sum(F(quadratureRule, edges, mesh, func))
 
 
@@ -94,7 +96,7 @@ def compute_energy_from_bcs(Uu, Ubc, p):
             loadPotential = TractionBC.compute_traction_potential_energy(self.mesh, U, self.lineQuadRule, self.mesh.sideSets['push'], lambda X: np.array([0.0, -F/self.pushArea]))
             return strainEnergy + loadPotential
 
-        self.compute_bc_reactions = jit(grad(compute_energy_from_bcs, 1))
+        self.compute_bc_reactions = jax.jit(jax.grad(compute_energy_from_bcs, 1))
 
         self.trSettings = EqSolver.get_settings(max_trust_iters=400, t1=0.4, t2=1.5, eta1=1e-8, eta2=0.2, eta3=0.8, over_iters=100)
 

examples/inverse/BuckleInverse.py

@@ -1,4 +1,5 @@
-from optimism.JaxConfig import *
+import jax
+from jax import numpy as np
 from jax import custom_jvp, custom_vjp
 
 from optimism import ReadMesh
@@ -197,7 +198,7 @@ def run(self):
         outputDisp = [0,]
         outputForce = [0,]
 
-        reaction_func = jit( lambda x,p: grad(self.energy_func,1)(x,p)[0][1] )
+        reaction_func = jax.jit( lambda x,p: jax.grad(self.energy_func,1)(x,p)[0][1] )
 
         N = 1
 
@@ -254,7 +255,7 @@ def loss(chi):
         opt_state = opt_init(chi)
 
         def step(step, opt_state):
-            value, grads = value_and_grad(loss)(get_params(opt_state))
+            value, grads = jax.value_and_grad(loss)(get_params(opt_state))
             opt_state = opt_update(step, grads, opt_state)
             return value, opt_state
 

examples/inverse/arch/ArchInverse.py

@@ -1,4 +1,5 @@
-from optimism.JaxConfig import *
+import jax
+from jax import numpy as np
 
 from optimism import ReadMesh
 from optimism.material import Neohookean
@@ -105,7 +106,7 @@ def energy_func(self, Uu, p):
 
 
     def reaction_func(self, Uu, p):
-        return grad(self.energy_func,1)(Uu,p)[0]
+        return jax.grad(self.energy_func,1)(Uu,p)[0]
 
 
     def compute_volume(self, p):
@@ -226,7 +227,7 @@ def debug_loss(chi):
         opt_state = opt_init(chi)
 
         def step(step, opt_state):
-            value, grads = value_and_grad(loss)(get_params(opt_state))
+            value, grads = jax.value_and_grad(loss)(get_params(opt_state))
             opt_state = opt_update(step, grads, opt_state)
             return value, opt_state
 

examples/material_test/MaterialTest.py

@@ -1,8 +1,8 @@
 """Demonstrate the material testing tools."""
 
+from jax import numpy as np
 from matplotlib import pyplot as plt
 
-from optimism.JaxConfig import *
 from optimism.material import MaterialUniaxialSimulator
 from optimism.material import J2Plastic
 

examples/sharp_notch_fracture/SharpNotchFracture.py

@@ -1,6 +1,7 @@
+from functools import partial
 import jax
+from jax import numpy as np
 
-from optimism.JaxConfig import *
 from optimism import AlSolver
 from optimism import BoundConstrainedSolver
 from optimism import BoundConstrainedObjective
@@ -102,7 +103,7 @@ def energy_for_rxns(Ubc, Uu, p):
             internalVariables = p[1]
             return self.bvpFunctions.compute_internal_energy(U, internalVariables)
 
-        self.compute_reactions = jit(grad(energy_for_rxns))
+        self.compute_reactions = jax.jit(jax.grad(energy_for_rxns))
 
 
     def plot_solution(self, U, p, lagrange, plotName):

examples/surfing_fracture/SurfingFracture.py

@@ -1,8 +1,9 @@
+from functools import partial
+import jax
+from jax import numpy as np
 from matplotlib import pyplot as plt
-from scipy.sparse import diags as sparse_diags
 import numpy as onp
-
-from optimism.JaxConfig import *
+from scipy.sparse import diags as sparse_diags
 
 from optimism import BoundConstrainedObjective
 from optimism import ConstrainedObjective
@@ -64,7 +65,7 @@ def apply_mode_I_field_at_point(X, K_I):
 
 
 def apply_mode_I_Bc(boundaryNodeCoords, K_I):
-    return vmap(apply_mode_I_field_at_point, (0, None))(boundaryNodeCoords, K_I)
+    return jax.vmap(apply_mode_I_field_at_point, (0, None))(boundaryNodeCoords, K_I)
 
 
 crackDirection = np.array([1.0,0.0])
@@ -84,7 +85,7 @@ def J_integral(U, internals, mesh, fs, edges, bvpFuncs):
     stresses = FunctionSpace.project_quadrature_field_to_element_field(fs, stresses)
     Ws = FunctionSpace.project_quadrature_field_to_element_field(fs, Ws)
 
-    computeJs = vmap(compute_J_integral_on_edge, (None, 0, 0, 0, 0))
+    computeJs = jax.vmap(compute_J_integral_on_edge, (None, 0, 0, 0, 0))
     return np.sum(computeJs(mesh,
                             edges,
                             Ws[edges[:,0]],
@@ -244,8 +245,7 @@ def get_ubcs(self, p):
         Xb = self.mesh.coords[self.mesh.nodeSets['external'],:]
         Xb = Xb.at[:,0].add(-origin)
         modeIBcs = apply_mode_I_Bc(Xb, KI)
-        index = (self.mesh.nodeSets['external'],:2)
-        V = V.at[index].set(modeIBcs)
+        V = V.at[self.mesh.nodeSets['external'],:2].set(modeIBcs)
 
         return self.dofManager.get_bc_values(V)
 

examples/tension_axisymmetric/TensionAxisymmetric.py

@@ -1,6 +1,8 @@
+from functools import partial
+import jax
+from jax import numpy as np
 from matplotlib import pyplot as plt
 
-from optimism.JaxConfig import *
 from optimism import EquationSolver as EqSolver
 from optimism import FunctionSpace
 from optimism import Interpolants
@@ -84,8 +86,8 @@ def energy_function(self, Uu, p):
         return self.mechanicsFunctions.compute_strain_energy(U, internalVariables)
 
 
-    @partial(jit, static_argnums=0)
-    @partial(value_and_grad, argnums=2)
+    @partial(jax.jit, static_argnums=0)
+    @partial(jax.value_and_grad, argnums=2)
     def compute_reactions_from_bcs(self, Uu, Ubc, internalVariables):
         U = self.dofManager.create_field(Uu, Ubc)
         return self.mechanicsFunctions.compute_strain_energy(U, internalVariables)

examples/tension_axisymmetric/TensionAxisymmetricFracture.py

@@ -1,6 +1,8 @@
+from functools import partial
+import jax
+from jax import numpy as np
 from scipy.sparse import diags as sp_sparse_diags
 
-from optimism.JaxConfig import *
 from optimism import AlSolver
 from optimism import BoundConstrainedSolver
 from optimism import BoundConstrainedObjective
@@ -39,7 +41,7 @@
 
 def compute_row_sum_gram_matrix(fs, dummyU, dummyInternalVars):
     func = lambda u, uGrad, q, x: u
-    compute = grad(FunctionSpace.integrate_over_block, 1)
+    compute = jax.grad(FunctionSpace.integrate_over_block, 1)
     return compute(fs, dummyU, dummyInternalVars, func, fs.mesh.blocks['block_1'])
 
 
@@ -112,8 +114,8 @@ def energy_function(self, Uu, p):
         return self.bvpFunctions.compute_internal_energy(U, internalVariables)
 
 
-    @partial(jit, static_argnums=0)
-    @partial(value_and_grad, argnums=2)
+    @partial(jax.jit, static_argnums=0)
+    @partial(jax.value_and_grad, argnums=2)
     def compute_reactions_from_bcs(self, Uu, Ubc, internalVariables):
         U = self.dofManager.create_field(Uu, Ubc)
         return self.bvpFunctions.compute_internal_energy(U, internalVariables)

examples/tension_axisymmetric/materialPointSimulation.py

@@ -1,8 +1,7 @@
 from matplotlib import pyplot as plt
 
-from optimism.JaxConfig import *
 from optimism import TensorMath
-from optimism.phasefield import SandiaModel as MaterialModel
+from optimism.phasefield import PhaseFieldLorentzPlastic as MaterialModel
 from optimism.phasefield.MaterialPointSimulator import MaterialPointSimulator
 
 from properties import props

examples/uniaxial/Uniaxial.py

@@ -1,6 +1,8 @@
+from functools import partial
+import jax
+from jax import numpy as np
 from matplotlib import pyplot as plt
 
-from optimism.JaxConfig import *
 from optimism import EquationSolver as EqSolver
 from optimism import FunctionSpace
 from optimism import Interpolants
@@ -82,8 +84,8 @@ def energy_function(self, Uu, p):
         return self.mechanicsFunctions.compute_strain_energy(U, internalVariables)
 
 
-    @partial(jit, static_argnums=0)
-    @partial(value_and_grad, argnums=2)
+    @partial(jax.jit, static_argnums=0)
+    @partial(jax.value_and_grad, argnums=2)
     def compute_reactions_from_bcs(self, Uu, Ubc, internalVariables):
         U = self.dofManager.create_field(Uu, Ubc)
         return self.mechanicsFunctions.compute_strain_energy(U, internalVariables)

examples/uniaxial_dynamic/UniaxialDynamic.py

@@ -1,4 +1,4 @@
-#import sys
+from jax import numpy as np
 
 from optimism.JaxConfig import *
 from optimism import EquationSolver
@@ -94,13 +94,6 @@ def energy_function(self, Uu, p):
         return self.dynamicsFunctions.compute_algorithmic_energy(U, UPredictor, internalVariables, dt)
 
 
-    # @partial(jit, static_argnums=0)
-    # @partial(value_and_grad, argnums=2)
-    # def compute_reactions_from_bcs(self, Uu, Ubc, internalVariables):
-    #     U = self.dofManager.create_field(Uu, Ubc)
-    #     return self.dynamicsFunctions.compute_output_potential_energies_and_stresses(U, internalVariables)
-
-
     def create_field(self, Uu, p):
         return self.dofManager.create_field(Uu, self.get_ubcs(p))
 

optimism/JaxConfig.py

@@ -1,37 +1,9 @@
 from functools import partial
 from collections import namedtuple
-import warnings
-with warnings.catch_warnings():
-    warnings.filterwarnings("ignore")
-    import jax.numpy as np
-    from jax import grad, jacrev, jacfwd, jvp, vjp, value_and_grad, hessian, ops, lax, make_jaxpr, linearize
-    from jax import custom_jvp, custom_vjp
-    from jax import jit as jaxJit
-    from jax import vmap as jaxVmap
-    from jax.lax import while_loop
-    from jax.config import config
-
-    
-config.update("jax_enable_x64", True)
-#config.update("jax_debug_nans", True)
-
-
-jaxDebug=False
-
-if jaxDebug:
-    def jit(f,static_argnums=None):
-        return f
-    
-    vmap = jaxVmap
-
-    def while_loop(cond_fun, body_fun, init_val):
-        val = init_val
-        while cond_fun(val):
-            val = body_fun(val)
-        return val
-else:
-    jit = jaxJit
-    vmap = jaxVmap
+import jax.numpy as np
+from jax import grad, jacrev, jacfwd, jvp, vjp, vmap, value_and_grad, hessian, ops, lax, make_jaxpr, linearize, jit
+from jax import custom_jvp, custom_vjp
+from jax.lax import while_loop
 
 
 def if_then_else(cond, val1, val2):

optimism/__init__.py

@@ -0,0 +1,15 @@
+from jax import config
+
+# force double precision floating point arithmetic
+# this is deprecated in jax, we'll have to find another way soon.
+config.update("jax_enable_x64", True)
+
+# silence warnings about no gpu/tpu
+config.update("jax_platforms", "cpu")
+
+# debugging options
+#config.update("jax_debug_nans", True)
+#config.update("jax_debug_infs", True)
+#config.update("jax_disable_jit", True)
+
+del config