From 17e75c40ce9cea4f12df96185d3c1891959551c9 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Tue, 24 Jun 2025 06:50:04 -0700 Subject: [PATCH 01/37] WIP Implement integration/evaluation of functions with arbitrary number and type of fields --- optimism/NewFunctionSpace.py | 118 +++++++++++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 optimism/NewFunctionSpace.py diff --git a/optimism/NewFunctionSpace.py b/optimism/NewFunctionSpace.py new file mode 100644 index 00000000..1b60f7f5 --- /dev/null +++ b/optimism/NewFunctionSpace.py @@ -0,0 +1,118 @@ +import abc +import jax +import jax.numpy as np + +from optimism import Interpolants +from optimism import Mesh +from optimism import QuadratureRule + +class Field(abc.ABC): + pass + + @abc.abstractmethod + def interpolate(self, U, shape, conn): + pass + + @property + @abc.abstractmethod + def element_axis(self): + pass + + @abc.abstractmethod + def compute_shape_functions(self, points): + pass + +class FEField(Field): + element_axis = None + quadpoint_axis = 0 + + def __init__(self, order, dim): + self.order = order + self.dim = dim + self.element, self.element1d = Interpolants.make_parent_elements(order) + + def interpolate(self, field, shape, conn): + e_vector = field[conn] + return shape@e_vector + + def compute_shape_functions(self, points): + return Interpolants.compute_shapes(self.element, points) + + +class DummyShapeFunctions: + values = None + gradients = None + + +class UniformScalarField(Field): + element_axis = None + quadpoint_axis = None + + def interpolate(cls, field, shape, conn): + return U + + def compute_shape_functions(self, points): + return DummyShapeFunctions() + + +class QuadratureField(Field): + element_axis = 0 + quadpoint_axis = 0 + + def interpolate(self, Q, shape, conn): + return Q + + def compute_shape_functions(self, points): + return DummyShapeFunctions() + +class FunctionSpace2: + def __init__(self, spaces, mesh, quadrature_rule): + self._spaces = spaces + self._mesh = mesh + self._quadrature_rule = quadrature_rule + self._shapes = [space.compute_shape_functions(quadrature_rule.xigauss) for space in spaces] + + def evaluate(self, f, *fields): + f_vmap_axis = None + conns_vmap_axis = 0 + compute_element_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, conns_vmap_axis) + tuple(space.element_axis for space in self._spaces)) + return compute_element_values(f, self._mesh.conns, *fields) + + def _evaluate_on_element(self, f, el_conn, *fields): + + f_args = [z[0].interpolate(z[1], z[2].values, el_conn) for z in zip(self._spaces, fields, self._shapes)] + f_batch = jax.vmap(f, tuple(space.quadpoint_axis for space in self._spaces)) + return f_batch(*f_args) + + +if __name__ == "__main__": + p = 1 + dim = 2 + + mesh = Mesh.construct_structured_mesh(2, 2, [0.0, 1.0], [0.0, 2.0], p) + quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) + + u_space = FEField(p, dim) + internal_variable_space = QuadratureField() + time_space = UniformScalarField() + spaces = [u_space, internal_variable_space, time_space] + + fs = FunctionSpace2([u_space, internal_variable_space, time_space], mesh, quad_rule) + # fs = FunctionSpace2([u_space], mesh, quad_rule) + # print(f"{fs._shapes[0].values}") + + def f(u, Q, t): + return np.dot(u, u)*Q[0]*t + + U = jax.vmap(lambda X: np.array([X[0], 0.0]))(mesh.coords) + print(f"{U=}") + Q = 2.0*np.ones((Mesh.num_elements(mesh), len(quad_rule), 1)) + t = 1.0 + + interpolated_values = time_space.interpolate(t, None, None) + print(f"{interpolated_values=}") + + val = fs.evaluate(f, U, Q, t) + print(f"func vals shape = {val.shape}") + print(f"elements = {Mesh.num_elements(mesh)}, quad points = {len(quad_rule)}") + print(f"{val=}") \ No newline at end of file From 4d4f26307b314107208f46cd19fd71fe65b3adf0 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Tue, 24 Jun 2025 07:04:36 -0700 Subject: [PATCH 02/37] Fix bug, output size is at least correct --- optimism/NewFunctionSpace.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/optimism/NewFunctionSpace.py b/optimism/NewFunctionSpace.py index 1b60f7f5..2f279a10 100644 --- a/optimism/NewFunctionSpace.py +++ b/optimism/NewFunctionSpace.py @@ -48,8 +48,8 @@ class UniformScalarField(Field): element_axis = None quadpoint_axis = None - def interpolate(cls, field, shape, conn): - return U + def interpolate(self, field, shape, conn): + return field def compute_shape_functions(self, points): return DummyShapeFunctions() @@ -59,6 +59,9 @@ class QuadratureField(Field): element_axis = 0 quadpoint_axis = 0 + def __init__(self, dim): + self.dim = dim + def interpolate(self, Q, shape, conn): return Q @@ -79,7 +82,6 @@ def evaluate(self, f, *fields): return compute_element_values(f, self._mesh.conns, *fields) def _evaluate_on_element(self, f, el_conn, *fields): - f_args = [z[0].interpolate(z[1], z[2].values, el_conn) for z in zip(self._spaces, fields, self._shapes)] f_batch = jax.vmap(f, tuple(space.quadpoint_axis for space in self._spaces)) return f_batch(*f_args) @@ -88,12 +90,13 @@ def _evaluate_on_element(self, f, el_conn, *fields): if __name__ == "__main__": p = 1 dim = 2 + internal_var_dim = 1 mesh = Mesh.construct_structured_mesh(2, 2, [0.0, 1.0], [0.0, 2.0], p) quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) u_space = FEField(p, dim) - internal_variable_space = QuadratureField() + internal_variable_space = QuadratureField(internal_var_dim) time_space = UniformScalarField() spaces = [u_space, internal_variable_space, time_space] @@ -112,6 +115,8 @@ def f(u, Q, t): interpolated_values = time_space.interpolate(t, None, None) print(f"{interpolated_values=}") + print(f"f() = {f(np.array([1.0, 0.0]), Q[0, 0], 1.0)}") + val = fs.evaluate(f, U, Q, t) print(f"func vals shape = {val.shape}") print(f"elements = {Mesh.num_elements(mesh)}, quad points = {len(quad_rule)}") From 9b4c7375975298e5e79ce5cd8a2e804d2fa850c3 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Wed, 25 Jun 2025 08:36:34 -0700 Subject: [PATCH 03/37] WIP Allow gradients of fields in kernel function --- optimism/NewFunctionSpace.py | 36 +++++++++++++++++++++++------------- 1 file changed, 23 insertions(+), 13 deletions(-) diff --git a/optimism/NewFunctionSpace.py b/optimism/NewFunctionSpace.py index 2f279a10..d27a9c1d 100644 --- a/optimism/NewFunctionSpace.py +++ b/optimism/NewFunctionSpace.py @@ -5,6 +5,7 @@ from optimism import Interpolants from optimism import Mesh from optimism import QuadratureRule +from optimism import VTKWriter class Field(abc.ABC): pass @@ -35,6 +36,9 @@ def interpolate(self, field, shape, conn): e_vector = field[conn] return shape@e_vector + def interpolate_gradient(self, field, dshape, conn): + return jax.vmap(lambda ev, dshp: np.tensordot(ev, dshp, (0, 0)), (None, 0))(field[conn], dshape) + def compute_shape_functions(self, points): return Interpolants.compute_shapes(self.element, points) @@ -78,12 +82,13 @@ def __init__(self, spaces, mesh, quadrature_rule): def evaluate(self, f, *fields): f_vmap_axis = None conns_vmap_axis = 0 - compute_element_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, conns_vmap_axis) + tuple(space.element_axis for space in self._spaces)) - return compute_element_values(f, self._mesh.conns, *fields) + compute_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, conns_vmap_axis) + tuple(space.element_axis for space in self._spaces)) + return compute_values(f, self._mesh.conns, *fields) def _evaluate_on_element(self, f, el_conn, *fields): f_args = [z[0].interpolate(z[1], z[2].values, el_conn) for z in zip(self._spaces, fields, self._shapes)] - f_batch = jax.vmap(f, tuple(space.quadpoint_axis for space in self._spaces)) + f_args += [self._spaces[0].interpolate_gradient(fields[0], self._shapes[0].gradients, el_conn)] + f_batch = jax.vmap(f, tuple(space.quadpoint_axis for space in self._spaces) + (0,)) return f_batch(*f_args) @@ -101,8 +106,6 @@ def _evaluate_on_element(self, f, el_conn, *fields): spaces = [u_space, internal_variable_space, time_space] fs = FunctionSpace2([u_space, internal_variable_space, time_space], mesh, quad_rule) - # fs = FunctionSpace2([u_space], mesh, quad_rule) - # print(f"{fs._shapes[0].values}") def f(u, Q, t): return np.dot(u, u)*Q[0]*t @@ -111,13 +114,20 @@ def f(u, Q, t): print(f"{U=}") Q = 2.0*np.ones((Mesh.num_elements(mesh), len(quad_rule), 1)) t = 1.0 - - interpolated_values = time_space.interpolate(t, None, None) - print(f"{interpolated_values=}") - print(f"f() = {f(np.array([1.0, 0.0]), Q[0, 0], 1.0)}") + # print(f"f() = {f(np.array([1.0, 0.0]), Q[0, 0], 1.0)}") + + # val = fs.evaluate(f, U, Q, t) + # print(f"func vals shape = {val.shape}") + # print(f"elements = {Mesh.num_elements(mesh)}, quad points = {len(quad_rule)}") + # print(f"{val=}") - val = fs.evaluate(f, U, Q, t) - print(f"func vals shape = {val.shape}") - print(f"elements = {Mesh.num_elements(mesh)}, quad points = {len(quad_rule)}") - print(f"{val=}") \ No newline at end of file + # writer = VTKWriter.VTKWriter(mesh, "new_fs") + # writer.write() + + def g(u, Q, t, du): + jax.debug.print("du={du}", du=du) + return np.tensordot(du, du) + + val = fs.evaluate(g, U, Q, t) + print(f"{val=}") From 2a0c004447daeebe39a12d0944275bae5436d99f Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Wed, 25 Jun 2025 11:08:00 -0700 Subject: [PATCH 04/37] Allow user to specify signature of q-function in terms of values and gradients of fields --- optimism/NewFunctionSpace.py | 77 +++++++++++++++++++++++++++--------- 1 file changed, 59 insertions(+), 18 deletions(-) diff --git a/optimism/NewFunctionSpace.py b/optimism/NewFunctionSpace.py index d27a9c1d..810761ff 100644 --- a/optimism/NewFunctionSpace.py +++ b/optimism/NewFunctionSpace.py @@ -1,4 +1,6 @@ import abc +import dataclasses +import functools import jax import jax.numpy as np @@ -11,7 +13,11 @@ class Field(abc.ABC): pass @abc.abstractmethod - def interpolate(self, U, shape, conn): + def interpolate(self, shape, U, conn): + pass + + @abc.abstractmethod + def interpolate_gradient(self, dshape, U, conn): pass @property @@ -23,6 +29,7 @@ def element_axis(self): def compute_shape_functions(self, points): pass + class FEField(Field): element_axis = None quadpoint_axis = 0 @@ -32,11 +39,11 @@ def __init__(self, order, dim): self.dim = dim self.element, self.element1d = Interpolants.make_parent_elements(order) - def interpolate(self, field, shape, conn): + def interpolate(self, shape, field, conn): e_vector = field[conn] return shape@e_vector - def interpolate_gradient(self, field, dshape, conn): + def interpolate_gradient(self, dshape, field, conn): return jax.vmap(lambda ev, dshp: np.tensordot(ev, dshp, (0, 0)), (None, 0))(field[conn], dshape) def compute_shape_functions(self, points): @@ -52,9 +59,12 @@ class UniformScalarField(Field): element_axis = None quadpoint_axis = None - def interpolate(self, field, shape, conn): + def interpolate(self, shape, field, conn): return field + def interpolate_gradient(self, dshape, field, conn): + raise NotImplementedError + def compute_shape_functions(self, points): return DummyShapeFunctions() @@ -66,18 +76,42 @@ class QuadratureField(Field): def __init__(self, dim): self.dim = dim - def interpolate(self, Q, shape, conn): + def interpolate(self, shape, Q, conn): return Q + def interpolate_gradient(self, dshape, Q, conn): + raise NotImplementedError + def compute_shape_functions(self, points): return DummyShapeFunctions() +@dataclasses.dataclass +class Value: + field: int + + +@dataclasses.dataclass +class Gradient: + field: int + + class FunctionSpace2: - def __init__(self, spaces, mesh, quadrature_rule): + def __init__(self, spaces, qfunction_signature, mesh, quadrature_rule): self._spaces = spaces self._mesh = mesh self._quadrature_rule = quadrature_rule self._shapes = [space.compute_shape_functions(quadrature_rule.xigauss) for space in spaces] + self._interpolators = [] + self._input_fields = [] + for input in qfunction_signature: + self._input_fields.append(input.field) + if type(input) is Value: + self._interpolators.append(functools.partial(self._spaces[input.field].interpolate, self._shapes[input.field].values)) + elif type(input) is Gradient: + self._interpolators.append(functools.partial(self._spaces[input.field].interpolate_gradient, self._shapes[input.field].gradients)) + else: + raise IndexError + def evaluate(self, f, *fields): f_vmap_axis = None @@ -86,14 +120,13 @@ def evaluate(self, f, *fields): return compute_values(f, self._mesh.conns, *fields) def _evaluate_on_element(self, f, el_conn, *fields): - f_args = [z[0].interpolate(z[1], z[2].values, el_conn) for z in zip(self._spaces, fields, self._shapes)] - f_args += [self._spaces[0].interpolate_gradient(fields[0], self._shapes[0].gradients, el_conn)] - f_batch = jax.vmap(f, tuple(space.quadpoint_axis for space in self._spaces) + (0,)) + f_args = [interp(fields[field_id], el_conn) for (interp, field_id) in zip(self._interpolators, self._input_fields)] + f_batch = jax.vmap(f, tuple(self._spaces[input].quadpoint_axis for input in self._input_fields)) return f_batch(*f_args) if __name__ == "__main__": - p = 1 + p = 2 dim = 2 internal_var_dim = 1 @@ -105,18 +138,20 @@ def _evaluate_on_element(self, f, el_conn, *fields): time_space = UniformScalarField() spaces = [u_space, internal_variable_space, time_space] - fs = FunctionSpace2([u_space, internal_variable_space, time_space], mesh, quad_rule) + inputs = (Value(0), Value(1), Value(2), Gradient(0)) + + fs = FunctionSpace2([u_space, internal_variable_space, time_space], inputs, mesh, quad_rule) def f(u, Q, t): return np.dot(u, u)*Q[0]*t - - U = jax.vmap(lambda X: np.array([X[0], 0.0]))(mesh.coords) - print(f"{U=}") - Q = 2.0*np.ones((Mesh.num_elements(mesh), len(quad_rule), 1)) - t = 1.0 # print(f"f() = {f(np.array([1.0, 0.0]), Q[0, 0], 1.0)}") + U = mesh.coords + U = U.at[:, 1].set(0.0) + Q = 2.0*np.ones((Mesh.num_elements(mesh), len(quad_rule), 1)) + t = 1.0 + # val = fs.evaluate(f, U, Q, t) # print(f"func vals shape = {val.shape}") # print(f"elements = {Mesh.num_elements(mesh)}, quad points = {len(quad_rule)}") @@ -126,8 +161,14 @@ def f(u, Q, t): # writer.write() def g(u, Q, t, du): - jax.debug.print("du={du}", du=du) - return np.tensordot(du, du) + return np.dot(u, u)*Q[0]*t*np.tensordot(du, du) + + test = fs._interpolators[0](U, mesh.conns[0]) + print(f"{test=}") + + print(fs._interpolators[1]) + test2 = fs._interpolators[1](Q, mesh.conns[0]) + print(f"{test2=}") val = fs.evaluate(g, U, Q, t) print(f"{val=}") From b23f483fcf6ff73a76d0b6803e5b40a39815ecbe Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Wed, 25 Jun 2025 11:16:37 -0700 Subject: [PATCH 05/37] Put missing required property (quad point vmap axis) in abstract base class --- optimism/NewFunctionSpace.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/optimism/NewFunctionSpace.py b/optimism/NewFunctionSpace.py index 810761ff..b189598b 100644 --- a/optimism/NewFunctionSpace.py +++ b/optimism/NewFunctionSpace.py @@ -25,6 +25,11 @@ def interpolate_gradient(self, dshape, U, conn): def element_axis(self): pass + @property + @abc.abstractmethod + def quadpoint_axis(self): + pass + @abc.abstractmethod def compute_shape_functions(self, points): pass From f1d85a79fcefe3cc8bb151f9507bdc343661ad84 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Thu, 26 Jun 2025 06:48:28 -0700 Subject: [PATCH 06/37] Add support for quadrature weight-like fields (quad point data on parent element only) --- optimism/NewFunctionSpace.py | 48 +++++++++++++++++++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/optimism/NewFunctionSpace.py b/optimism/NewFunctionSpace.py index b189598b..bb1977ae 100644 --- a/optimism/NewFunctionSpace.py +++ b/optimism/NewFunctionSpace.py @@ -85,7 +85,22 @@ def interpolate(self, shape, Q, conn): return Q def interpolate_gradient(self, dshape, Q, conn): - raise NotImplementedError + raise NotImplementedError(f"Gradients not supported for {type(self).__name__}") + + def compute_shape_functions(self, points): + return DummyShapeFunctions() + + +class ParametricElementField(Field): + """Holds things like quadrature weights and shape function values.""" + element_axis = None + quadpoint_axis = 0 + + def interpolate(self, shape, field, conn): + return field + + def interpolate_gradient(self, dshape, U, conn): + raise NotImplementedError(f"Gradients not supported for {type(self).__name__}") def compute_shape_functions(self, points): return DummyShapeFunctions() @@ -177,3 +192,34 @@ def g(u, Q, t, du): val = fs.evaluate(g, U, Q, t) print(f"{val=}") + + def h(dX_dxi, du_dxi, w): + E = 1.0 + nu = 0.0 + mu = 0.5*E/(1 + nu) + lam = E*nu/(1 + nu)/(1 - 2*nu) + dxi_dX = np.linalg.inv(dX_dxi) + du_dX = du_dxi@dxi_dX + dV = np.linalg.det(dX_dxi)*w + strain = 0.5*(du_dX + du_dX.T) + energy_density = mu*np.tensordot(strain, strain) + 0.5*lam*np.trace(strain)**2 + return energy_density*dV + + def area(dX_dxi, du_dxi, w): + return np.linalg.det(dX_dxi)*w + + quad_weight_space = ParametricElementField() + + spaces = [u_space, u_space, quad_weight_space] + qfunction_signature = [Gradient(0), Gradient(1), Value(2)] + + fs2 = FunctionSpace2(spaces, qfunction_signature, mesh, quad_rule) + energies = fs2.evaluate(h, mesh.coords, U, quad_rule.wgauss) + print(f"{energies=}") + + def potential(X, U): + return np.sum(fs2.evaluate(h, X, U, quad_rule.wgauss)) + + compute_force = jax.grad(potential, 1) + R = compute_force(mesh.coords, U) + print(f"{R=}") From dddaa335d6a0e0e7289556d99fd6184bc9db2548 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Thu, 26 Jun 2025 07:22:26 -0700 Subject: [PATCH 07/37] Tidy and give more appropriate names --- optimism/NewFunctionSpace.py | 55 +++++++++++++++++++++--------------- 1 file changed, 32 insertions(+), 23 deletions(-) diff --git a/optimism/NewFunctionSpace.py b/optimism/NewFunctionSpace.py index bb1977ae..601fe395 100644 --- a/optimism/NewFunctionSpace.py +++ b/optimism/NewFunctionSpace.py @@ -35,14 +35,15 @@ def compute_shape_functions(self, points): pass -class FEField(Field): +class PkField(Field): + """Standard Lagrange polynomial finite element fields.""" element_axis = None quadpoint_axis = 0 - def __init__(self, order, dim): - self.order = order + def __init__(self, k, dim): + self.order = k self.dim = dim - self.element, self.element1d = Interpolants.make_parent_elements(order) + self.element, self.element1d = Interpolants.make_parent_elements(k) def interpolate(self, shape, field, conn): e_vector = field[conn] @@ -60,7 +61,8 @@ class DummyShapeFunctions: gradients = None -class UniformScalarField(Field): +class UniformField(Field): + """A unique value for the whole mesh (things like time).""" element_axis = None quadpoint_axis = None @@ -68,13 +70,15 @@ def interpolate(self, shape, field, conn): return field def interpolate_gradient(self, dshape, field, conn): - raise NotImplementedError + grad_shape = field.shape + dshape.shape[-1] + return np.zeros(grad_shape) def compute_shape_functions(self, points): return DummyShapeFunctions() class QuadratureField(Field): + """Arrays defined directly at quadrature points (things like internal variables).""" element_axis = 0 quadpoint_axis = 0 @@ -92,7 +96,7 @@ def compute_shape_functions(self, points): class ParametricElementField(Field): - """Holds things like quadrature weights and shape function values.""" + """Things like quadrature weights and shape function values.""" element_axis = None quadpoint_axis = 0 @@ -115,23 +119,28 @@ class Gradient: field: int -class FunctionSpace2: +def pushforward(du_dxi, dX_dxi): + return du_dxi@np.linalg.inv(dX_dxi) + + +def _make_interpolation_function(input, spaces, shapes): + """Helper function to choose correct field space interpolation function for each input.""" + if type(input) is Value: + return functools.partial(spaces[input.field].interpolate, shapes[input.field].values) + elif type(input) is Gradient: + return functools.partial(spaces[input.field].interpolate_gradient, shapes[input.field].gradients) + else: + raise IndexError(f"Field space {input.field} exceeds number of field spaces, which is {len(spaces)}.") + + +class FieldEvaluator: def __init__(self, spaces, qfunction_signature, mesh, quadrature_rule): self._spaces = spaces self._mesh = mesh self._quadrature_rule = quadrature_rule self._shapes = [space.compute_shape_functions(quadrature_rule.xigauss) for space in spaces] - self._interpolators = [] - self._input_fields = [] - for input in qfunction_signature: - self._input_fields.append(input.field) - if type(input) is Value: - self._interpolators.append(functools.partial(self._spaces[input.field].interpolate, self._shapes[input.field].values)) - elif type(input) is Gradient: - self._interpolators.append(functools.partial(self._spaces[input.field].interpolate_gradient, self._shapes[input.field].gradients)) - else: - raise IndexError - + self._input_fields = tuple(input.field for input in qfunction_signature) + self._interpolators = tuple(_make_interpolation_function(input, self._spaces, self._shapes) for input in qfunction_signature) def evaluate(self, f, *fields): f_vmap_axis = None @@ -153,14 +162,14 @@ def _evaluate_on_element(self, f, el_conn, *fields): mesh = Mesh.construct_structured_mesh(2, 2, [0.0, 1.0], [0.0, 2.0], p) quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) - u_space = FEField(p, dim) + u_space = PkField(p, dim) internal_variable_space = QuadratureField(internal_var_dim) - time_space = UniformScalarField() + time_space = UniformField() spaces = [u_space, internal_variable_space, time_space] inputs = (Value(0), Value(1), Value(2), Gradient(0)) - fs = FunctionSpace2([u_space, internal_variable_space, time_space], inputs, mesh, quad_rule) + fs = FieldEvaluator([u_space, internal_variable_space, time_space], inputs, mesh, quad_rule) def f(u, Q, t): return np.dot(u, u)*Q[0]*t @@ -213,7 +222,7 @@ def area(dX_dxi, du_dxi, w): spaces = [u_space, u_space, quad_weight_space] qfunction_signature = [Gradient(0), Gradient(1), Value(2)] - fs2 = FunctionSpace2(spaces, qfunction_signature, mesh, quad_rule) + fs2 = FieldEvaluator(spaces, qfunction_signature, mesh, quad_rule) energies = fs2.evaluate(h, mesh.coords, U, quad_rule.wgauss) print(f"{energies=}") From cd22ac0106137a8f8fb38a9bad32fd128e55fb5d Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Fri, 27 Jun 2025 10:42:21 -0700 Subject: [PATCH 08/37] Put in more input validation in constructor --- optimism/NewFunctionSpace.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/optimism/NewFunctionSpace.py b/optimism/NewFunctionSpace.py index 601fe395..8932f7ef 100644 --- a/optimism/NewFunctionSpace.py +++ b/optimism/NewFunctionSpace.py @@ -125,12 +125,14 @@ def pushforward(du_dxi, dX_dxi): def _make_interpolation_function(input, spaces, shapes): """Helper function to choose correct field space interpolation function for each input.""" + if input.field >= len(spaces): + raise IndexError(f"Field space {input.field} exceeds number of field spaces, which is {len(spaces) - 1}.") if type(input) is Value: return functools.partial(spaces[input.field].interpolate, shapes[input.field].values) elif type(input) is Gradient: return functools.partial(spaces[input.field].interpolate_gradient, shapes[input.field].gradients) else: - raise IndexError(f"Field space {input.field} exceeds number of field spaces, which is {len(spaces)}.") + raise TypeError("Type of object in qfunction signature is invalid.") class FieldEvaluator: From 0d4d8ac784be0b0f27b7eb80c8a408a85e738bc9 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Sun, 6 Jul 2025 02:56:23 -0700 Subject: [PATCH 09/37] Implement field evaluation in another way that allows transforming the shape functions at the element level --- optimism/FieldEvaluator.py | 185 +++++++++++++++++++++++++++++++++++++ 1 file changed, 185 insertions(+) create mode 100644 optimism/FieldEvaluator.py diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py new file mode 100644 index 00000000..0bce2712 --- /dev/null +++ b/optimism/FieldEvaluator.py @@ -0,0 +1,185 @@ +import abc +import dataclasses +import functools +import jax +import jax.numpy as np + +from optimism import Interpolants +from optimism import Mesh +from optimism import QuadratureRule +from optimism import VTKWriter + +class Field(abc.ABC): + pass + + @abc.abstractmethod + def interpolate(self, shape, U, conn, jac): + pass + + @abc.abstractmethod + def interpolate_gradient(self, dshape, U, conn, jac): + pass + + @property + @abc.abstractmethod + def element_axis(self): + pass + + @property + @abc.abstractmethod + def quadpoint_axis(self): + pass + + @abc.abstractmethod + def compute_shape_functions(self, points): + pass + + @abc.abstractmethod + def map_shape_functions(self, shapes, jacs): + return shapes + + +class PkField(Field): + """Standard Lagrange polynomial finite element fields.""" + element_axis = None + quadpoint_axis = 0 + + def __init__(self, k, dim): + self.order = k + self.dim = dim + self.element, self.element1d = Interpolants.make_parent_elements(k) + + def interpolate(self, shape, field, conn): + e_vector = field[conn] + return shape.values@e_vector + + def interpolate_gradient(self, shape, field, conn): + return jax.vmap(lambda ev, dshp: np.tensordot(ev, dshp, (0, 0)), (None, 0))(field[conn], shape.gradients) + + def compute_shape_functions(self, points): + return Interpolants.compute_shapes(self.element, points) + + def map_shape_functions(self, shapes, jacs): + shape_grads = jax.vmap(lambda dshp, J: dshp@np.linalg.inv(J))(shapes.gradients, jacs) + return Interpolants.ShapeFunctions(shapes.values, shape_grads) + + +class UniformField(Field): + """A unique value for the whole mesh (things like time).""" + element_axis = None + quadpoint_axis = None + + def interpolate(self, shape, field, conn): + return field + + def interpolate_gradient(self, shape, field, conn): + grad_shape = field.shape + shape.gradients.shape[-1] + return np.zeros(grad_shape) + + def compute_shape_functions(self, points): + return Interpolants.ShapeFunctions(np.array([]), np.array([])) + + def map_shape_functions(self, shapes, jacs): + return super().map_shape_functions(shapes, jacs) + + +@dataclasses.dataclass +class Value: + """Sentinel to indicate you want to interpolate the value of the field.""" + field: int + + +@dataclasses.dataclass +class Gradient: + """Sentinel to indicate you want to interpolate the gradient of the field.""" + field: int + + +def _choose_interpolation_function(input, spaces): + """Helper function to choose correct field space interpolation function for each input.""" + if input.field >= len(spaces): + raise IndexError(f"Field space {input.field} exceeds range, which is {len(spaces) - 1}.") + if type(input) is Value: + return spaces[input.field].interpolate + elif type(input) is Gradient: + return spaces[input.field].interpolate_gradient + else: + raise TypeError("Type of object in qfunction signature is invalid.") + + +class FieldEvaluator: + def __init__(self, spaces, qfunction_signature, mesh, quadrature_rule): + # the coord space should live on the Mesh eventually, or be metadata of the coords themselves + self._coord_space = PkField(mesh.parentElement.degree, mesh.coords.shape[1]) + self._coord_shapes = self._coord_space.compute_shape_functions(quadrature_rule.xigauss) + + self._spaces = spaces + self._mesh = mesh + self._quadrature_rule = quadrature_rule + self._shapes = [space.compute_shape_functions(quadrature_rule.xigauss) for space in spaces] + self._input_fields = tuple(input.field for input in qfunction_signature) + self._interpolators = tuple(_choose_interpolation_function(input, self._spaces) for input in qfunction_signature) + + def evaluate(self, f, coords, *fields): + f_vmap_axis = None + conns_vmap_axis = 0 + coords_vmap_axis = None + compute_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, conns_vmap_axis, coords_vmap_axis) + tuple(space.element_axis for space in self._spaces)) + return compute_values(f, self._mesh.conns, coords, *fields) + + def _evaluate_on_element(self, f, el_conn, coords, *fields): + jacs = self._coord_space.interpolate_gradient(self._coord_shapes, coords, el_conn) + shapes = [space.map_shape_functions(shape, jacs) for (space, shape) in zip(self._spaces, self._shapes)] + f_args = [interp(shapes[field_id], fields[field_id], el_conn) for (interp, field_id) in zip(self._interpolators, self._input_fields)] + f_batch = jax.vmap(f, tuple(self._spaces[input].quadpoint_axis for input in self._input_fields)) + return f_batch(*f_args) + + def _integrate_over_element(self, f, el_conn, coords, *fields): + jacs = self._coord_space.interpolate_gradient(self._coord_shapes, coords, el_conn) + shapes = [space.map_shape_functions(shape, jacs) for (space, shape) in zip(self._spaces, self._shapes)] + dVs = jax.vmap(lambda J, w: np.linalg.det(J)*w)(jacs, self._quadrature_rule.wgauss) + f_args = [interp(shapes[field_id], fields[field_id], el_conn) for (interp, field_id) in zip(self._interpolators, self._input_fields)] + f_batch = jax.vmap(f, tuple(self._spaces[input].quadpoint_axis for input in self._input_fields)) + f_vals = f_batch(*f_args) + return np.dot(f_vals, dVs) + + def integrate(self, f, coords, *fields): + f_vmap_axis = None + conns_vmap_axis = 0 + coords_vmap_axis = None + compute_values = jax.vmap(self._integrate_over_element, (f_vmap_axis, conns_vmap_axis, coords_vmap_axis) + tuple(space.element_axis for space in self._spaces)) + return np.sum(compute_values(f, self._mesh.conns, coords, *fields)) + +if __name__ == "__main__": + p = 2 + dim = 2 + + mesh = Mesh.construct_structured_mesh(2, 2, [0.0, 3.0], [0.0, 2.0], p) + quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) + + u_space = PkField(p, dim) + time_space = UniformField() + spaces = [u_space, time_space] + DISPLACEMENT = 0 + TIME = 1 + + inputs = (Value(DISPLACEMENT), Value(TIME), Gradient(DISPLACEMENT)) + + def scaled_disp_grad(u, t, dudX): + return t*dudX + + fs = FieldEvaluator([u_space, time_space], inputs, mesh, quad_rule) + + target_disp_grad = np.array([[0.1, 0.01], + [0.05, 0.3]]) + U = np.einsum('aj, ij', mesh.coords, target_disp_grad + np.identity(2)) - mesh.coords + time = 2.0 + + disp_grads = fs.evaluate(scaled_disp_grad, mesh.coords, U, time) + print(f"{disp_grads=}") + + def area(u, t, dudX): + return 1.0 + + A = fs.integrate(area, mesh.coords, U, time) + print(f"{A=}") From b5d40d46f790b0842fb62508a613d67d269864a6 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Sun, 6 Jul 2025 05:03:59 -0700 Subject: [PATCH 10/37] Implement first unit tests --- optimism/test/test_FieldEvaluator.py | 52 ++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) create mode 100644 optimism/test/test_FieldEvaluator.py diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py new file mode 100644 index 00000000..08c625f0 --- /dev/null +++ b/optimism/test/test_FieldEvaluator.py @@ -0,0 +1,52 @@ +import pytest +import jax.numpy as np + +from optimism import Mesh +from optimism.FieldEvaluator import * + +class TestBasics: + p = 2 + dim = 2 + mesh = Mesh.construct_structured_mesh(2, 2, [0.0, 3.0], [0.0, 2.0], p) + quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) + + def test_area(self): + spaces = [PkField(self.p, self.dim)] + inputs = [Value(0)] + field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + U = np.zeros_like(self.mesh.coords) + def f(u): + return 1.0 + area = field_evaluator.integrate(f, self.mesh.coords, U) + assert pytest.approx(area) == 6 + + def test_area_another_way(self): + spaces = [PkField(self.p, self.dim)] + inputs = [Gradient(0)] + field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + + def f(dXdX): + return np.trace(dXdX)/self.dim + + area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.coords) + assert pytest.approx(area) == 6 + + def test_gradient_interpolation(self): + p = 2 + spaces = [PkField(p, self.dim)] + inputs = [Gradient(0)] + field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + + target_disp_grad = np.array([[0.1, 0.01], + [0.05, 0.3]]) + #coords_linear = self.mesh.coords[self.mesh.simplexNodesOrdinals, :] + U = np.einsum('aj, ij', self.mesh.coords, target_disp_grad + np.identity(2)) - self.mesh.coords + + def f(dudX): + return dudX + + disp_grads = field_evaluator.evaluate(f, self.mesh.coords, U) + + for H in disp_grads.reshape(-1, 2, 2): + assert pytest.approx(H) == target_disp_grad + From e135b0309abacfc596dd730282793dc3bd3e272b Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Sun, 6 Jul 2025 05:07:34 -0700 Subject: [PATCH 11/37] Implement uniform fields and quadrature fields --- optimism/FieldEvaluator.py | 43 +++++++++++++++++++++++++++++++++++--- 1 file changed, 40 insertions(+), 3 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index 0bce2712..26db067f 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -73,8 +73,28 @@ def interpolate(self, shape, field, conn): return field def interpolate_gradient(self, shape, field, conn): - grad_shape = field.shape + shape.gradients.shape[-1] - return np.zeros(grad_shape) + raise NotImplementedError(f"Gradients not supported for {type(self).__name__}") + + def compute_shape_functions(self, points): + return Interpolants.ShapeFunctions(np.array([]), np.array([])) + + def map_shape_functions(self, shapes, jacs): + return super().map_shape_functions(shapes, jacs) + + +class QuadratureField(Field): + """Arrays defined directly at quadrature points (things like internal variables).""" + element_axis = 0 + quadpoint_axis = 0 + + def __init__(self, dim): + self.dim = dim + + def interpolate(self, shape, field, conn): + return field + + def interpolate_gradient(self, shape, field, conn): + raise NotImplementedError(f"Gradients not supported for {type(self).__name__}") def compute_shape_functions(self, points): return Interpolants.ShapeFunctions(np.array([]), np.array([])) @@ -109,7 +129,7 @@ def _choose_interpolation_function(input, spaces): class FieldEvaluator: def __init__(self, spaces, qfunction_signature, mesh, quadrature_rule): - # the coord space should live on the Mesh eventually, or be metadata of the coords themselves + # the coord space should live on the Mesh self._coord_space = PkField(mesh.parentElement.degree, mesh.coords.shape[1]) self._coord_shapes = self._coord_space.compute_shape_functions(quadrature_rule.xigauss) @@ -183,3 +203,20 @@ def area(u, t, dudX): A = fs.integrate(area, mesh.coords, U, time) print(f"{A=}") + + + # Example 2, intepolate positions to quadrature points + inputs = [Value(0), Gradient(0)] + fs = FieldEvaluator([u_space], inputs, mesh, quad_rule) + + def position(X, dXdX): + return X + + qp_coords = fs.evaluate(position, mesh.coords, mesh.coords) + print(f"{qp_coords=}") + + def area2(X, dXdX): + dim = dXdX.shape[0] + return np.trace(dXdX)/dim + A2 = fs.integrate(area2, mesh.coords, mesh.coords) + print(f"area = {A2}") From 561735919289e4100d6900919c60708c04f4f22a Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Sun, 6 Jul 2025 07:16:45 -0700 Subject: [PATCH 12/37] Remove temperory tests that were located directly in module --- optimism/FieldEvaluator.py | 51 -------------------------------------- 1 file changed, 51 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index 26db067f..f9757beb 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -169,54 +169,3 @@ def integrate(self, f, coords, *fields): coords_vmap_axis = None compute_values = jax.vmap(self._integrate_over_element, (f_vmap_axis, conns_vmap_axis, coords_vmap_axis) + tuple(space.element_axis for space in self._spaces)) return np.sum(compute_values(f, self._mesh.conns, coords, *fields)) - -if __name__ == "__main__": - p = 2 - dim = 2 - - mesh = Mesh.construct_structured_mesh(2, 2, [0.0, 3.0], [0.0, 2.0], p) - quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) - - u_space = PkField(p, dim) - time_space = UniformField() - spaces = [u_space, time_space] - DISPLACEMENT = 0 - TIME = 1 - - inputs = (Value(DISPLACEMENT), Value(TIME), Gradient(DISPLACEMENT)) - - def scaled_disp_grad(u, t, dudX): - return t*dudX - - fs = FieldEvaluator([u_space, time_space], inputs, mesh, quad_rule) - - target_disp_grad = np.array([[0.1, 0.01], - [0.05, 0.3]]) - U = np.einsum('aj, ij', mesh.coords, target_disp_grad + np.identity(2)) - mesh.coords - time = 2.0 - - disp_grads = fs.evaluate(scaled_disp_grad, mesh.coords, U, time) - print(f"{disp_grads=}") - - def area(u, t, dudX): - return 1.0 - - A = fs.integrate(area, mesh.coords, U, time) - print(f"{A=}") - - - # Example 2, intepolate positions to quadrature points - inputs = [Value(0), Gradient(0)] - fs = FieldEvaluator([u_space], inputs, mesh, quad_rule) - - def position(X, dXdX): - return X - - qp_coords = fs.evaluate(position, mesh.coords, mesh.coords) - print(f"{qp_coords=}") - - def area2(X, dXdX): - dim = dXdX.shape[0] - return np.trace(dXdX)/dim - A2 = fs.integrate(area2, mesh.coords, mesh.coords) - print(f"area = {A2}") From 1c7631af10c5610e53d93305772972fcb5a660d5 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Sun, 6 Jul 2025 07:30:24 -0700 Subject: [PATCH 13/37] Imrpove clarity of tests --- optimism/test/test_FieldEvaluator.py | 50 ++++++++++++++++------------ 1 file changed, 28 insertions(+), 22 deletions(-) diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index 08c625f0..65683bca 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -10,28 +10,8 @@ class TestBasics: mesh = Mesh.construct_structured_mesh(2, 2, [0.0, 3.0], [0.0, 2.0], p) quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) - def test_area(self): - spaces = [PkField(self.p, self.dim)] - inputs = [Value(0)] - field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) - U = np.zeros_like(self.mesh.coords) - def f(u): - return 1.0 - area = field_evaluator.integrate(f, self.mesh.coords, U) - assert pytest.approx(area) == 6 - - def test_area_another_way(self): - spaces = [PkField(self.p, self.dim)] - inputs = [Gradient(0)] - field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) - - def f(dXdX): - return np.trace(dXdX)/self.dim - - area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.coords) - assert pytest.approx(area) == 6 - - def test_gradient_interpolation(self): + def test_gradient_evaluation(self): + "Check the gradient of an affine field" p = 2 spaces = [PkField(p, self.dim)] inputs = [Gradient(0)] @@ -50,3 +30,29 @@ def f(dudX): for H in disp_grads.reshape(-1, 2, 2): assert pytest.approx(H) == target_disp_grad + def test_trivial_integral(self): + spaces = [PkField(self.p, self.dim)] + inputs = [Value(0)] + field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + U = np.zeros_like(self.mesh.coords) + def f(u): + return 1.0 + area = field_evaluator.integrate(f, self.mesh.coords, U) + assert pytest.approx(area) == 6 + + def test_integral_with_one_nodal_field(self): + "Computes area in a non-trivial way, checking consistency of gradient and integral operators." + spaces = [PkField(self.p, self.dim)] + POSITION = 0 + # We're taking the gradient of position, which is just the identity tensor + inputs = [Gradient(POSITION)] + field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + + def f(dXdX): + # dXdX == identity + return np.trace(dXdX)/self.dim + + area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.coords) + assert pytest.approx(area) == 6 + + From 8f1581bce2b86f384509b8e17b12854e0485533a Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Wed, 16 Jul 2025 13:47:55 -0700 Subject: [PATCH 14/37] Allow different order fields by storing a connectivity table for each field space --- optimism/FieldEvaluator.py | 62 +++++++++++++++++++++++++++- optimism/test/test_FieldEvaluator.py | 8 +++- 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index f9757beb..aeeaf80a 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -44,10 +44,12 @@ class PkField(Field): element_axis = None quadpoint_axis = 0 - def __init__(self, k, dim): + def __init__(self, k, dim, mesh): self.order = k self.dim = dim self.element, self.element1d = Interpolants.make_parent_elements(k) + self.mesh = mesh + self.conns = self._make_connectivity() def interpolate(self, shape, field, conn): e_vector = field[conn] @@ -62,6 +64,53 @@ def compute_shape_functions(self, points): def map_shape_functions(self, shapes, jacs): shape_grads = jax.vmap(lambda dshp, J: dshp@np.linalg.inv(J))(shapes.gradients, jacs) return Interpolants.ShapeFunctions(shapes.values, shape_grads) + + def _make_connectivity(self): + mesh_conns = self.mesh.conns + if self.order == 1: + return mesh_conns + + conns = np.zeros((mesh_conns.shape[0], self.element.coordinates.shape[0]), dtype=int) + + # step 1/3: vertex nodes + conns = conns.at[:, self.element.vertexNodes].set(mesh_conns) + nodeOrdinalOffset = mesh_conns.max() + 1 # offset for later node numbering + + # step 2/3: mid-edge nodes (excluding vertices) + _, edges = Mesh.create_edges(mesh_conns) + + nNodesPerEdge = self.element1d.interiorNodes.size + for e, edge in enumerate(edges): + edgeNodeOrdinals = nodeOrdinalOffset + np.arange(e*nNodesPerEdge,(e+1)*nNodesPerEdge) + + elemLeft = edge[0] + sideLeft = edge[1] + edgeMasterNodes = self.element.faceNodes[sideLeft][self.element1d.interiorNodes] + conns = conns.at[elemLeft, edgeMasterNodes].set(edgeNodeOrdinals) + + elemRight = edge[2] + if elemRight >= 0: + sideRight = edge[3] + edgeMasterNodes = self.element.faceNodes[sideRight][self.element1d.interiorNodes] + conns = conns.at[elemRight, edgeMasterNodes].set(np.flip(edgeNodeOrdinals)) + + nEdges = edges.shape[0] + nodeOrdinalOffset += nEdges*nNodesPerEdge # for offset of interior node numbering + + # step 3/3: interior nodes + nInNodesPerTri = self.element.interiorNodes.shape[0] + if nInNodesPerTri > 0: + + def add_element_interior_nodes(conn, newNodeOrdinals): + return conn.at[self.element.interiorNodes].set(newNodeOrdinals) + + nTri = conns.shape[0] + newNodeOrdinals = np.arange(nTri*nInNodesPerTri).reshape(nTri,nInNodesPerTri) \ + + nodeOrdinalOffset + + conns = jax.vmap(add_element_interior_nodes)(conns, newNodeOrdinals) + return conns + class UniformField(Field): @@ -169,3 +218,14 @@ def integrate(self, f, coords, *fields): coords_vmap_axis = None compute_values = jax.vmap(self._integrate_over_element, (f_vmap_axis, conns_vmap_axis, coords_vmap_axis) + tuple(space.element_axis for space in self._spaces)) return np.sum(compute_values(f, self._mesh.conns, coords, *fields)) + +if __name__ == "__main__": + p = 3 + dim = 2 + p_coords = 1 + mesh = Mesh.construct_structured_mesh(2, 2, [0.0, 3.0], [0.0, 2.0], p_coords) + disp_space = PkField(p, dim, mesh) + conns = disp_space._make_connectivity() + print(f"{conns=}") + print(f"{mesh.conns=}") + print(f"{mesh.coords=}") diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index 65683bca..0643d0d4 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -54,5 +54,9 @@ def f(dXdX): area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.coords) assert pytest.approx(area) == 6 - - + + def test_conn(self): + p = 1 + disp_space = PkField(p, self.dim, self.mesh) + conns = disp_space._make_connectivity() + print(f"{conns=}") From a8704c4c035e3d6a583ee675ebf5c440b3547aa1 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Tue, 5 Aug 2025 15:38:54 -0700 Subject: [PATCH 15/37] Make tests pass with recent change to Pk space (holds mesh) --- optimism/FieldEvaluator.py | 15 +++++---------- optimism/test/test_FieldEvaluator.py | 21 ++++++++++++--------- 2 files changed, 17 insertions(+), 19 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index aeeaf80a..d3ee8b55 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -10,8 +10,6 @@ from optimism import VTKWriter class Field(abc.ABC): - pass - @abc.abstractmethod def interpolate(self, shape, U, conn, jac): pass @@ -179,7 +177,7 @@ def _choose_interpolation_function(input, spaces): class FieldEvaluator: def __init__(self, spaces, qfunction_signature, mesh, quadrature_rule): # the coord space should live on the Mesh - self._coord_space = PkField(mesh.parentElement.degree, mesh.coords.shape[1]) + self._coord_space = PkField(mesh.parentElement.degree, mesh.coords.shape[1], mesh) self._coord_shapes = self._coord_space.compute_shape_functions(quadrature_rule.xigauss) self._spaces = spaces @@ -196,6 +194,9 @@ def evaluate(self, f, coords, *fields): compute_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, conns_vmap_axis, coords_vmap_axis) + tuple(space.element_axis for space in self._spaces)) return compute_values(f, self._mesh.conns, coords, *fields) + def integrate(self, f, coords, *fields): + return np.sum(self.evaluate(f, coords, *fields)) + def _evaluate_on_element(self, f, el_conn, coords, *fields): jacs = self._coord_space.interpolate_gradient(self._coord_shapes, coords, el_conn) shapes = [space.map_shape_functions(shape, jacs) for (space, shape) in zip(self._spaces, self._shapes)] @@ -211,13 +212,7 @@ def _integrate_over_element(self, f, el_conn, coords, *fields): f_batch = jax.vmap(f, tuple(self._spaces[input].quadpoint_axis for input in self._input_fields)) f_vals = f_batch(*f_args) return np.dot(f_vals, dVs) - - def integrate(self, f, coords, *fields): - f_vmap_axis = None - conns_vmap_axis = 0 - coords_vmap_axis = None - compute_values = jax.vmap(self._integrate_over_element, (f_vmap_axis, conns_vmap_axis, coords_vmap_axis) + tuple(space.element_axis for space in self._spaces)) - return np.sum(compute_values(f, self._mesh.conns, coords, *fields)) + if __name__ == "__main__": p = 3 diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index 0643d0d4..100d492c 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -5,15 +5,17 @@ from optimism.FieldEvaluator import * class TestBasics: - p = 2 + p = 1 dim = 2 - mesh = Mesh.construct_structured_mesh(2, 2, [0.0, 3.0], [0.0, 2.0], p) + length = 3.0 + width = 2.0 + mesh = Mesh.construct_structured_mesh(2, 2, [0.0, length], [0.0, width], p) quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) def test_gradient_evaluation(self): "Check the gradient of an affine field" - p = 2 - spaces = [PkField(p, self.dim)] + p = 1 + spaces = [PkField(p, self.dim, self.mesh)] inputs = [Gradient(0)] field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) @@ -31,29 +33,30 @@ def f(dudX): assert pytest.approx(H) == target_disp_grad def test_trivial_integral(self): - spaces = [PkField(self.p, self.dim)] + spaces = [PkField(self.p, self.dim, self.mesh)] inputs = [Value(0)] field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) U = np.zeros_like(self.mesh.coords) def f(u): return 1.0 area = field_evaluator.integrate(f, self.mesh.coords, U) - assert pytest.approx(area) == 6 + assert pytest.approx(area) == self.length*self.width def test_integral_with_one_nodal_field(self): "Computes area in a non-trivial way, checking consistency of gradient and integral operators." - spaces = [PkField(self.p, self.dim)] + spaces = [PkField(self.p, self.dim, self.mesh)] POSITION = 0 # We're taking the gradient of position, which is just the identity tensor inputs = [Gradient(POSITION)] field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) def f(dXdX): - # dXdX == identity + # note dXdX == identity, so + # trace(dXdX)/dim = 1 return np.trace(dXdX)/self.dim area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.coords) - assert pytest.approx(area) == 6 + assert pytest.approx(area) == self.length*self.width def test_conn(self): p = 1 From 1d1b74281160a758311144687989c6167aac43b6 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Tue, 5 Aug 2025 16:51:02 -0700 Subject: [PATCH 16/37] Prevent (for now) use of Pk space with higher order meshes, until mesh class is updated Mesh class needs to keep track of simplex connectivity --- optimism/FieldEvaluator.py | 7 ++++++- optimism/test/test_FieldEvaluator.py | 8 ++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index d3ee8b55..dfd2b634 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -43,6 +43,11 @@ class PkField(Field): quadpoint_axis = 0 def __init__(self, k, dim, mesh): + # temporarily restrict to first order meshes. + # The logic here will work with higher-order meshes (and even curved + # meshes), but building the connectivity table for higher-order + # fields requires knowing the simplex connectivity. + assert mesh.parentElement.degree == 1 self.order = k self.dim = dim self.element, self.element1d = Interpolants.make_parent_elements(k) @@ -176,7 +181,7 @@ def _choose_interpolation_function(input, spaces): class FieldEvaluator: def __init__(self, spaces, qfunction_signature, mesh, quadrature_rule): - # the coord space should live on the Mesh + # the coord space should live on the Mesh eventually self._coord_space = PkField(mesh.parentElement.degree, mesh.coords.shape[1], mesh) self._coord_shapes = self._coord_space.compute_shape_functions(quadrature_rule.xigauss) diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index 100d492c..44431152 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -5,11 +5,11 @@ from optimism.FieldEvaluator import * class TestBasics: - p = 1 + coord_degree = 1 dim = 2 length = 3.0 width = 2.0 - mesh = Mesh.construct_structured_mesh(2, 2, [0.0, length], [0.0, width], p) + mesh = Mesh.construct_structured_mesh(2, 2, [0.0, length], [0.0, width], coord_degree) quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) def test_gradient_evaluation(self): @@ -33,7 +33,7 @@ def f(dudX): assert pytest.approx(H) == target_disp_grad def test_trivial_integral(self): - spaces = [PkField(self.p, self.dim, self.mesh)] + spaces = [PkField(self.coord_degree, self.dim, self.mesh)] inputs = [Value(0)] field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) U = np.zeros_like(self.mesh.coords) @@ -44,7 +44,7 @@ def f(u): def test_integral_with_one_nodal_field(self): "Computes area in a non-trivial way, checking consistency of gradient and integral operators." - spaces = [PkField(self.p, self.dim, self.mesh)] + spaces = [PkField(self.coord_degree, self.dim, self.mesh)] POSITION = 0 # We're taking the gradient of position, which is just the identity tensor inputs = [Gradient(POSITION)] From 733e3483c85579322c71d2f6991977e4ae09b2a6 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Thu, 7 Aug 2025 11:18:34 -0700 Subject: [PATCH 17/37] Remove field dimension from Field classes, they don't need to know this --- optimism/FieldEvaluator.py | 8 ++------ optimism/test/test_FieldEvaluator.py | 22 ++++++++++++++++++---- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index dfd2b634..3be3e34b 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -42,14 +42,13 @@ class PkField(Field): element_axis = None quadpoint_axis = 0 - def __init__(self, k, dim, mesh): + def __init__(self, k, mesh): # temporarily restrict to first order meshes. # The logic here will work with higher-order meshes (and even curved # meshes), but building the connectivity table for higher-order # fields requires knowing the simplex connectivity. assert mesh.parentElement.degree == 1 self.order = k - self.dim = dim self.element, self.element1d = Interpolants.make_parent_elements(k) self.mesh = mesh self.conns = self._make_connectivity() @@ -139,9 +138,6 @@ class QuadratureField(Field): element_axis = 0 quadpoint_axis = 0 - def __init__(self, dim): - self.dim = dim - def interpolate(self, shape, field, conn): return field @@ -182,7 +178,7 @@ def _choose_interpolation_function(input, spaces): class FieldEvaluator: def __init__(self, spaces, qfunction_signature, mesh, quadrature_rule): # the coord space should live on the Mesh eventually - self._coord_space = PkField(mesh.parentElement.degree, mesh.coords.shape[1], mesh) + self._coord_space = PkField(mesh.parentElement.degree, mesh) self._coord_shapes = self._coord_space.compute_shape_functions(quadrature_rule.xigauss) self._spaces = spaces diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index 44431152..95665bd9 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -15,7 +15,7 @@ class TestBasics: def test_gradient_evaluation(self): "Check the gradient of an affine field" p = 1 - spaces = [PkField(p, self.dim, self.mesh)] + spaces = [PkField(p, self.mesh)] inputs = [Gradient(0)] field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) @@ -33,7 +33,7 @@ def f(dudX): assert pytest.approx(H) == target_disp_grad def test_trivial_integral(self): - spaces = [PkField(self.coord_degree, self.dim, self.mesh)] + spaces = [PkField(self.coord_degree, self.mesh)] inputs = [Value(0)] field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) U = np.zeros_like(self.mesh.coords) @@ -44,7 +44,7 @@ def f(u): def test_integral_with_one_nodal_field(self): "Computes area in a non-trivial way, checking consistency of gradient and integral operators." - spaces = [PkField(self.coord_degree, self.dim, self.mesh)] + spaces = [PkField(self.coord_degree, self.mesh)] POSITION = 0 # We're taking the gradient of position, which is just the identity tensor inputs = [Gradient(POSITION)] @@ -60,6 +60,20 @@ def f(dXdX): def test_conn(self): p = 1 - disp_space = PkField(p, self.dim, self.mesh) + disp_space = PkField(p, self.mesh) conns = disp_space._make_connectivity() print(f"{conns=}") + + def test_helmholtz(self): + space = PkField(1, self.mesh) + inputs = [Value(0), Gradient(0)] + + def f(u, dudX): + return 0.5*(u*u + np.dot(dudX, dudX)) + + target_disp_grad = np.array([[0.1, 0.01], + [0.05, 0.3]]) + U = np.einsum('aj, ij', self.mesh.coords, target_disp_grad + np.identity(2)) - self.mesh.coords + field_evaluator = FieldEvaluator([space], inputs, self.mesh, self.quad_rule) + energy = field_evaluator.integrate(f, self.mesh.coords, U) + print(f"{energy=}") From 75ff05569cf69d41f4c5b3da0a8a28fe95d386cc Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Thu, 7 Aug 2025 11:24:24 -0700 Subject: [PATCH 18/37] Add a test with a quadrature field --- optimism/test/test_FieldEvaluator.py | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index 95665bd9..f347c6ed 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -57,23 +57,20 @@ def f(dXdX): area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.coords) assert pytest.approx(area) == self.length*self.width - - def test_conn(self): - p = 1 - disp_space = PkField(p, self.mesh) - conns = disp_space._make_connectivity() - print(f"{conns=}") def test_helmholtz(self): - space = PkField(1, self.mesh) - inputs = [Value(0), Gradient(0)] + spaces = [PkField(1, self.mesh), QuadratureField()] + inputs = [Value(0), Gradient(0), Value(1)] + field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) - def f(u, dudX): - return 0.5*(u*u + np.dot(dudX, dudX)) + def f(u, dudX, q): + return 0.5*q[0]*(u*u + np.dot(dudX, dudX)) target_disp_grad = np.array([[0.1, 0.01], [0.05, 0.3]]) U = np.einsum('aj, ij', self.mesh.coords, target_disp_grad + np.identity(2)) - self.mesh.coords - field_evaluator = FieldEvaluator([space], inputs, self.mesh, self.quad_rule) - energy = field_evaluator.integrate(f, self.mesh.coords, U) + + Q = 2*np.ones((Mesh.num_elements(self.mesh), len(self.quad_rule), 1)) + + energy = field_evaluator.integrate(f, self.mesh.coords, U, Q) print(f"{energy=}") From d1f67797858a12d502ce9401c1c1c472e1a89c8d Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Thu, 7 Aug 2025 14:40:03 -0700 Subject: [PATCH 19/37] Test helmholtz against analytical solution --- optimism/test/test_FieldEvaluator.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index f347c6ed..4146da62 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -8,8 +8,8 @@ class TestBasics: coord_degree = 1 dim = 2 length = 3.0 - width = 2.0 - mesh = Mesh.construct_structured_mesh(2, 2, [0.0, length], [0.0, width], coord_degree) + height = 2.0 + mesh = Mesh.construct_structured_mesh(2, 2, [0.0, length], [0.0, height], coord_degree) quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) def test_gradient_evaluation(self): @@ -40,7 +40,7 @@ def test_trivial_integral(self): def f(u): return 1.0 area = field_evaluator.integrate(f, self.mesh.coords, U) - assert pytest.approx(area) == self.length*self.width + assert pytest.approx(area) == self.length*self.height def test_integral_with_one_nodal_field(self): "Computes area in a non-trivial way, checking consistency of gradient and integral operators." @@ -56,21 +56,22 @@ def f(dXdX): return np.trace(dXdX)/self.dim area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.coords) - assert pytest.approx(area) == self.length*self.width + assert pytest.approx(area) == self.length*self.height def test_helmholtz(self): spaces = [PkField(1, self.mesh), QuadratureField()] + inputs = [Value(0), Gradient(0), Value(1)] field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) def f(u, dudX, q): return 0.5*q[0]*(u*u + np.dot(dudX, dudX)) - target_disp_grad = np.array([[0.1, 0.01], - [0.05, 0.3]]) - U = np.einsum('aj, ij', self.mesh.coords, target_disp_grad + np.identity(2)) - self.mesh.coords + target_grad = np.array([0.1, 0.01]) + U = self.mesh.coords@target_grad + 2.0 Q = 2*np.ones((Mesh.num_elements(self.mesh), len(self.quad_rule), 1)) energy = field_evaluator.integrate(f, self.mesh.coords, U, Q) - print(f"{energy=}") + assert energy == pytest.approx(28.0994) + From d1195d25053220ac6d4cff610dda910c34b53c8e Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Fri, 8 Aug 2025 14:00:40 -0700 Subject: [PATCH 20/37] Change how loop over elements is implemented so that you can use fields of different polynomial order --- optimism/FieldEvaluator.py | 45 +++++++++++----------------- optimism/test/test_FieldEvaluator.py | 5 ++-- 2 files changed, 21 insertions(+), 29 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index 3be3e34b..fec86e79 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -11,16 +11,11 @@ class Field(abc.ABC): @abc.abstractmethod - def interpolate(self, shape, U, conn, jac): + def interpolate(self, shape, U, el_id): pass @abc.abstractmethod - def interpolate_gradient(self, dshape, U, conn, jac): - pass - - @property - @abc.abstractmethod - def element_axis(self): + def interpolate_gradient(self, dshape, U, el_id): pass @property @@ -39,7 +34,6 @@ def map_shape_functions(self, shapes, jacs): class PkField(Field): """Standard Lagrange polynomial finite element fields.""" - element_axis = None quadpoint_axis = 0 def __init__(self, k, mesh): @@ -53,12 +47,12 @@ def __init__(self, k, mesh): self.mesh = mesh self.conns = self._make_connectivity() - def interpolate(self, shape, field, conn): - e_vector = field[conn] + def interpolate(self, shape, field, el_id): + e_vector = field[self.conns[el_id]] return shape.values@e_vector - def interpolate_gradient(self, shape, field, conn): - return jax.vmap(lambda ev, dshp: np.tensordot(ev, dshp, (0, 0)), (None, 0))(field[conn], shape.gradients) + def interpolate_gradient(self, shape, field, el_id): + return jax.vmap(lambda ev, dshp: np.tensordot(ev, dshp, (0, 0)), (None, 0))(field[self.conns[el_id]], shape.gradients) def compute_shape_functions(self, points): return Interpolants.compute_shapes(self.element, points) @@ -117,13 +111,12 @@ def add_element_interior_nodes(conn, newNodeOrdinals): class UniformField(Field): """A unique value for the whole mesh (things like time).""" - element_axis = None quadpoint_axis = None - def interpolate(self, shape, field, conn): - return field + def interpolate(self, shape, U, el_id): + return U - def interpolate_gradient(self, shape, field, conn): + def interpolate_gradient(self, shape, U, el_id): raise NotImplementedError(f"Gradients not supported for {type(self).__name__}") def compute_shape_functions(self, points): @@ -135,13 +128,12 @@ def map_shape_functions(self, shapes, jacs): class QuadratureField(Field): """Arrays defined directly at quadrature points (things like internal variables).""" - element_axis = 0 quadpoint_axis = 0 - def interpolate(self, shape, field, conn): - return field + def interpolate(self, shape, field, el_id): + return field[el_id] - def interpolate_gradient(self, shape, field, conn): + def interpolate_gradient(self, shape, field, el_id): raise NotImplementedError(f"Gradients not supported for {type(self).__name__}") def compute_shape_functions(self, points): @@ -190,18 +182,17 @@ def __init__(self, spaces, qfunction_signature, mesh, quadrature_rule): def evaluate(self, f, coords, *fields): f_vmap_axis = None - conns_vmap_axis = 0 - coords_vmap_axis = None - compute_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, conns_vmap_axis, coords_vmap_axis) + tuple(space.element_axis for space in self._spaces)) - return compute_values(f, self._mesh.conns, coords, *fields) + elems_in_block = np.arange(Mesh.num_elements(self._mesh)) + compute_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, 0, None) + tuple(None for field in fields)) + return compute_values(f, elems_in_block, coords, *fields) def integrate(self, f, coords, *fields): return np.sum(self.evaluate(f, coords, *fields)) - def _evaluate_on_element(self, f, el_conn, coords, *fields): - jacs = self._coord_space.interpolate_gradient(self._coord_shapes, coords, el_conn) + def _evaluate_on_element(self, f, el_id, coords, *fields): + jacs = self._coord_space.interpolate_gradient(self._coord_shapes, coords, el_id) shapes = [space.map_shape_functions(shape, jacs) for (space, shape) in zip(self._spaces, self._shapes)] - f_args = [interp(shapes[field_id], fields[field_id], el_conn) for (interp, field_id) in zip(self._interpolators, self._input_fields)] + f_args = [interp(shapes[field_id], fields[field_id], el_id) for (interp, field_id) in zip(self._interpolators, self._input_fields)] f_batch = jax.vmap(f, tuple(self._spaces[input].quadpoint_axis for input in self._input_fields)) return f_batch(*f_args) diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index 4146da62..b3a9011b 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -59,7 +59,7 @@ def f(dXdX): assert pytest.approx(area) == self.length*self.height def test_helmholtz(self): - spaces = [PkField(1, self.mesh), QuadratureField()] + spaces = [PkField(2, self.mesh), QuadratureField()] inputs = [Value(0), Gradient(0), Value(1)] field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) @@ -68,7 +68,8 @@ def f(u, dudX, q): return 0.5*q[0]*(u*u + np.dot(dudX, dudX)) target_grad = np.array([0.1, 0.01]) - U = self.mesh.coords@target_grad + 2.0 + high_order_mesh = Mesh.create_higher_order_mesh_from_simplex_mesh(self.mesh, 2) + U = high_order_mesh.coords@target_grad + 2.0 Q = 2*np.ones((Mesh.num_elements(self.mesh), len(self.quad_rule), 1)) From cd51e324b7ef148ca2484a87a54a483b54c28d6b Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Sun, 10 Aug 2025 14:39:42 -0700 Subject: [PATCH 21/37] Add coordinates to all spaces, get integrals with higher order fields working --- optimism/FieldEvaluator.py | 49 +++++++++++++++------------- optimism/test/test_FieldEvaluator.py | 7 ++-- 2 files changed, 29 insertions(+), 27 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index fec86e79..510055c3 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -45,7 +45,7 @@ def __init__(self, k, mesh): self.order = k self.element, self.element1d = Interpolants.make_parent_elements(k) self.mesh = mesh - self.conns = self._make_connectivity() + self.conns, self.coords = self._make_connectivity() def interpolate(self, shape, field, el_id): e_vector = field[self.conns[el_id]] @@ -62,18 +62,20 @@ def map_shape_functions(self, shapes, jacs): return Interpolants.ShapeFunctions(shapes.values, shape_grads) def _make_connectivity(self): - mesh_conns = self.mesh.conns if self.order == 1: - return mesh_conns + return self.mesh.conns, self.mesh.coords - conns = np.zeros((mesh_conns.shape[0], self.element.coordinates.shape[0]), dtype=int) + conns = np.zeros((Mesh.num_elements(self.mesh), self.element.coordinates.shape[0]), dtype=int) # step 1/3: vertex nodes - conns = conns.at[:, self.element.vertexNodes].set(mesh_conns) - nodeOrdinalOffset = mesh_conns.max() + 1 # offset for later node numbering + conns = conns.at[:, self.element.vertexNodes].set(self.mesh.conns) + nodeOrdinalOffset = self.mesh.conns.max() + 1 # offset for later node numbering # step 2/3: mid-edge nodes (excluding vertices) - _, edges = Mesh.create_edges(mesh_conns) + edgeConns, edges = Mesh.create_edges(self.mesh.conns) + A = np.column_stack((1.0 - self.element1d.coordinates[self.element1d.interiorNodes], + self.element1d.coordinates[self.element1d.interiorNodes])) + edgeCoords = jax.vmap(lambda edgeConn: np.dot(A, self.mesh.coords[edgeConn, :]))(edgeConns) nNodesPerEdge = self.element1d.interiorNodes.size for e, edge in enumerate(edges): @@ -96,6 +98,11 @@ def _make_connectivity(self): # step 3/3: interior nodes nInNodesPerTri = self.element.interiorNodes.shape[0] if nInNodesPerTri > 0: + N0 = self.element.coordinates[self.element.interiorNodes, 0] + N1 = self.element.coordinates[self.element.interiorNodes, 1] + N2 = 1.0 - N0 - N1 + A = np.column_stack((N0, N1, N2)) + interiorCoords = jax.vmap(lambda triConn: np.dot(A, self.mesh.coords[triConn]))(self.mesh.conns) def add_element_interior_nodes(conn, newNodeOrdinals): return conn.at[self.element.interiorNodes].set(newNodeOrdinals) @@ -105,7 +112,11 @@ def add_element_interior_nodes(conn, newNodeOrdinals): + nodeOrdinalOffset conns = jax.vmap(add_element_interior_nodes)(conns, newNodeOrdinals) - return conns + else: + interiorCoords = np.zeros((0, 2)) + + coords = np.vstack((self.mesh.coords, edgeCoords.reshape(-1,2), interiorCoords.reshape(-1,2))) + return conns, coords @@ -187,7 +198,10 @@ def evaluate(self, f, coords, *fields): return compute_values(f, elems_in_block, coords, *fields) def integrate(self, f, coords, *fields): - return np.sum(self.evaluate(f, coords, *fields)) + f_vmap_axis = None + elems_in_block = np.arange(Mesh.num_elements(self._mesh)) + integrate = jax.vmap(self._integrate_over_element, (f_vmap_axis, 0, None) + tuple(None for field in fields)) + return np.sum(integrate(f, elems_in_block, coords, *fields)) def _evaluate_on_element(self, f, el_id, coords, *fields): jacs = self._coord_space.interpolate_gradient(self._coord_shapes, coords, el_id) @@ -196,23 +210,12 @@ def _evaluate_on_element(self, f, el_id, coords, *fields): f_batch = jax.vmap(f, tuple(self._spaces[input].quadpoint_axis for input in self._input_fields)) return f_batch(*f_args) - def _integrate_over_element(self, f, el_conn, coords, *fields): - jacs = self._coord_space.interpolate_gradient(self._coord_shapes, coords, el_conn) + def _integrate_over_element(self, f, el_id, coords, *fields): + jacs = self._coord_space.interpolate_gradient(self._coord_shapes, coords, el_id) shapes = [space.map_shape_functions(shape, jacs) for (space, shape) in zip(self._spaces, self._shapes)] dVs = jax.vmap(lambda J, w: np.linalg.det(J)*w)(jacs, self._quadrature_rule.wgauss) - f_args = [interp(shapes[field_id], fields[field_id], el_conn) for (interp, field_id) in zip(self._interpolators, self._input_fields)] + f_args = [interp(shapes[field_id], fields[field_id], el_id) for (interp, field_id) in zip(self._interpolators, self._input_fields)] f_batch = jax.vmap(f, tuple(self._spaces[input].quadpoint_axis for input in self._input_fields)) f_vals = f_batch(*f_args) return np.dot(f_vals, dVs) - -if __name__ == "__main__": - p = 3 - dim = 2 - p_coords = 1 - mesh = Mesh.construct_structured_mesh(2, 2, [0.0, 3.0], [0.0, 2.0], p_coords) - disp_space = PkField(p, dim, mesh) - conns = disp_space._make_connectivity() - print(f"{conns=}") - print(f"{mesh.conns=}") - print(f"{mesh.coords=}") diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index b3a9011b..6df7eaad 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -10,7 +10,7 @@ class TestBasics: length = 3.0 height = 2.0 mesh = Mesh.construct_structured_mesh(2, 2, [0.0, length], [0.0, height], coord_degree) - quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) + quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(4) def test_gradient_evaluation(self): "Check the gradient of an affine field" @@ -68,11 +68,10 @@ def f(u, dudX, q): return 0.5*q[0]*(u*u + np.dot(dudX, dudX)) target_grad = np.array([0.1, 0.01]) - high_order_mesh = Mesh.create_higher_order_mesh_from_simplex_mesh(self.mesh, 2) - U = high_order_mesh.coords@target_grad + 2.0 + U = spaces[0].coords@target_grad + 2.0 Q = 2*np.ones((Mesh.num_elements(self.mesh), len(self.quad_rule), 1)) energy = field_evaluator.integrate(f, self.mesh.coords, U, Q) + print(f"{energy:.12e}") assert energy == pytest.approx(28.0994) - From ac2cb84c1f4311e9f24d031a12010a9e1abfd7e5 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Sun, 10 Aug 2025 17:07:26 -0700 Subject: [PATCH 22/37] Add comments and rename things for clarity --- optimism/FieldEvaluator.py | 33 +++++++++++++++++++++------------ 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index 510055c3..8b57af4e 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -38,14 +38,13 @@ class PkField(Field): def __init__(self, k, mesh): # temporarily restrict to first order meshes. - # The logic here will work with higher-order meshes (and even curved - # meshes), but building the connectivity table for higher-order - # fields requires knowing the simplex connectivity. + # Building the connectivity table for higher-order fields requires knowing the + # simplex connectivity. The mesh object must be updated to store that. assert mesh.parentElement.degree == 1 self.order = k self.element, self.element1d = Interpolants.make_parent_elements(k) self.mesh = mesh - self.conns, self.coords = self._make_connectivity() + self.conns, self.coords = self._make_connectivity_and_coordinates() def interpolate(self, shape, field, el_id): e_vector = field[self.conns[el_id]] @@ -61,7 +60,7 @@ def map_shape_functions(self, shapes, jacs): shape_grads = jax.vmap(lambda dshp, J: dshp@np.linalg.inv(J))(shapes.gradients, jacs) return Interpolants.ShapeFunctions(shapes.values, shape_grads) - def _make_connectivity(self): + def _make_connectivity_and_coordinates(self): if self.order == 1: return self.mesh.conns, self.mesh.coords @@ -71,6 +70,10 @@ def _make_connectivity(self): conns = conns.at[:, self.element.vertexNodes].set(self.mesh.conns) nodeOrdinalOffset = self.mesh.conns.max() + 1 # offset for later node numbering + # The non-vertex nodes are placed using linear interpolation. When we add meshes that + # have higher-order coordinate spaces, we must remember to update this part with + # the higher-order interpolation. + # step 2/3: mid-edge nodes (excluding vertices) edgeConns, edges = Mesh.create_edges(self.mesh.conns) A = np.column_stack((1.0 - self.element1d.coordinates[self.element1d.interiorNodes], @@ -155,15 +158,22 @@ def map_shape_functions(self, shapes, jacs): @dataclasses.dataclass -class Value: - """Sentinel to indicate you want to interpolate the value of the field.""" +class FieldInterpolation: + """Abstract base class for specific types of field interpolations. + + This class should not be instantiated, only derived from. + """ field: int -@dataclasses.dataclass -class Gradient: +class Value(FieldInterpolation): + """Sentinel to indicate you want to interpolate the value of the field.""" + pass + + +class Gradient(FieldInterpolation): """Sentinel to indicate you want to interpolate the gradient of the field.""" - field: int + pass def _choose_interpolation_function(input, spaces): @@ -179,7 +189,7 @@ def _choose_interpolation_function(input, spaces): class FieldEvaluator: - def __init__(self, spaces, qfunction_signature, mesh, quadrature_rule): + def __init__(self, spaces: tuple[Field], qfunction_signature: tuple[FieldInterpolation], mesh: Mesh.Mesh, quadrature_rule: QuadratureRule.QuadratureRule) -> None: # the coord space should live on the Mesh eventually self._coord_space = PkField(mesh.parentElement.degree, mesh) self._coord_shapes = self._coord_space.compute_shape_functions(quadrature_rule.xigauss) @@ -218,4 +228,3 @@ def _integrate_over_element(self, f, el_id, coords, *fields): f_batch = jax.vmap(f, tuple(self._spaces[input].quadpoint_axis for input in self._input_fields)) f_vals = f_batch(*f_args) return np.dot(f_vals, dVs) - From 846e1aaa97f19c7fa472377dfa8470cac5809ab2 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Sun, 10 Aug 2025 17:58:46 -0700 Subject: [PATCH 23/37] Add more type hints and input error checking --- optimism/FieldEvaluator.py | 5 ++++- optimism/test/test_FieldEvaluator.py | 24 ++++++++++++++++-------- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index 8b57af4e..edb94d9e 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -36,7 +36,8 @@ class PkField(Field): """Standard Lagrange polynomial finite element fields.""" quadpoint_axis = 0 - def __init__(self, k, mesh): + def __init__(self, k: int, mesh: Mesh.Mesh) -> None: + assert k > 0, "Polynomial degree must be positive" # temporarily restrict to first order meshes. # Building the connectivity table for higher-order fields requires knowing the # simplex connectivity. The mesh object must be updated to store that. @@ -190,6 +191,8 @@ def _choose_interpolation_function(input, spaces): class FieldEvaluator: def __init__(self, spaces: tuple[Field], qfunction_signature: tuple[FieldInterpolation], mesh: Mesh.Mesh, quadrature_rule: QuadratureRule.QuadratureRule) -> None: + for input in qfunction_signature: + assert 0 <= input.field < len(spaces), """Field index in qfunction signature outside valid range.""" # the coord space should live on the Mesh eventually self._coord_space = PkField(mesh.parentElement.degree, mesh) self._coord_shapes = self._coord_space.compute_shape_functions(quadrature_rule.xigauss) diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index 6df7eaad..7564212e 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -15,8 +15,8 @@ class TestBasics: def test_gradient_evaluation(self): "Check the gradient of an affine field" p = 1 - spaces = [PkField(p, self.mesh)] - inputs = [Gradient(0)] + spaces = PkField(p, self.mesh), + inputs = Gradient(0), field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) target_disp_grad = np.array([[0.1, 0.01], @@ -33,8 +33,8 @@ def f(dudX): assert pytest.approx(H) == target_disp_grad def test_trivial_integral(self): - spaces = [PkField(self.coord_degree, self.mesh)] - inputs = [Value(0)] + spaces = PkField(self.coord_degree, self.mesh), + inputs = Value(0), field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) U = np.zeros_like(self.mesh.coords) def f(u): @@ -44,10 +44,10 @@ def f(u): def test_integral_with_one_nodal_field(self): "Computes area in a non-trivial way, checking consistency of gradient and integral operators." - spaces = [PkField(self.coord_degree, self.mesh)] + spaces = PkField(self.coord_degree, self.mesh), POSITION = 0 # We're taking the gradient of position, which is just the identity tensor - inputs = [Gradient(POSITION)] + inputs = Gradient(POSITION), field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) def f(dXdX): @@ -59,14 +59,15 @@ def f(dXdX): assert pytest.approx(area) == self.length*self.height def test_helmholtz(self): - spaces = [PkField(2, self.mesh), QuadratureField()] + spaces = PkField(2, self.mesh), QuadratureField() - inputs = [Value(0), Gradient(0), Value(1)] + inputs = Value(0), Gradient(0), Value(1) field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) def f(u, dudX, q): return 0.5*q[0]*(u*u + np.dot(dudX, dudX)) + # u(X, Y) = 0.1*X + 0.01*Y + 2 target_grad = np.array([0.1, 0.01]) U = spaces[0].coords@target_grad + 2.0 @@ -75,3 +76,10 @@ def f(u, dudX, q): energy = field_evaluator.integrate(f, self.mesh.coords, U, Q) print(f"{energy:.12e}") assert energy == pytest.approx(28.0994) + + def test_nonexistent_field_id_gets_error(self): + spaces = PkField(self.coord_degree, self.mesh), + inputs = Gradient(1), # there is no field 1 + with pytest.raises(AssertionError): + field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + From 83ecaa44064fc0e46068b4cf842db5f0d836e20f Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Sun, 10 Aug 2025 18:00:41 -0700 Subject: [PATCH 24/37] Remove earlier attempt --- optimism/NewFunctionSpace.py | 236 ----------------------------------- 1 file changed, 236 deletions(-) delete mode 100644 optimism/NewFunctionSpace.py diff --git a/optimism/NewFunctionSpace.py b/optimism/NewFunctionSpace.py deleted file mode 100644 index 8932f7ef..00000000 --- a/optimism/NewFunctionSpace.py +++ /dev/null @@ -1,236 +0,0 @@ -import abc -import dataclasses -import functools -import jax -import jax.numpy as np - -from optimism import Interpolants -from optimism import Mesh -from optimism import QuadratureRule -from optimism import VTKWriter - -class Field(abc.ABC): - pass - - @abc.abstractmethod - def interpolate(self, shape, U, conn): - pass - - @abc.abstractmethod - def interpolate_gradient(self, dshape, U, conn): - pass - - @property - @abc.abstractmethod - def element_axis(self): - pass - - @property - @abc.abstractmethod - def quadpoint_axis(self): - pass - - @abc.abstractmethod - def compute_shape_functions(self, points): - pass - - -class PkField(Field): - """Standard Lagrange polynomial finite element fields.""" - element_axis = None - quadpoint_axis = 0 - - def __init__(self, k, dim): - self.order = k - self.dim = dim - self.element, self.element1d = Interpolants.make_parent_elements(k) - - def interpolate(self, shape, field, conn): - e_vector = field[conn] - return shape@e_vector - - def interpolate_gradient(self, dshape, field, conn): - return jax.vmap(lambda ev, dshp: np.tensordot(ev, dshp, (0, 0)), (None, 0))(field[conn], dshape) - - def compute_shape_functions(self, points): - return Interpolants.compute_shapes(self.element, points) - - -class DummyShapeFunctions: - values = None - gradients = None - - -class UniformField(Field): - """A unique value for the whole mesh (things like time).""" - element_axis = None - quadpoint_axis = None - - def interpolate(self, shape, field, conn): - return field - - def interpolate_gradient(self, dshape, field, conn): - grad_shape = field.shape + dshape.shape[-1] - return np.zeros(grad_shape) - - def compute_shape_functions(self, points): - return DummyShapeFunctions() - - -class QuadratureField(Field): - """Arrays defined directly at quadrature points (things like internal variables).""" - element_axis = 0 - quadpoint_axis = 0 - - def __init__(self, dim): - self.dim = dim - - def interpolate(self, shape, Q, conn): - return Q - - def interpolate_gradient(self, dshape, Q, conn): - raise NotImplementedError(f"Gradients not supported for {type(self).__name__}") - - def compute_shape_functions(self, points): - return DummyShapeFunctions() - - -class ParametricElementField(Field): - """Things like quadrature weights and shape function values.""" - element_axis = None - quadpoint_axis = 0 - - def interpolate(self, shape, field, conn): - return field - - def interpolate_gradient(self, dshape, U, conn): - raise NotImplementedError(f"Gradients not supported for {type(self).__name__}") - - def compute_shape_functions(self, points): - return DummyShapeFunctions() - -@dataclasses.dataclass -class Value: - field: int - - -@dataclasses.dataclass -class Gradient: - field: int - - -def pushforward(du_dxi, dX_dxi): - return du_dxi@np.linalg.inv(dX_dxi) - - -def _make_interpolation_function(input, spaces, shapes): - """Helper function to choose correct field space interpolation function for each input.""" - if input.field >= len(spaces): - raise IndexError(f"Field space {input.field} exceeds number of field spaces, which is {len(spaces) - 1}.") - if type(input) is Value: - return functools.partial(spaces[input.field].interpolate, shapes[input.field].values) - elif type(input) is Gradient: - return functools.partial(spaces[input.field].interpolate_gradient, shapes[input.field].gradients) - else: - raise TypeError("Type of object in qfunction signature is invalid.") - - -class FieldEvaluator: - def __init__(self, spaces, qfunction_signature, mesh, quadrature_rule): - self._spaces = spaces - self._mesh = mesh - self._quadrature_rule = quadrature_rule - self._shapes = [space.compute_shape_functions(quadrature_rule.xigauss) for space in spaces] - self._input_fields = tuple(input.field for input in qfunction_signature) - self._interpolators = tuple(_make_interpolation_function(input, self._spaces, self._shapes) for input in qfunction_signature) - - def evaluate(self, f, *fields): - f_vmap_axis = None - conns_vmap_axis = 0 - compute_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, conns_vmap_axis) + tuple(space.element_axis for space in self._spaces)) - return compute_values(f, self._mesh.conns, *fields) - - def _evaluate_on_element(self, f, el_conn, *fields): - f_args = [interp(fields[field_id], el_conn) for (interp, field_id) in zip(self._interpolators, self._input_fields)] - f_batch = jax.vmap(f, tuple(self._spaces[input].quadpoint_axis for input in self._input_fields)) - return f_batch(*f_args) - - -if __name__ == "__main__": - p = 2 - dim = 2 - internal_var_dim = 1 - - mesh = Mesh.construct_structured_mesh(2, 2, [0.0, 1.0], [0.0, 2.0], p) - quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) - - u_space = PkField(p, dim) - internal_variable_space = QuadratureField(internal_var_dim) - time_space = UniformField() - spaces = [u_space, internal_variable_space, time_space] - - inputs = (Value(0), Value(1), Value(2), Gradient(0)) - - fs = FieldEvaluator([u_space, internal_variable_space, time_space], inputs, mesh, quad_rule) - - def f(u, Q, t): - return np.dot(u, u)*Q[0]*t - - # print(f"f() = {f(np.array([1.0, 0.0]), Q[0, 0], 1.0)}") - - U = mesh.coords - U = U.at[:, 1].set(0.0) - Q = 2.0*np.ones((Mesh.num_elements(mesh), len(quad_rule), 1)) - t = 1.0 - - # val = fs.evaluate(f, U, Q, t) - # print(f"func vals shape = {val.shape}") - # print(f"elements = {Mesh.num_elements(mesh)}, quad points = {len(quad_rule)}") - # print(f"{val=}") - - # writer = VTKWriter.VTKWriter(mesh, "new_fs") - # writer.write() - - def g(u, Q, t, du): - return np.dot(u, u)*Q[0]*t*np.tensordot(du, du) - - test = fs._interpolators[0](U, mesh.conns[0]) - print(f"{test=}") - - print(fs._interpolators[1]) - test2 = fs._interpolators[1](Q, mesh.conns[0]) - print(f"{test2=}") - - val = fs.evaluate(g, U, Q, t) - print(f"{val=}") - - def h(dX_dxi, du_dxi, w): - E = 1.0 - nu = 0.0 - mu = 0.5*E/(1 + nu) - lam = E*nu/(1 + nu)/(1 - 2*nu) - dxi_dX = np.linalg.inv(dX_dxi) - du_dX = du_dxi@dxi_dX - dV = np.linalg.det(dX_dxi)*w - strain = 0.5*(du_dX + du_dX.T) - energy_density = mu*np.tensordot(strain, strain) + 0.5*lam*np.trace(strain)**2 - return energy_density*dV - - def area(dX_dxi, du_dxi, w): - return np.linalg.det(dX_dxi)*w - - quad_weight_space = ParametricElementField() - - spaces = [u_space, u_space, quad_weight_space] - qfunction_signature = [Gradient(0), Gradient(1), Value(2)] - - fs2 = FieldEvaluator(spaces, qfunction_signature, mesh, quad_rule) - energies = fs2.evaluate(h, mesh.coords, U, quad_rule.wgauss) - print(f"{energies=}") - - def potential(X, U): - return np.sum(fs2.evaluate(h, X, U, quad_rule.wgauss)) - - compute_force = jax.grad(potential, 1) - R = compute_force(mesh.coords, U) - print(f"{R=}") From 5b90259be0d9d621721e368ab5572cce7d70a37b Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Sun, 10 Aug 2025 18:11:37 -0700 Subject: [PATCH 25/37] Add test that code can jit and support jax gradients --- optimism/test/test_FieldEvaluator.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index 7564212e..460793bc 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -83,3 +83,25 @@ def test_nonexistent_field_id_gets_error(self): with pytest.raises(AssertionError): field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + def test_jit_and_grad(self): + k = 2 + spaces = PkField(k, self.mesh), + inputs = Gradient(0), + field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + + def f(dudX): + return 0.5*np.dot(dudX, dudX) + + @jax.jit + def energy(U): + return field_evaluator.integrate(f, self.mesh.coords, U) + + target_grad = np.array([0.1, 0.01]) + V = spaces[0].coords@target_grad + + e = energy(V) + assert e == pytest.approx(0.5*np.dot(target_grad, target_grad)*self.length*self.height) + + force = jax.jit(jax.grad(energy)) + F = force(V) + assert sum(F) == pytest.approx(0.0) From e9cdb748da95d1cbee5ed0920caef3ec92a470aa Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Mon, 11 Aug 2025 10:10:49 -0700 Subject: [PATCH 26/37] Specify block of elements to integrate over --- optimism/FieldEvaluator.py | 10 ++++------ optimism/test/test_FieldEvaluator.py | 12 ++++++------ 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index edb94d9e..d04ccf73 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -204,17 +204,15 @@ def __init__(self, spaces: tuple[Field], qfunction_signature: tuple[FieldInterpo self._input_fields = tuple(input.field for input in qfunction_signature) self._interpolators = tuple(_choose_interpolation_function(input, self._spaces) for input in qfunction_signature) - def evaluate(self, f, coords, *fields): + def evaluate(self, f, coords, block, *fields): f_vmap_axis = None - elems_in_block = np.arange(Mesh.num_elements(self._mesh)) compute_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, 0, None) + tuple(None for field in fields)) - return compute_values(f, elems_in_block, coords, *fields) + return compute_values(f, block, coords, *fields) - def integrate(self, f, coords, *fields): + def integrate(self, f, coords, block, *fields): f_vmap_axis = None - elems_in_block = np.arange(Mesh.num_elements(self._mesh)) integrate = jax.vmap(self._integrate_over_element, (f_vmap_axis, 0, None) + tuple(None for field in fields)) - return np.sum(integrate(f, elems_in_block, coords, *fields)) + return np.sum(integrate(f, block, coords, *fields)) def _evaluate_on_element(self, f, el_id, coords, *fields): jacs = self._coord_space.interpolate_gradient(self._coord_shapes, coords, el_id) diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index 460793bc..6d839694 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -26,8 +26,8 @@ def test_gradient_evaluation(self): def f(dudX): return dudX - - disp_grads = field_evaluator.evaluate(f, self.mesh.coords, U) + + disp_grads = field_evaluator.evaluate(f, self.mesh.coords, self.mesh.blocks['block_0'], U) for H in disp_grads.reshape(-1, 2, 2): assert pytest.approx(H) == target_disp_grad @@ -39,7 +39,7 @@ def test_trivial_integral(self): U = np.zeros_like(self.mesh.coords) def f(u): return 1.0 - area = field_evaluator.integrate(f, self.mesh.coords, U) + area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], U) assert pytest.approx(area) == self.length*self.height def test_integral_with_one_nodal_field(self): @@ -55,7 +55,7 @@ def f(dXdX): # trace(dXdX)/dim = 1 return np.trace(dXdX)/self.dim - area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.coords) + area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], self.mesh.coords) assert pytest.approx(area) == self.length*self.height def test_helmholtz(self): @@ -73,7 +73,7 @@ def f(u, dudX, q): Q = 2*np.ones((Mesh.num_elements(self.mesh), len(self.quad_rule), 1)) - energy = field_evaluator.integrate(f, self.mesh.coords, U, Q) + energy = field_evaluator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], U, Q) print(f"{energy:.12e}") assert energy == pytest.approx(28.0994) @@ -94,7 +94,7 @@ def f(dudX): @jax.jit def energy(U): - return field_evaluator.integrate(f, self.mesh.coords, U) + return field_evaluator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], U) target_grad = np.array([0.1, 0.01]) V = spaces[0].coords@target_grad From 9ccfeb4330820e8e1d28e812ac23e7fe0a8078c4 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Mon, 11 Aug 2025 13:47:59 -0700 Subject: [PATCH 27/37] write docstrigns --- optimism/FieldEvaluator.py | 47 ++++++++++++++++++++++++++++++++++---- 1 file changed, 43 insertions(+), 4 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index d04ccf73..627b68bd 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -190,26 +190,65 @@ def _choose_interpolation_function(input, spaces): class FieldEvaluator: - def __init__(self, spaces: tuple[Field], qfunction_signature: tuple[FieldInterpolation], mesh: Mesh.Mesh, quadrature_rule: QuadratureRule.QuadratureRule) -> None: + def __init__(self, input_spaces: tuple[Field], qfunction_signature: tuple[FieldInterpolation], + mesh: Mesh.Mesh, quadrature_rule: QuadratureRule.QuadratureRule) -> None: + """Entity that can evaluate and integrate functions at points in a mesh. + + Args: + spaces: Collection of ``Field`` objects describing the inputs + qfunction_signature: + mesh: Finite mesh over which to evaluate/integrate + quad_rule: Quadrature rule to define evaluation points and integration weights + """ for input in qfunction_signature: - assert 0 <= input.field < len(spaces), """Field index in qfunction signature outside valid range.""" + assert 0 <= input.field < len(input_spaces), """Field index in qfunction signature outside valid range.""" # the coord space should live on the Mesh eventually self._coord_space = PkField(mesh.parentElement.degree, mesh) self._coord_shapes = self._coord_space.compute_shape_functions(quadrature_rule.xigauss) - self._spaces = spaces + self._spaces = input_spaces self._mesh = mesh self._quadrature_rule = quadrature_rule - self._shapes = [space.compute_shape_functions(quadrature_rule.xigauss) for space in spaces] + self._shapes = [space.compute_shape_functions(quadrature_rule.xigauss) for space in input_spaces] self._input_fields = tuple(input.field for input in qfunction_signature) self._interpolators = tuple(_choose_interpolation_function(input, self._spaces) for input in qfunction_signature) def evaluate(self, f, coords, block, *fields): + """Evaluate a function at quadrature points. + + Args: + f: Integrand function. + coords: Spatial coordinates of the mesh, used for parametric mapping + of gradients. Can be different values than given in the + constructor, but the shape must be the same. + block: Element ids identifying the domain of integration. + *fields: Input fields to the functional. Must match the + specifications of the ``inputs`` agrument to the constructor. For + performance reasons, this is not checked. + + Returns: + A ``QuadratureField`` with the value of ``f`` at every quadrature point in the mesh. + """ f_vmap_axis = None compute_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, 0, None) + tuple(None for field in fields)) return compute_values(f, block, coords, *fields) def integrate(self, f, coords, block, *fields): + """Integrate a function over a mesh block. + + Args: + f: Integrand function. + coords: Spatial coordinates of the mesh, used for parametric mapping + of gradients. Can be different values than given in the + constructor, but the shape must be the same. + block: Element ids identifying the domain of integration. + *fields: Input fields to the functional. Must match the + specifications of the ``inputs`` agrument to the constructor. For + performance reasons, this is not checked. + + Returns: + The integral of ``f`` over the domain. + """ f_vmap_axis = None integrate = jax.vmap(self._integrate_over_element, (f_vmap_axis, 0, None) + tuple(None for field in fields)) return np.sum(integrate(f, block, coords, *fields)) From 8d1f9a8afd861ec09596e2793bdc493ec859b569 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Mon, 11 Aug 2025 13:52:41 -0700 Subject: [PATCH 28/37] Remove unused imports --- optimism/FieldEvaluator.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index 627b68bd..5c36bf68 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -1,13 +1,12 @@ import abc import dataclasses -import functools + import jax import jax.numpy as np from optimism import Interpolants from optimism import Mesh from optimism import QuadratureRule -from optimism import VTKWriter class Field(abc.ABC): @abc.abstractmethod @@ -235,7 +234,7 @@ def evaluate(self, f, coords, block, *fields): def integrate(self, f, coords, block, *fields): """Integrate a function over a mesh block. - + Args: f: Integrand function. coords: Spatial coordinates of the mesh, used for parametric mapping From d109cd9575d8fc990916758a80d73b4d3203e4f1 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Mon, 11 Aug 2025 14:25:06 -0700 Subject: [PATCH 29/37] Add more type annotations --- optimism/FieldEvaluator.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldEvaluator.py index 5c36bf68..317389fe 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldEvaluator.py @@ -1,8 +1,11 @@ import abc +from collections.abc import Callable import dataclasses import jax import jax.numpy as np +from jax import Array +from jaxtyping import Int, Float from optimism import Interpolants from optimism import Mesh @@ -212,7 +215,7 @@ def __init__(self, input_spaces: tuple[Field], qfunction_signature: tuple[FieldI self._input_fields = tuple(input.field for input in qfunction_signature) self._interpolators = tuple(_choose_interpolation_function(input, self._spaces) for input in qfunction_signature) - def evaluate(self, f, coords, block, *fields): + def evaluate(self, f: Callable, coords: Float[Array, "nnode ndim"], block: Int[Array, "nelem"], *fields: Array) -> Array: """Evaluate a function at quadrature points. Args: @@ -232,7 +235,7 @@ def evaluate(self, f, coords, block, *fields): compute_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, 0, None) + tuple(None for field in fields)) return compute_values(f, block, coords, *fields) - def integrate(self, f, coords, block, *fields): + def integrate(self, f, coords, block, *fields) -> Array: """Integrate a function over a mesh block. Args: From bf528d87265ec1ee3be1f62f7c2a171867103da3 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Mon, 11 Aug 2025 14:55:26 -0700 Subject: [PATCH 30/37] Rename variables in tests for increased clarity --- optimism/test/test_FieldEvaluator.py | 32 ++++++++++++++-------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldEvaluator.py index 6d839694..e6e17df3 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldEvaluator.py @@ -14,14 +14,13 @@ class TestBasics: def test_gradient_evaluation(self): "Check the gradient of an affine field" - p = 1 - spaces = PkField(p, self.mesh), - inputs = Gradient(0), - field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + k = 1 + spaces = PkField(k, self.mesh), + integrand_signature = Gradient(0), + field_evaluator = FieldEvaluator(spaces, integrand_signature, self.mesh, self.quad_rule) target_disp_grad = np.array([[0.1, 0.01], [0.05, 0.3]]) - #coords_linear = self.mesh.coords[self.mesh.simplexNodesOrdinals, :] U = np.einsum('aj, ij', self.mesh.coords, target_disp_grad + np.identity(2)) - self.mesh.coords def f(dudX): @@ -34,8 +33,8 @@ def f(dudX): def test_trivial_integral(self): spaces = PkField(self.coord_degree, self.mesh), - inputs = Value(0), - field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + integrand_signature = Value(0), + field_evaluator = FieldEvaluator(spaces, integrand_signature, self.mesh, self.quad_rule) U = np.zeros_like(self.mesh.coords) def f(u): return 1.0 @@ -43,12 +42,12 @@ def f(u): assert pytest.approx(area) == self.length*self.height def test_integral_with_one_nodal_field(self): - "Computes area in a non-trivial way, checking consistency of gradient and integral operators." - spaces = PkField(self.coord_degree, self.mesh), + "Computes area in a non-trivial way, checking consistency of gradient interpolation." + integrand_signature = PkField(self.coord_degree, self.mesh), POSITION = 0 # We're taking the gradient of position, which is just the identity tensor inputs = Gradient(POSITION), - field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + field_evaluator = FieldEvaluator(integrand_signature, inputs, self.mesh, self.quad_rule) def f(dXdX): # note dXdX == identity, so @@ -59,10 +58,11 @@ def f(dXdX): assert pytest.approx(area) == self.length*self.height def test_helmholtz(self): + "Tests interpolation, gradient, and simple use of a QuadratureField" spaces = PkField(2, self.mesh), QuadratureField() - inputs = Value(0), Gradient(0), Value(1) - field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + integrand_signature = Value(0), Gradient(0), Value(1) + field_evaluator = FieldEvaluator(spaces, integrand_signature, self.mesh, self.quad_rule) def f(u, dudX, q): return 0.5*q[0]*(u*u + np.dot(dudX, dudX)) @@ -79,15 +79,15 @@ def f(u, dudX, q): def test_nonexistent_field_id_gets_error(self): spaces = PkField(self.coord_degree, self.mesh), - inputs = Gradient(1), # there is no field 1 + integrand_signature = Gradient(1), # there is no field 1 with pytest.raises(AssertionError): - field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + field_evaluator = FieldEvaluator(spaces, integrand_signature, self.mesh, self.quad_rule) def test_jit_and_grad(self): k = 2 spaces = PkField(k, self.mesh), - inputs = Gradient(0), - field_evaluator = FieldEvaluator(spaces, inputs, self.mesh, self.quad_rule) + integrand_signature = Gradient(0), + field_evaluator = FieldEvaluator(spaces, integrand_signature, self.mesh, self.quad_rule) def f(dudX): return 0.5*np.dot(dudX, dudX) From 7518a87f933959769df167b6c10bb80010cac43e Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Mon, 11 Aug 2025 15:34:54 -0700 Subject: [PATCH 31/37] Rename class to FieldOperator --- .../{FieldEvaluator.py => FieldOperator.py} | 8 ++++--- ...ieldEvaluator.py => test_FieldOperator.py} | 24 +++++++++---------- 2 files changed, 17 insertions(+), 15 deletions(-) rename optimism/{FieldEvaluator.py => FieldOperator.py} (97%) rename optimism/test/{test_FieldEvaluator.py => test_FieldOperator.py} (74%) diff --git a/optimism/FieldEvaluator.py b/optimism/FieldOperator.py similarity index 97% rename from optimism/FieldEvaluator.py rename to optimism/FieldOperator.py index 317389fe..12fdfa75 100644 --- a/optimism/FieldEvaluator.py +++ b/optimism/FieldOperator.py @@ -191,7 +191,7 @@ def _choose_interpolation_function(input, spaces): raise TypeError("Type of object in qfunction signature is invalid.") -class FieldEvaluator: +class FieldOperator: def __init__(self, input_spaces: tuple[Field], qfunction_signature: tuple[FieldInterpolation], mesh: Mesh.Mesh, quadrature_rule: QuadratureRule.QuadratureRule) -> None: """Entity that can evaluate and integrate functions at points in a mesh. @@ -215,7 +215,8 @@ def __init__(self, input_spaces: tuple[Field], qfunction_signature: tuple[FieldI self._input_fields = tuple(input.field for input in qfunction_signature) self._interpolators = tuple(_choose_interpolation_function(input, self._spaces) for input in qfunction_signature) - def evaluate(self, f: Callable, coords: Float[Array, "nnode ndim"], block: Int[Array, "nelem"], *fields: Array) -> Array: + def evaluate(self, f: Callable, coords: Float[Array, "nnode ndim"], + block: Int[Array, "nelem"], *fields: Array) -> Array: """Evaluate a function at quadrature points. Args: @@ -235,7 +236,8 @@ def evaluate(self, f: Callable, coords: Float[Array, "nnode ndim"], block: Int[A compute_values = jax.vmap(self._evaluate_on_element, (f_vmap_axis, 0, None) + tuple(None for field in fields)) return compute_values(f, block, coords, *fields) - def integrate(self, f, coords, block, *fields) -> Array: + def integrate(self, f: Callable, coords: Float[Array, "nnode ndim"], + block: Int[Array, "nelem"], *fields: Array) -> Array: """Integrate a function over a mesh block. Args: diff --git a/optimism/test/test_FieldEvaluator.py b/optimism/test/test_FieldOperator.py similarity index 74% rename from optimism/test/test_FieldEvaluator.py rename to optimism/test/test_FieldOperator.py index e6e17df3..4057d13b 100644 --- a/optimism/test/test_FieldEvaluator.py +++ b/optimism/test/test_FieldOperator.py @@ -2,7 +2,7 @@ import jax.numpy as np from optimism import Mesh -from optimism.FieldEvaluator import * +from optimism.FieldOperator import * class TestBasics: coord_degree = 1 @@ -17,7 +17,7 @@ def test_gradient_evaluation(self): k = 1 spaces = PkField(k, self.mesh), integrand_signature = Gradient(0), - field_evaluator = FieldEvaluator(spaces, integrand_signature, self.mesh, self.quad_rule) + field_operator = FieldOperator(spaces, integrand_signature, self.mesh, self.quad_rule) target_disp_grad = np.array([[0.1, 0.01], [0.05, 0.3]]) @@ -26,7 +26,7 @@ def test_gradient_evaluation(self): def f(dudX): return dudX - disp_grads = field_evaluator.evaluate(f, self.mesh.coords, self.mesh.blocks['block_0'], U) + disp_grads = field_operator.evaluate(f, self.mesh.coords, self.mesh.blocks['block_0'], U) for H in disp_grads.reshape(-1, 2, 2): assert pytest.approx(H) == target_disp_grad @@ -34,11 +34,11 @@ def f(dudX): def test_trivial_integral(self): spaces = PkField(self.coord_degree, self.mesh), integrand_signature = Value(0), - field_evaluator = FieldEvaluator(spaces, integrand_signature, self.mesh, self.quad_rule) + field_operator = FieldOperator(spaces, integrand_signature, self.mesh, self.quad_rule) U = np.zeros_like(self.mesh.coords) def f(u): return 1.0 - area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], U) + area = field_operator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], U) assert pytest.approx(area) == self.length*self.height def test_integral_with_one_nodal_field(self): @@ -47,14 +47,14 @@ def test_integral_with_one_nodal_field(self): POSITION = 0 # We're taking the gradient of position, which is just the identity tensor inputs = Gradient(POSITION), - field_evaluator = FieldEvaluator(integrand_signature, inputs, self.mesh, self.quad_rule) + field_operator = FieldOperator(integrand_signature, inputs, self.mesh, self.quad_rule) def f(dXdX): # note dXdX == identity, so # trace(dXdX)/dim = 1 return np.trace(dXdX)/self.dim - area = field_evaluator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], self.mesh.coords) + area = field_operator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], self.mesh.coords) assert pytest.approx(area) == self.length*self.height def test_helmholtz(self): @@ -62,7 +62,7 @@ def test_helmholtz(self): spaces = PkField(2, self.mesh), QuadratureField() integrand_signature = Value(0), Gradient(0), Value(1) - field_evaluator = FieldEvaluator(spaces, integrand_signature, self.mesh, self.quad_rule) + field_operator = FieldOperator(spaces, integrand_signature, self.mesh, self.quad_rule) def f(u, dudX, q): return 0.5*q[0]*(u*u + np.dot(dudX, dudX)) @@ -73,7 +73,7 @@ def f(u, dudX, q): Q = 2*np.ones((Mesh.num_elements(self.mesh), len(self.quad_rule), 1)) - energy = field_evaluator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], U, Q) + energy = field_operator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], U, Q) print(f"{energy:.12e}") assert energy == pytest.approx(28.0994) @@ -81,20 +81,20 @@ def test_nonexistent_field_id_gets_error(self): spaces = PkField(self.coord_degree, self.mesh), integrand_signature = Gradient(1), # there is no field 1 with pytest.raises(AssertionError): - field_evaluator = FieldEvaluator(spaces, integrand_signature, self.mesh, self.quad_rule) + field_operator = FieldOperator(spaces, integrand_signature, self.mesh, self.quad_rule) def test_jit_and_grad(self): k = 2 spaces = PkField(k, self.mesh), integrand_signature = Gradient(0), - field_evaluator = FieldEvaluator(spaces, integrand_signature, self.mesh, self.quad_rule) + field_operator = FieldOperator(spaces, integrand_signature, self.mesh, self.quad_rule) def f(dudX): return 0.5*np.dot(dudX, dudX) @jax.jit def energy(U): - return field_evaluator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], U) + return field_operator.integrate(f, self.mesh.coords, self.mesh.blocks['block_0'], U) target_grad = np.array([0.1, 0.01]) V = spaces[0].coords@target_grad From 843516165ddb1016179ed2b2240c6422d5223909 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Mon, 11 Aug 2025 16:31:23 -0700 Subject: [PATCH 32/37] Rename test class consistently --- optimism/test/test_FieldOperator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimism/test/test_FieldOperator.py b/optimism/test/test_FieldOperator.py index 4057d13b..c7abf7bc 100644 --- a/optimism/test/test_FieldOperator.py +++ b/optimism/test/test_FieldOperator.py @@ -4,7 +4,7 @@ from optimism import Mesh from optimism.FieldOperator import * -class TestBasics: +class TestFieldOperator: coord_degree = 1 dim = 2 length = 3.0 From 3589dbfd0c20366acb0a240f20137d675713804a Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Tue, 12 Aug 2025 16:02:22 -0700 Subject: [PATCH 33/37] Add DG field with one test --- optimism/FieldOperator.py | 50 ++++++++++++++++++++++++++--- optimism/test/test_FieldOperator.py | 27 ++++++++++++++++ 2 files changed, 72 insertions(+), 5 deletions(-) diff --git a/optimism/FieldOperator.py b/optimism/FieldOperator.py index 12fdfa75..47b13e9c 100644 --- a/optimism/FieldOperator.py +++ b/optimism/FieldOperator.py @@ -17,7 +17,7 @@ def interpolate(self, shape, U, el_id): pass @abc.abstractmethod - def interpolate_gradient(self, dshape, U, el_id): + def interpolate_gradient(self, shape, U, el_id): pass @property @@ -125,6 +125,46 @@ def add_element_interior_nodes(conn, newNodeOrdinals): return conns, coords +class DG_PkField(Field): + quadpoint_axis = 0 + + def __init__(self, k: int, mesh: Mesh.Mesh) -> None: + assert k >= 0, "Polynomial degree must be >= 0" + assert mesh.parentElement.degree == 1 + self.order = k + self.element, self.element1d = Interpolants.make_parent_elements(k) + self.mesh = mesh + self.conns, self.coords = self._make_connectivity_and_coordinates() + self.field_shape = Mesh.num_elements(mesh), self.element.num_nodes + + def _make_connectivity_and_coordinates(self): + nnodes = Mesh.num_elements(self.mesh)*self.element.num_nodes + conns = np.arange(nnodes).reshape(-1, self.element.num_nodes) + + def set_elem_coords(simplex_element_conn): + Xs = self.mesh.coords[simplex_element_conn] + J = np.column_stack((Xs[0] - Xs[2], Xs[1] - Xs[2])) + Jxi = self.element.coordinates@J.T + b = Xs[0] - Jxi[0] + return self.element.coordinates@J.T + b + + coords = jax.vmap(set_elem_coords)(self.mesh.conns) + return conns, coords + + def compute_shape_functions(self, points): + return Interpolants.compute_shapes(self.element, points) + + def interpolate(self, shape, U, el_id): + e_vector = U[el_id] + return shape.values@e_vector + + def interpolate_gradient(self, shape, U, el_id): + return jax.vmap(lambda ev, dshp: np.tensordot(ev, dshp, (0, 0)), (None, 0))(U[el_id], shape.gradients) + + def map_shape_functions(self, shapes, jacs): + shape_grads = jax.vmap(lambda dshp, J: dshp@np.linalg.inv(J))(shapes.gradients, jacs) + return Interpolants.ShapeFunctions(shapes.values, shape_grads) + class UniformField(Field): """A unique value for the whole mesh (things like time).""" @@ -192,13 +232,13 @@ def _choose_interpolation_function(input, spaces): class FieldOperator: - def __init__(self, input_spaces: tuple[Field], qfunction_signature: tuple[FieldInterpolation], + def __init__(self, input_spaces: tuple[Field, ...], qfunction_signature: tuple[FieldInterpolation, ...], mesh: Mesh.Mesh, quadrature_rule: QuadratureRule.QuadratureRule) -> None: """Entity that can evaluate and integrate functions at points in a mesh. Args: spaces: Collection of ``Field`` objects describing the inputs - qfunction_signature: + qfunction_signature: Tells the ``FieldOperator`` what quantitites to interpolate to the quadrature points. mesh: Finite mesh over which to evaluate/integrate quad_rule: Quadrature rule to define evaluation points and integration weights """ @@ -216,7 +256,7 @@ def __init__(self, input_spaces: tuple[Field], qfunction_signature: tuple[FieldI self._interpolators = tuple(_choose_interpolation_function(input, self._spaces) for input in qfunction_signature) def evaluate(self, f: Callable, coords: Float[Array, "nnode ndim"], - block: Int[Array, "nelem"], *fields: Array) -> Array: + block: Int[Array, "nelem"], *fields: Array | tuple[Array, ...]) -> Array: """Evaluate a function at quadrature points. Args: @@ -237,7 +277,7 @@ def evaluate(self, f: Callable, coords: Float[Array, "nnode ndim"], return compute_values(f, block, coords, *fields) def integrate(self, f: Callable, coords: Float[Array, "nnode ndim"], - block: Int[Array, "nelem"], *fields: Array) -> Array: + block: Int[Array, "nelem"], *fields: Array | tuple[Array, ...]) -> Array: """Integrate a function over a mesh block. Args: diff --git a/optimism/test/test_FieldOperator.py b/optimism/test/test_FieldOperator.py index c7abf7bc..4fbbf924 100644 --- a/optimism/test/test_FieldOperator.py +++ b/optimism/test/test_FieldOperator.py @@ -105,3 +105,30 @@ def energy(U): force = jax.jit(jax.grad(energy)) F = force(V) assert sum(F) == pytest.approx(0.0) + + def test_dg_interpolation(self): + "Check interpolation of DG field" + k = 1 + mesh = Mesh.construct_structured_mesh(3, 3, [0.0, 2.0], [0.0, 1.0]) + space = DG_PkField(k, mesh) + + # make a DG field that is u = 0.1*x on x < 1, u = 2 on x > 1 + def field_values(el_coords): + centroid = np.mean(el_coords, axis=0) + return np.where(centroid[0] < 1.0, 0.1*el_coords[:, 0], 2.0*np.ones_like(el_coords[:, 0])) + + U = jax.vmap(field_values)(space.coords) + + coord_space = PkField(1, mesh) + field_operator = FieldOperator((space, coord_space), (Value(0), Value(1)), mesh, self.quad_rule) + def f(u, x): + return u, x + + # elements in left-hand side are [0:4] -> u = 0.1*x + Uq, Xq = field_operator.evaluate(f, mesh.coords, np.arange(4), U, mesh.coords) + Uq_expected = 0.1*Xq[..., 0] + assert Uq == pytest.approx(Uq_expected) + + # elements in right-hand side [4:7] u = 2.0 + Uq, Xq = field_operator.evaluate(f, mesh.coords, np.arange(4, 8), U, mesh.coords) + assert Uq == pytest.approx(2.0) From e4a348205359ab1904e3c8fe2baded71ea3e3dcb Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Tue, 12 Aug 2025 21:37:12 -0700 Subject: [PATCH 34/37] Test interpolation of gradient for DG space --- optimism/test/test_FieldOperator.py | 51 ++++++++++++++++++++++------- 1 file changed, 40 insertions(+), 11 deletions(-) diff --git a/optimism/test/test_FieldOperator.py b/optimism/test/test_FieldOperator.py index 4fbbf924..6fff941a 100644 --- a/optimism/test/test_FieldOperator.py +++ b/optimism/test/test_FieldOperator.py @@ -106,11 +106,14 @@ def energy(U): F = force(V) assert sum(F) == pytest.approx(0.0) + +class TestDGField: + quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(1) + mesh = Mesh.construct_structured_mesh(3, 3, [0.0, 2.0], [0.0, 1.0]) + def test_dg_interpolation(self): - "Check interpolation of DG field" k = 1 - mesh = Mesh.construct_structured_mesh(3, 3, [0.0, 2.0], [0.0, 1.0]) - space = DG_PkField(k, mesh) + space = DG_PkField(k, self.mesh) # make a DG field that is u = 0.1*x on x < 1, u = 2 on x > 1 def field_values(el_coords): @@ -119,16 +122,42 @@ def field_values(el_coords): U = jax.vmap(field_values)(space.coords) - coord_space = PkField(1, mesh) - field_operator = FieldOperator((space, coord_space), (Value(0), Value(1)), mesh, self.quad_rule) + coord_space = PkField(1, self.mesh) + field_operator = FieldOperator((space, coord_space), (Value(0), Value(1)), self.mesh, self.quad_rule) def f(u, x): return u, x - + # elements in left-hand side are [0:4] -> u = 0.1*x - Uq, Xq = field_operator.evaluate(f, mesh.coords, np.arange(4), U, mesh.coords) - Uq_expected = 0.1*Xq[..., 0] - assert Uq == pytest.approx(Uq_expected) + u_q, x_q = field_operator.evaluate(f, self.mesh.coords, np.arange(4), U, self.mesh.coords) + Uq_expected = 0.1*x_q[..., 0] + assert u_q == pytest.approx(Uq_expected) # elements in right-hand side [4:7] u = 2.0 - Uq, Xq = field_operator.evaluate(f, mesh.coords, np.arange(4, 8), U, mesh.coords) - assert Uq == pytest.approx(2.0) + u_q, x_q = field_operator.evaluate(f, self.mesh.coords, np.arange(4, 8), U, self.mesh.coords) + assert u_q == pytest.approx(2.0) + + def test_dg_gradient_interpolation(self): + k = 1 + self.mesh = Mesh.construct_structured_mesh(3, 3, [0.0, 2.0], [0.0, 1.0]) + space = DG_PkField(k, self.mesh) + + # make a DG field that is u = 0.1*x on x < 1, u = 2 on x > 1 + def field_values(el_coords): + centroid = np.mean(el_coords, axis=0) + return np.where(centroid[0] < 1.0, 0.1*el_coords[:, 0], 2.0*np.ones_like(el_coords[:, 0])) + + U = jax.vmap(field_values)(space.coords) + + coord_space = PkField(1, self.mesh) + field_operator = FieldOperator((space, coord_space), (Gradient(0), Value(1)), self.mesh, self.quad_rule) + def f(dudX, x): + return dudX, x + + # elements in left-hand side are [0:4] -> grad u = [0.1, 0.0] + dudX_q, _ = field_operator.evaluate(f, self.mesh.coords, np.arange(4), U, self.mesh.coords) + for dudX in dudX_q.reshape(-1, 2): + assert dudX == pytest.approx(np.array([0.1, 0.0])) + + # elements in right-hand side [4:7] -> grad u = 0 + dudX_q, _ = field_operator.evaluate(f, self.mesh.coords, np.arange(4, 8), U, self.mesh.coords) + assert dudX_q == pytest.approx(0.0) From 965f55c7ff72505d1e60d7ae3ebeeef2d2cfe5fd Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Wed, 13 Aug 2025 11:55:57 -0700 Subject: [PATCH 35/37] Test DG P0 field --- optimism/test/test_FieldOperator.py | 44 +++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/optimism/test/test_FieldOperator.py b/optimism/test/test_FieldOperator.py index 6fff941a..75f5edec 100644 --- a/optimism/test/test_FieldOperator.py +++ b/optimism/test/test_FieldOperator.py @@ -161,3 +161,47 @@ def f(dudX, x): # elements in right-hand side [4:7] -> grad u = 0 dudX_q, _ = field_operator.evaluate(f, self.mesh.coords, np.arange(4, 8), U, self.mesh.coords) assert dudX_q == pytest.approx(0.0) + +def test_parameterized_elasticity(): + mesh = Mesh.construct_structured_mesh(5, 3, [0.0, 1.0], [0.0, 1.0]) + ne = Mesh.num_elements(mesh) + blocks = {'all': np.arange(Mesh.num_elements(mesh)), + 'left': np.arange(ne//2), + 'right': np.arange(ne//2, ne)} + quad_rule = QuadratureRule.create_quadrature_rule_on_triangle(2) + spaces = PkField(1, mesh), DG_PkField(0, mesh) + integrand_signature = Gradient(0), Value(1) + field_operator = FieldOperator(spaces, integrand_signature, mesh, quad_rule) + + lam = 3.0 + + def f(dudX, mu): + strain = 0.5*(dudX + dudX.T) + return mu*np.tensordot(strain, strain) + 0.5*lam*np.trace(strain)**2 + + target_disp_grad = np.array([[0.1, 0.01], + [0.05, 0.3]]) + U = np.einsum('aj, ij', mesh.coords, target_disp_grad + np.identity(2)) - mesh.coords + + mu_left = 1.0 + mu_right = 2.0 + mu = np.zeros(spaces[1].field_shape) + mu = mu.at[blocks['left']].set(mu_left) + mu = mu.at[blocks['right']].set(mu_right) + + def energy(U, mu): + return field_operator.integrate(f, mesh.coords, blocks['all'], U, mu) + + R = jax.grad(energy, 0)(U, mu) + assert R[6] == pytest.approx(np.zeros(2)) + assert R[8] == pytest.approx(np.zeros(2)) + + + stresses = field_operator.evaluate(jax.grad(f, 0), mesh.coords, blocks['all'], U, mu) + strain = 0.5*(target_disp_grad + target_disp_grad.T) + stress_left_exact = 2.0*mu_left*strain + lam*np.trace(strain)*np.identity(2) + stress_right_exact = 2.0*mu_right*strain + lam*np.trace(strain)*np.identity(2) + for stress in stresses[blocks['left']].reshape(-1, 2, 2): + assert stress == pytest.approx(stress_left_exact) + for stress in stresses[blocks['right']].reshape(-1, 2, 2): + assert stress == pytest.approx(stress_right_exact) From 74d1413abdd3b13d9ba276c981f521f32a1251a7 Mon Sep 17 00:00:00 2001 From: Brandon Talamini Date: Fri, 22 Aug 2025 09:12:31 -0700 Subject: [PATCH 36/37] Make type annotations work with older versions of Python with a future import --- optimism/FieldOperator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/optimism/FieldOperator.py b/optimism/FieldOperator.py index 47b13e9c..7545e45b 100644 --- a/optimism/FieldOperator.py +++ b/optimism/FieldOperator.py @@ -1,3 +1,4 @@ +from __future__ import annotations import abc from collections.abc import Callable import dataclasses From cca6878e517e621a9fab7f6a0adfbcfa9f69b838 Mon Sep 17 00:00:00 2001 From: Craig Hamel <31457225+cmhamel@users.noreply.github.com> Date: Thu, 28 Aug 2025 16:48:09 -0400 Subject: [PATCH 37/37] Update ci-build.yml --- .github/workflows/ci-build.yml | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/.github/workflows/ci-build.yml b/.github/workflows/ci-build.yml index 94385938..38fce0da 100644 --- a/.github/workflows/ci-build.yml +++ b/.github/workflows/ci-build.yml @@ -36,18 +36,18 @@ jobs: pip install -e ".[sparse,test]" python -m pytest -n auto optimism --cov=optimism -Wignore # we can also add the flag -n auto for parallel testing - - name: docs - run: | - pip install -e ".[docs,sparse,test]" - cd docs - sphinx-apidoc -o source/ ../optimism -P - make html - - name: Deploy to GitHub Pages - uses: peaceiris/actions-gh-pages@v3 - with: - github_token: ${{ secrets.GITHUB_TOKEN }} - publish_dir: docs/build/html # Adjust this if your output directory is different - publish_branch: gh-pages # The branch to deploy to + # - name: docs + # run: | + # pip install -e ".[docs,sparse,test]" + # cd docs + # sphinx-apidoc -o source/ ../optimism -P + # make html + # - name: Deploy to GitHub Pages + # uses: peaceiris/actions-gh-pages@v3 + # with: + # github_token: ${{ secrets.GITHUB_TOKEN }} + # publish_dir: docs/build/html # Adjust this if your output directory is different + # publish_branch: gh-pages # The branch to deploy to - name: codecov uses: codecov/codecov-action@v4 with: